cache-dit 0.3.2__py3-none-any.whl → 1.0.14__py3-none-any.whl

cache_dit/caching/cache_interface.py

@@ -0,0 +1,358 @@
+import torch
+from typing import Any, Tuple, List, Union, Optional
+from diffusers import DiffusionPipeline, ModelMixin
+from cache_dit.caching.cache_types import CacheType
+from cache_dit.caching.block_adapters import BlockAdapter
+from cache_dit.caching.block_adapters import BlockAdapterRegistry
+from cache_dit.caching.cache_adapters import CachedAdapter
+from cache_dit.caching.cache_contexts import BasicCacheConfig
+from cache_dit.caching.cache_contexts import DBCacheConfig
+from cache_dit.caching.cache_contexts import DBPruneConfig
+from cache_dit.caching.cache_contexts import CalibratorConfig
+from cache_dit.caching.params_modifier import ParamsModifier
+from cache_dit.parallelism import ParallelismConfig
+from cache_dit.parallelism import enable_parallelism
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def enable_cache(
+    # DiffusionPipeline or BlockAdapter
+    pipe_or_adapter: Union[
+        DiffusionPipeline,
+        BlockAdapter,
+        # Transformer-only
+        torch.nn.Module,
+        ModelMixin,
+    ],
+    # BasicCacheConfig, DBCacheConfig, DBPruneConfig, etc.
+    cache_config: Optional[
+        Union[
+            BasicCacheConfig,
+            DBCacheConfig,
+            DBPruneConfig,
+        ]
+    ] = None,
+    # Calibrator config: TaylorSeerCalibratorConfig, etc.
+    calibrator_config: Optional[CalibratorConfig] = None,
+    # Modify cache context params for specific blocks.
+    params_modifiers: Optional[
+        Union[
+            ParamsModifier,
+            List[ParamsModifier],
+            List[List[ParamsModifier]],
+        ]
+    ] = None,
+    # Config for Parallelism
+    parallelism_config: Optional[ParallelismConfig] = None,
+    # Other cache context kwargs: Deprecated cache kwargs
+    **kwargs,
+) -> Union[
+    DiffusionPipeline,
+    # Transformer-only
+    torch.nn.Module,
+    ModelMixin,
+    BlockAdapter,
+]:
+    r"""
+    The `enable_cache` function serves as a unified caching interface designed to optimize the performance
+    of diffusion transformer models by implementing an intelligent caching mechanism known as `DBCache`.
+    This API is engineered to be compatible with nearly `all` diffusion transformer architectures that
+    feature transformer blocks adhering to standard input-output patterns, eliminating the need for
+    architecture-specific modifications.
+
+    By strategically caching intermediate outputs of transformer blocks during the diffusion process,
+    `DBCache` significantly reduces redundant computations without compromising generation quality.
+    The caching mechanism works by tracking residual differences between consecutive steps, allowing
+    the model to reuse previously computed features when these differences fall below a configurable
+    threshold. This approach maintains a balance between computational efficiency and output precision.
+
+    The default configuration (`F8B0, 8 warmup steps, unlimited cached steps`) is carefully tuned to
+    provide an optimal tradeoff for most common use cases. The "F8B0" configuration indicates that
+    the first 8 transformer blocks are used to compute stable feature differences, while no final
+    blocks are employed for additional fusion. The warmup phase ensures the model establishes
+    sufficient feature representation before caching begins, preventing potential degradation of
+    output quality.
+
+    This function seamlessly integrates with both standard diffusion pipelines and custom block
+    adapters, making it versatile for various deployment scenarios—from research prototyping to
+    production environments where inference speed is critical. By abstracting the complexity of
+    caching logic behind a simple interface, it enables developers to enhance model performance
+    with minimal code changes.
+
+    Args:
+        pipe_or_adapter (`DiffusionPipeline`, `BlockAdapter` or `Transformer`, *required*):
+            The standard Diffusion Pipeline or custom BlockAdapter (from cache-dit or user-defined).
+            For example: cache_dit.enable_cache(FluxPipeline(...)).
+
+        cache_config (`BasicCacheConfig`, *required*, defaults to BasicCacheConfig()):
+            Basic DBCache config for cache context, defaults to BasicCacheConfig(). The configurable params listed belows:
+                Fn_compute_blocks: (`int`, *required*, defaults to 8):
+                    Specifies that `DBCache` uses the**first n**Transformer blocks to fit the information at time step t,
+                    enabling the calculation of a more stable L1 difference and delivering more accurate information
+                    to subsequent blocks. Please check https://github.com/vipshop/cache-dit/blob/main/docs/DBCache.md
+                    for more details of DBCache.
+                Bn_compute_blocks: (`int`, *required*, defaults to 0):
+                    Further fuses approximate information in the **last n** Transformer blocks to enhance
+                    prediction accuracy. These blocks act as an auto-scaler for approximate hidden states
+                    that use residual cache.
+                residual_diff_threshold (`float`, *required*, defaults to 0.08):
+                    the value of residual diff threshold, a higher value leads to faster performance at the
+                    cost of lower precision.
+                max_warmup_steps (`int`, *required*, defaults to 8):
+                    DBCache does not apply the caching strategy when the number of running steps is less than
+                    or equal to this value, ensuring the model sufficiently learns basic features during warmup.
+                warmup_interval (`int`, *required*, defaults to 1):
+                    Skip interval in warmup steps, e.g., when warmup_interval is 2, only 0, 2, 4, ... steps
+                    in warmup steps will be computed, others will use dynamic cache.
+                max_cached_steps (`int`, *required*, defaults to -1):
+                    DBCache disables the caching strategy when the previous cached steps exceed this value to
+                    prevent precision degradation.
+                max_continuous_cached_steps (`int`, *required*, defaults to -1):
+                    DBCache disables the caching strategy when the previous continous cached steps exceed this value to
+                    prevent precision degradation.
+                enable_separate_cfg (`bool`, *required*,  defaults to None):
+                    Whether to do separate cfg or not, such as Wan 2.1, Qwen-Image. For model that fused CFG
+                    and non-CFG into single forward step, should set enable_separate_cfg as False, for example:
+                    CogVideoX, HunyuanVideo, Mochi, etc.
+                cfg_compute_first (`bool`, *required*,  defaults to False):
+                    Whether to compute cfg forward first, default is False, meaning:
+                    0, 2, 4, ..., -> non-CFG step;
+                    1, 3, 5, ... -> CFG step.
+                cfg_diff_compute_separate (`bool`, *required*,  defaults to True):
+                    Whether to compute separate difference values for CFG and non-CFG steps, default is True.
+                    If False, we will use the computed difference from the current non-CFG transformer step
+                    for the current CFG step.
+                num_inference_steps (`int`, *optional*, defaults to None):
+                    num_inference_steps for DiffusionPipeline, used to adjust some internal settings
+                    for better caching performance. For example, we will refresh the cache once the
+                    executed steps exceed num_inference_steps if num_inference_steps is provided.
+
+        calibrator_config (`CalibratorConfig`, *optional*, defaults to None):
+            Config for calibrator. If calibrator_config is not None, it means the user wants to use DBCache
+            with a specific calibrator, such as taylorseer, foca, and so on.
+
+        params_modifiers ('ParamsModifier', *optional*, defaults to None):
+            Modify cache context params for specific blocks. The configurable params listed belows:
+                cache_config (`BasicCacheConfig`, *required*, defaults to BasicCacheConfig()):
+                    The same as 'cache_config' param in cache_dit.enable_cache() interface.
+                calibrator_config (`CalibratorConfig`, *optional*, defaults to None):
+                    The same as 'calibrator_config' param in cache_dit.enable_cache() interface.
+                **kwargs: (`dict`, *optional*, defaults to {}):
+                    The same as 'kwargs' param in cache_dit.enable_cache() interface.
+
+        parallelism_config (`ParallelismConfig`, *optional*, defaults to None):
+            Config for Parallelism. If parallelism_config is not None, it means the user wants to enable
+            parallelism for cache-dit. Please check https://github.com/vipshop/cache-dit/blob/main/src/cache_dit/parallelism/parallel_config.py
+            for more details of ParallelismConfig.
+                backend: (`ParallelismBackend`, *required*, defaults to "ParallelismBackend.NATIVE_DIFFUSER"):
+                    Parallelism backend, currently only NATIVE_DIFFUSER and NVTIVE_PYTORCH are supported.
+                    For context parallelism, only NATIVE_DIFFUSER backend is supported, for tensor parallelism,
+                    only NATIVE_PYTORCH backend is supported.
+                ulysses_size: (`int`, *optional*, defaults to None):
+                    The size of Ulysses cluster. If ulysses_size is not None, enable Ulysses style parallelism.
+                    This setting is only valid when backend is NATIVE_DIFFUSER.
+                ring_size: (`int`, *optional*, defaults to None):
+                    The size of ring for ring parallelism. If ring_size is not None, enable ring attention.
+                    This setting is only valid when backend is NATIVE_DIFFUSER.
+                tp_size: (`int`, *optional*, defaults to None):
+                    The size of tensor parallelism. If tp_size is not None, enable tensor parallelism.
+                    This setting is only valid when backend is NATIVE_PYTORCH.
+                parallel_kwargs: (`dict`, *optional*, defaults to {}):
+                    Additional kwargs for parallelism backends. For example, for NATIVE_DIFFUSER backend,
+                    it can include `cp_plan` and `attention_backend` arguments for `Context Parallelism`.
+
+        kwargs (`dict`, *optional*, defaults to {})
+            Other cache context kwargs, please check https://github.com/vipshop/cache-dit/blob/main/src/cache_dit/caching/cache_contexts/cache_context.py
+            for more details.
+
+    Examples:
+    ```py
+    >>> import cache_dit
+    >>> from diffusers import DiffusionPipeline
+    >>> pipe = DiffusionPipeline.from_pretrained("Qwen/Qwen-Image") # Can be any diffusion pipeline
+    >>> cache_dit.enable_cache(pipe) # One-line code with default cache options.
+    >>> output = pipe(...) # Just call the pipe as normal.
+    >>> stats = cache_dit.summary(pipe) # Then, get the summary of cache acceleration stats.
+    >>> cache_dit.disable_cache(pipe) # Disable cache and run original pipe.
+    """
+    # Precheck for compatibility of different configurations
+    if cache_config is None:
+        if parallelism_config is None:
+            # Set default cache config only when parallelism is not enabled
+            logger.info("cache_config is None, using default DBCacheConfig")
+            cache_config = DBCacheConfig()
+        else:
+            # Allow empty cache_config when parallelism is enabled
+            logger.warning(
+                "Parallelism is enabled and cache_config is None. Please manually "
+                "set cache_config to avoid potential compatibility issues. "
+                "Otherwise, cache will not be enabled."
+            )
+
+    # Collect cache context kwargs
+    context_kwargs = {}
+    if (cache_type := context_kwargs.get("cache_type", None)) is not None:
+        if cache_type == CacheType.NONE:
+            return pipe_or_adapter
+
+    # NOTE: Deprecated cache config params. These parameters are now retained
+    # for backward compatibility but will be removed in the future.
+    deprecated_kwargs = {
+        "Fn_compute_blocks": kwargs.get("Fn_compute_blocks", None),
+        "Bn_compute_blocks": kwargs.get("Bn_compute_blocks", None),
+        "max_warmup_steps": kwargs.get("max_warmup_steps", None),
+        "max_cached_steps": kwargs.get("max_cached_steps", None),
+        "max_continuous_cached_steps": kwargs.get(
+            "max_continuous_cached_steps", None
+        ),
+        "residual_diff_threshold": kwargs.get("residual_diff_threshold", None),
+        "enable_separate_cfg": kwargs.get("enable_separate_cfg", None),
+        "cfg_compute_first": kwargs.get("cfg_compute_first", None),
+        "cfg_diff_compute_separate": kwargs.get(
+            "cfg_diff_compute_separate", None
+        ),
+    }
+
+    deprecated_kwargs = {
+        k: v for k, v in deprecated_kwargs.items() if v is not None
+    }
+
+    if deprecated_kwargs:
+        logger.warning(
+            "Manually settup DBCache context without BasicCacheConfig is "
+            "deprecated and will be removed in the future, please use "
+            "`cache_config` parameter instead!"
+        )
+        if cache_config is not None:
+            cache_config.update(**deprecated_kwargs)
+        else:
+            cache_config = BasicCacheConfig(**deprecated_kwargs)
+
+    if cache_config is not None:
+        context_kwargs["cache_config"] = cache_config
+
+    # NOTE: Deprecated taylorseer params. These parameters are now retained
+    # for backward compatibility but will be removed in the future.
+    if cache_config is not None and (
+        kwargs.get("enable_taylorseer", None) is not None
+        or kwargs.get("enable_encoder_taylorseer", None) is not None
+    ):
+        logger.warning(
+            "Manually settup TaylorSeer calibrator without TaylorSeerCalibratorConfig is "
+            "deprecated and will be removed in the future, please use "
+            "`calibrator_config` parameter instead!"
+        )
+        from cache_dit.caching.cache_contexts.calibrators import (
+            TaylorSeerCalibratorConfig,
+        )
+
+        calibrator_config = TaylorSeerCalibratorConfig(
+            enable_calibrator=kwargs.get("enable_taylorseer"),
+            enable_encoder_calibrator=kwargs.get("enable_encoder_taylorseer"),
+            calibrator_cache_type=kwargs.get(
+                "taylorseer_cache_type", "residual"
+            ),
+            taylorseer_order=kwargs.get("taylorseer_order", 1),
+        )
+
+    if calibrator_config is not None:
+        context_kwargs["calibrator_config"] = calibrator_config
+
+    if params_modifiers is not None:
+        context_kwargs["params_modifiers"] = params_modifiers
+
+    if cache_config is not None:
+        if isinstance(
+            pipe_or_adapter,
+            (DiffusionPipeline, BlockAdapter, torch.nn.Module, ModelMixin),
+        ):
+            pipe_or_adapter = CachedAdapter.apply(
+                pipe_or_adapter,
+                **context_kwargs,
+            )
+        else:
+            raise ValueError(
+                f"type: {type(pipe_or_adapter)} is not valid, "
+                "Please pass DiffusionPipeline or BlockAdapter"
+                "for the 1's position param: pipe_or_adapter"
+            )
+    else:
+        logger.warning(
+            "cache_config is None, skip enabling cache for "
+            f"{pipe_or_adapter.__class__.__name__}."
+        )
+
+    # NOTE: Users should always enable parallelism after applying
+    # cache to avoid hooks conflict.
+    if parallelism_config is not None:
+        assert isinstance(
+            parallelism_config, ParallelismConfig
+        ), "parallelism_config should be of type ParallelismConfig."
+
+        transformers = []
+        if isinstance(pipe_or_adapter, DiffusionPipeline):
+            adapter = BlockAdapterRegistry.get_adapter(pipe_or_adapter)
+            if adapter is None:
+                assert hasattr(pipe_or_adapter, "transformer"), (
+                    "The given DiffusionPipeline does not have "
+                    "a 'transformer' attribute, cannot enable "
+                    "parallelism."
+                )
+                transformers = [pipe_or_adapter.transformer]
+            else:
+                adapter = BlockAdapter.normalize(adapter, unique=False)
+                transformers = BlockAdapter.flatten(adapter.transformer)
+        else:
+            if not BlockAdapter.is_normalized(pipe_or_adapter):
+                pipe_or_adapter = BlockAdapter.normalize(
+                    pipe_or_adapter, unique=False
+                )
+            transformers = BlockAdapter.flatten(pipe_or_adapter.transformer)
+
+        if len(transformers) == 0:
+            logger.warning(
+                "No transformer is detected in the "
+                "BlockAdapter, skip enabling parallelism."
+            )
+            return pipe_or_adapter
+
+        if len(transformers) > 1:
+            logger.warning(
+                "Multiple transformers are detected in the "
+                "BlockAdapter, all transfomers will be "
+                "enabled for parallelism."
+            )
+        for i, transformer in enumerate(transformers):
+            # Enable parallelism for the transformer inplace
+            transformers[i] = enable_parallelism(
+                transformer, parallelism_config
+            )
+    return pipe_or_adapter
+
+
+def disable_cache(
+    pipe_or_adapter: Union[
+        DiffusionPipeline,
+        BlockAdapter,
+    ],
+):
+    CachedAdapter.maybe_release_hooks(pipe_or_adapter)
+    logger.warning(
+        f"Cache Acceleration is disabled for: "
+        f"{pipe_or_adapter.__class__.__name__}."
+    )
+
+
+def supported_pipelines(
+    **kwargs,
+) -> Tuple[int, List[str]]:
+    return BlockAdapterRegistry.supported_pipelines(**kwargs)
+
+
+def get_adapter(
+    pipe: DiffusionPipeline | str | Any,
+) -> BlockAdapter:
+    return BlockAdapterRegistry.get_adapter(pipe)

cache_dit/{cache_factory → caching}/cache_types.py

@@ -6,7 +6,8 @@ logger = init_logger(__name__)
 
 class CacheType(Enum):
     NONE = "NONE"
-    DBCache = "Dual_Block_Cache"
+    DBCache = "DBCache"  # "Dual_Block_Cache"
+    DBPrune = "DBPrune"  # "Dynamic_Block_Prune"
 
     @staticmethod
     def type(type_hint: "CacheType | str") -> "CacheType":
@@ -14,6 +15,9 @@ class CacheType(Enum):
             return type_hint
         return cache_type(type_hint)
 
+    def __str__(self) -> str:
+        return self.value
+
 
 def cache_type(type_hint: "CacheType | str") -> "CacheType":
     if type_hint is None:
@@ -21,7 +25,6 @@ def cache_type(type_hint: "CacheType | str") -> "CacheType":
 
     if isinstance(type_hint, CacheType):
         return type_hint
-
     elif type_hint.upper() in (
         "DUAL_BLOCK_CACHE",
         "DB_CACHE",
@@ -29,6 +32,20 @@ def cache_type(type_hint: "CacheType | str") -> "CacheType":
         "DB",
     ):
         return CacheType.DBCache
+    elif type_hint.upper() in (
+        "DYNAMIC_BLOCK_PRUNE",
+        "DB_PRUNE",
+        "DBPRUNE",
+        "DBP",
+    ):
+        return CacheType.DBPrune
+    elif type_hint.upper() in (
+        "NONE",
+        "NO_CACHE",
+        "NOCACHE",
+        "NC",
+    ):
+        return CacheType.NONE
     return CacheType.NONE
 
 

cache_dit/{cache_factory → caching}/forward_pattern.py

@@ -20,33 +20,33 @@ class ForwardPattern(Enum):
 
     Pattern_0 = (
         True,  # Return_H_First
-        False, # Return_H_Only
-        False, # Forward_H_only
+        False,  # Return_H_Only
+        False,  # Forward_H_only
         ("hidden_states", "encoder_hidden_states"),  # In
         ("hidden_states", "encoder_hidden_states"),  # Out
         True,  # Supported
     )
 
     Pattern_1 = (
-        False, # Return_H_First
-        False, # Return_H_Only
-        False, # Forward_H_only
+        False,  # Return_H_First
+        False,  # Return_H_Only
+        False,  # Forward_H_only
         ("hidden_states", "encoder_hidden_states"),  # In
         ("encoder_hidden_states", "hidden_states"),  # Out
         True,  # Supported
     )
 
     Pattern_2 = (
-        False, # Return_H_First
+        False,  # Return_H_First
         True,  # Return_H_Only
-        False, # Forward_H_only
+        False,  # Forward_H_only
         ("hidden_states", "encoder_hidden_states"),  # In
-        ("hidden_states",),                          # Out
+        ("hidden_states",),  # Out
         True,  # Supported
     )
 
     Pattern_3 = (
-        False, # Return_H_First
+        False,  # Return_H_First
         True,  # Return_H_Only
         True,  # Forward_H_only
         ("hidden_states",),  # In
@@ -56,18 +56,18 @@ class ForwardPattern(Enum):
 
     Pattern_4 = (
         True,  # Return_H_First
-        False, # Return_H_Only
+        False,  # Return_H_Only
         True,  # Forward_H_only
-        ("hidden_states",),                          # In
+        ("hidden_states",),  # In
         ("hidden_states", "encoder_hidden_states"),  # Out
         True,  # Supported
     )
 
     Pattern_5 = (
-        False, # Return_H_First
-        False, # Return_H_Only
+        False,  # Return_H_First
+        False,  # Return_H_Only
         True,  # Forward_H_only
-        ("hidden_states",),                          # In
+        ("hidden_states",),  # In
         ("encoder_hidden_states", "hidden_states"),  # Out
         True,  # Supported
     )

cache_dit/{cache_factory → caching}/params_modifier.py

@@ -1,7 +1,7 @@
 from typing import Optional
 
-from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
-from cache_dit.cache_factory.cache_contexts import CalibratorConfig
+from cache_dit.caching.cache_contexts import BasicCacheConfig
+from cache_dit.caching.cache_contexts import CalibratorConfig
 
 from cache_dit.logger import init_logger
 
@@ -11,7 +11,7 @@ logger = init_logger(__name__)
 class ParamsModifier:
     def __init__(
         self,
-        # Basic DBCache config: BasicCacheConfig
+        # BasicCacheConfig, DBCacheConfig, DBPruneConfig, etc.
         cache_config: BasicCacheConfig = None,
         # Calibrator config: TaylorSeerCalibratorConfig, etc.
         calibrator_config: Optional[CalibratorConfig] = None,
@@ -22,7 +22,7 @@ class ParamsModifier:
 
         # WARNING: Deprecated cache config params. These parameters are now retained
         # for backward compatibility but will be removed in the future.
-        deprecated_cache_kwargs = {
+        deprecated_kwargs = {
             "Fn_compute_blocks": kwargs.get("Fn_compute_blocks", None),
             "Bn_compute_blocks": kwargs.get("Bn_compute_blocks", None),
             "max_warmup_steps": kwargs.get("max_warmup_steps", None),
@@ -40,20 +40,20 @@ class ParamsModifier:
             ),
         }
 
-        deprecated_cache_kwargs = {
-            k: v for k, v in deprecated_cache_kwargs.items() if v is not None
+        deprecated_kwargs = {
+            k: v for k, v in deprecated_kwargs.items() if v is not None
         }
 
-        if deprecated_cache_kwargs:
+        if deprecated_kwargs:
             logger.warning(
                 "Manually settup DBCache context without BasicCacheConfig is "
                 "deprecated and will be removed in the future, please use "
                 "`cache_config` parameter instead!"
             )
             if cache_config is not None:
-                cache_config.update(**deprecated_cache_kwargs)
+                cache_config.update(**deprecated_kwargs)
             else:
-                cache_config = BasicCacheConfig(**deprecated_cache_kwargs)
+                cache_config = BasicCacheConfig(**deprecated_kwargs)
 
         if cache_config is not None:
             self._context_kwargs["cache_config"] = cache_config
@@ -68,7 +68,7 @@ class ParamsModifier:
                 "deprecated and will be removed in the future, please use "
                 "`calibrator_config` parameter instead!"
             )
-            from cache_dit.cache_factory.cache_contexts.calibrators import (
+            from cache_dit.caching.cache_contexts.calibrators import (
                 TaylorSeerCalibratorConfig,
             )
 

cache_dit/caching/patch_functors/__init__.py

@@ -0,0 +1,15 @@
+from cache_dit.caching.patch_functors.functor_base import PatchFunctor
+from cache_dit.caching.patch_functors.functor_dit import DiTPatchFunctor
+from cache_dit.caching.patch_functors.functor_flux import FluxPatchFunctor
+from cache_dit.caching.patch_functors.functor_chroma import (
+    ChromaPatchFunctor,
+)
+from cache_dit.caching.patch_functors.functor_hidream import (
+    HiDreamPatchFunctor,
+)
+from cache_dit.caching.patch_functors.functor_hunyuan_dit import (
+    HunyuanDiTPatchFunctor,
+)
+from cache_dit.caching.patch_functors.functor_qwen_image_controlnet import (
+    QwenImageControlNetPatchFunctor,
+)

cache_dit/{cache_factory → caching}/patch_functors/functor_chroma.py

@@ -13,7 +13,7 @@ from diffusers.utils import (
     unscale_lora_layers,
 )
 
-from cache_dit.cache_factory.patch_functors.functor_base import (
+from cache_dit.caching.patch_functors.functor_base import (
     PatchFunctor,
 )
 from cache_dit.logger import init_logger

cache_dit/{cache_factory → caching}/patch_functors/functor_dit.py

@@ -6,7 +6,7 @@ from diffusers.models.transformers.dit_transformer_2d import (
     DiTTransformer2DModel,
     Transformer2DModelOutput,
 )
-from cache_dit.cache_factory.patch_functors.functor_base import (
+from cache_dit.caching.patch_functors.functor_base import (
     PatchFunctor,
 )
 from cache_dit.logger import init_logger

cache_dit/{cache_factory → caching}/patch_functors/functor_flux.py

@@ -14,7 +14,7 @@ from diffusers.utils import (
     unscale_lora_layers,
 )
 
-from cache_dit.cache_factory.patch_functors.functor_base import (
+from cache_dit.caching.patch_functors.functor_base import (
     PatchFunctor,
 )
 from cache_dit.logger import init_logger

cache_dit/{cache_factory → caching}/patch_functors/functor_hidream.py

@@ -13,7 +13,7 @@ from diffusers.utils import (
     scale_lora_layers,
     unscale_lora_layers,
 )
-from cache_dit.cache_factory.patch_functors.functor_base import (
+from cache_dit.caching.patch_functors.functor_base import (
     PatchFunctor,
 )
 from cache_dit.logger import init_logger
@@ -362,9 +362,7 @@ def __patch_transformer_forward__(
     )
     if hidden_states_masks is not None:
         # NOTE: Patched
-        cur_llama31_encoder_hidden_states = llama31_encoder_hidden_states[
-            self.double_stream_blocks[-1].block._block_id
-        ]
+        cur_llama31_encoder_hidden_states = llama31_encoder_hidden_states[0]
         encoder_attention_mask_ones = torch.ones(
             (
                 batch_size,

cache_dit/{cache_factory → caching}/patch_functors/functor_hunyuan_dit.py

@@ -5,7 +5,7 @@ from diffusers.models.transformers.hunyuan_transformer_2d import (
     HunyuanDiTBlock,
     Transformer2DModelOutput,
 )
-from cache_dit.cache_factory.patch_functors.functor_base import (
+from cache_dit.caching.patch_functors.functor_base import (
     PatchFunctor,
 )
 from cache_dit.logger import init_logger