cache-dit 1.0.0__py3-none-any.whl → 1.0.2__py3-none-any.whl

cache_dit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '1.0.0'
-__version_tuple__ = version_tuple = (1, 0, 0)
+__version__ = version = '1.0.2'
+__version_tuple__ = version_tuple = (1, 0, 2)
 
 __commit_id__ = commit_id = None

cache_dit/cache_factory/block_adapters/__init__.py

@@ -143,14 +143,30 @@ def qwenimage_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import QwenImageTransformer2DModel
 
     assert isinstance(pipe.transformer, QwenImageTransformer2DModel)
-    return BlockAdapter(
-        pipe=pipe,
-        transformer=pipe.transformer,
-        blocks=pipe.transformer.transformer_blocks,
-        forward_pattern=ForwardPattern.Pattern_1,
-        has_separate_cfg=True,
-        **kwargs,
-    )
+
+    pipe_cls_name: str = pipe.__class__.__name__
+    if pipe_cls_name.startswith("QwenImageControlNet"):
+        from cache_dit.cache_factory.patch_functors import (
+            QwenImageControlNetPatchFunctor,
+        )
+
+        return BlockAdapter(
+            pipe=pipe,
+            transformer=pipe.transformer,
+            blocks=pipe.transformer.transformer_blocks,
+            forward_pattern=ForwardPattern.Pattern_1,
+            patch_functor=QwenImageControlNetPatchFunctor(),
+            has_separate_cfg=True,
+        )
+    else:
+        return BlockAdapter(
+            pipe=pipe,
+            transformer=pipe.transformer,
+            blocks=pipe.transformer.transformer_blocks,
+            forward_pattern=ForwardPattern.Pattern_1,
+            has_separate_cfg=True,
+            **kwargs,
+        )
 
 
 @BlockAdapterRegistry.register("LTX")

cache_dit/cache_factory/cache_adapters/cache_adapter.py

@@ -14,10 +14,6 @@ from cache_dit.cache_factory.cache_contexts import CachedContextManager
 from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
 from cache_dit.cache_factory.cache_contexts import CalibratorConfig
 from cache_dit.cache_factory.cache_blocks import CachedBlocks
-from cache_dit.cache_factory.cache_blocks import (
-    patch_cached_stats,
-    remove_cached_stats,
-)
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
@@ -167,7 +163,7 @@ class CachedAdapter:
         cls,
         block_adapter: BlockAdapter,
         **cache_context_kwargs,
-    ) -> DiffusionPipeline:
+    ) -> Tuple[List[str], List[Dict[str, Any]]]:
 
         BlockAdapter.assert_normalized(block_adapter)
 
@@ -221,7 +217,7 @@ class CachedAdapter:
 
         cls.apply_params_hooks(block_adapter, contexts_kwargs)
 
-        return block_adapter.pipe
+        return flatten_contexts, contexts_kwargs
 
     @classmethod
     def modify_context_params(
@@ -470,6 +466,10 @@ class CachedAdapter:
         cls,
         block_adapter: BlockAdapter,
     ):
+        from cache_dit.cache_factory.cache_blocks import (
+            patch_cached_stats,
+        )
+
         cache_manager = block_adapter.pipe._cache_manager
 
         for i in range(len(block_adapter.transformer)):
@@ -557,6 +557,10 @@ class CachedAdapter:
         )
 
         # release stats hooks
+        from cache_dit.cache_factory.cache_blocks import (
+            remove_cached_stats,
+        )
+
         cls.release_hooks(
             pipe_or_adapter,
             remove_cached_stats,

cache_dit/cache_factory/cache_blocks/pattern_3_4_5.py

@@ -1,6 +1,9 @@
 import torch
 
 from cache_dit.cache_factory import ForwardPattern
+from cache_dit.cache_factory.cache_contexts.cache_manager import (
+    CacheNotExistError,
+)
 from cache_dit.cache_factory.cache_blocks.pattern_base import (
     CachedBlocks_Pattern_Base,
 )
@@ -16,6 +19,70 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
         ForwardPattern.Pattern_5,
     ]
 
+    def call_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        # Call all blocks to process the hidden states without cache.
+        new_encoder_hidden_states = None
+        for block in self.transformer_blocks:
+            hidden_states = block(
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+            hidden_states, new_encoder_hidden_states = self._process_outputs(
+                hidden_states
+            )
+
+        return hidden_states, new_encoder_hidden_states
+
+    @torch.compiler.disable
+    def _process_outputs(
+        self, hidden_states: torch.Tensor | tuple
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        # Process the outputs for the block.
+        new_encoder_hidden_states = None
+        if not isinstance(hidden_states, torch.Tensor):  # Pattern 4, 5
+            if len(hidden_states) == 2:
+                if isinstance(hidden_states[1], torch.Tensor):
+                    hidden_states, new_encoder_hidden_states = hidden_states
+                    if not self.forward_pattern.Return_H_First:
+                        hidden_states, new_encoder_hidden_states = (
+                            new_encoder_hidden_states,
+                            hidden_states,
+                        )
+                elif isinstance(hidden_states[0], torch.Tensor):
+                    hidden_states = hidden_states[0]
+                else:
+                    raise ValueError("Unexpected hidden_states format.")
+            else:
+                assert (
+                    len(hidden_states) == 1
+                ), f"Unexpected output length: {len(hidden_states)}"
+                hidden_states = hidden_states[0]
+        return hidden_states, new_encoder_hidden_states
+
+    @torch.compiler.disable
+    def _forward_outputs(
+        self,
+        hidden_states: torch.Tensor,
+        new_encoder_hidden_states: torch.Tensor | None,
+    ) -> (
+        torch.Tensor
+        | tuple[torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, None]
+    ):
+        if self.forward_pattern.Return_H_Only:
+            return hidden_states
+        else:
+            if self.forward_pattern.Return_H_First:
+                return (hidden_states, new_encoder_hidden_states)
+            else:
+                return (new_encoder_hidden_states, hidden_states)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -23,8 +90,19 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
         **kwargs,
     ):
         # Use it's own cache context.
-        self.cache_manager.set_context(self.cache_context)
-        self._check_cache_params()
+        try:
+            self.cache_manager.set_context(self.cache_context)
+            self._check_cache_params()
+        except CacheNotExistError as e:
+            logger.warning(f"Cache context not exist: {e}, skip cache.")
+            hidden_states, new_encoder_hidden_states = self.call_blocks(
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+            return self._forward_outputs(
+                hidden_states, new_encoder_hidden_states
+            )
 
         original_hidden_states = hidden_states
         # Call first `n` blocks to process the hidden states for
@@ -35,7 +113,9 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
             **kwargs,
         )
 
-        Fn_hidden_states_residual = hidden_states - original_hidden_states
+        Fn_hidden_states_residual = hidden_states - original_hidden_states.to(
+            hidden_states.device
+        )
         del original_hidden_states
 
         self.cache_manager.mark_step_begin()
@@ -147,15 +227,7 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
 
         torch._dynamo.graph_break()
 
-        return (
-            hidden_states
-            if self.forward_pattern.Return_H_Only
-            else (
-                (hidden_states, new_encoder_hidden_states)
-                if self.forward_pattern.Return_H_First
-                else (new_encoder_hidden_states, hidden_states)
-            )
-        )
+        return self._forward_outputs(hidden_states, new_encoder_hidden_states)
 
     def call_Fn_blocks(
         self,
@@ -170,13 +242,9 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 *args,
                 **kwargs,
             )
-            if not isinstance(hidden_states, torch.Tensor):  # Pattern 4, 5
-                hidden_states, new_encoder_hidden_states = hidden_states
-                if not self.forward_pattern.Return_H_First:
-                    hidden_states, new_encoder_hidden_states = (
-                        new_encoder_hidden_states,
-                        hidden_states,
-                    )
+            hidden_states, new_encoder_hidden_states = self._process_outputs(
+                hidden_states
+            )
 
         return hidden_states, new_encoder_hidden_states
 
@@ -194,16 +262,16 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 *args,
                 **kwargs,
             )
-            if not isinstance(hidden_states, torch.Tensor):  # Pattern 4, 5
-                hidden_states, new_encoder_hidden_states = hidden_states
-                if not self.forward_pattern.Return_H_First:
-                    hidden_states, new_encoder_hidden_states = (
-                        new_encoder_hidden_states,
-                        hidden_states,
-                    )
+
+            hidden_states, new_encoder_hidden_states = self._process_outputs(
+                hidden_states
+            )
+
         # compute hidden_states residual
         hidden_states = hidden_states.contiguous()
-        hidden_states_residual = hidden_states - original_hidden_states
+        hidden_states_residual = hidden_states - original_hidden_states.to(
+            hidden_states.device
+        )
 
         return (
             hidden_states,
@@ -227,12 +295,9 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 *args,
                 **kwargs,
             )
-            if not isinstance(hidden_states, torch.Tensor):  # Pattern 4,5
-                hidden_states, new_encoder_hidden_states = hidden_states
-                if not self.forward_pattern.Return_H_First:
-                    hidden_states, new_encoder_hidden_states = (
-                        new_encoder_hidden_states,
-                        hidden_states,
-                    )
+
+            hidden_states, new_encoder_hidden_states = self._process_outputs(
+                hidden_states
+            )
 
         return hidden_states, new_encoder_hidden_states

cache_dit/cache_factory/cache_blocks/pattern_base.py

@@ -1,12 +1,11 @@
 import inspect
-import asyncio
 import torch
 import torch.distributed as dist
 
-from typing import List
 from cache_dit.cache_factory.cache_contexts.cache_context import CachedContext
 from cache_dit.cache_factory.cache_contexts.cache_manager import (
     CachedContextManager,
+    CacheNotExistError,
 )
 from cache_dit.cache_factory import ForwardPattern
 from cache_dit.logger import init_logger
@@ -47,7 +46,6 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         self.cache_prefix = cache_prefix
         self.cache_context = cache_context
         self.cache_manager = cache_manager
-        self.pending_tasks: List[asyncio.Task] = []
 
         self._check_forward_pattern()
         logger.info(
@@ -111,6 +109,62 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
             f"the number of transformer blocks {len(self.transformer_blocks)}"
         )
 
+    def call_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        # Call all blocks to process the hidden states without cache.
+        for block in self.transformer_blocks:
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+            if not isinstance(hidden_states, torch.Tensor):
+                hidden_states, encoder_hidden_states = hidden_states
+                if not self.forward_pattern.Return_H_First:
+                    hidden_states, encoder_hidden_states = (
+                        encoder_hidden_states,
+                        hidden_states,
+                    )
+
+        return hidden_states, encoder_hidden_states
+
+    @torch.compiler.disable
+    def _process_outputs(
+        self,
+        hidden_states: torch.Tensor | tuple,
+        encoder_hidden_states: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if not isinstance(hidden_states, torch.Tensor):
+            hidden_states, encoder_hidden_states = hidden_states
+            if not self.forward_pattern.Return_H_First:
+                hidden_states, encoder_hidden_states = (
+                    encoder_hidden_states,
+                    hidden_states,
+                )
+        return hidden_states, encoder_hidden_states
+
+    @torch.compiler.disable
+    def _forward_outputs(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
+        return (
+            hidden_states
+            if self.forward_pattern.Return_H_Only
+            else (
+                (hidden_states, encoder_hidden_states)
+                if self.forward_pattern.Return_H_First
+                else (encoder_hidden_states, hidden_states)
+            )
+        )
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -119,8 +173,19 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         **kwargs,
     ):
         # Use it's own cache context.
-        self.cache_manager.set_context(self.cache_context)
-        self._check_cache_params()
+        try:
+            self.cache_manager.set_context(self.cache_context)
+            self._check_cache_params()
+        except CacheNotExistError as e:
+            logger.warning(f"Cache context not exist: {e}, skip cache.")
+            # Call all blocks to process the hidden states.
+            hidden_states, encoder_hidden_states = self.call_blocks(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+            return self._forward_outputs(hidden_states, encoder_hidden_states)
 
         original_hidden_states = hidden_states
         # Call first `n` blocks to process the hidden states for
@@ -239,15 +304,7 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         # patch cached stats for blocks or remove it.
         torch._dynamo.graph_break()
 
-        return (
-            hidden_states
-            if self.forward_pattern.Return_H_Only
-            else (
-                (hidden_states, encoder_hidden_states)
-                if self.forward_pattern.Return_H_First
-                else (encoder_hidden_states, hidden_states)
-            )
-        )
+        return self._forward_outputs(hidden_states, encoder_hidden_states)
 
     @torch.compiler.disable
     def _is_parallelized(self):
@@ -322,13 +379,9 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
                 *args,
                 **kwargs,
             )
-            if not isinstance(hidden_states, torch.Tensor):
-                hidden_states, encoder_hidden_states = hidden_states
-                if not self.forward_pattern.Return_H_First:
-                    hidden_states, encoder_hidden_states = (
-                        encoder_hidden_states,
-                        hidden_states,
-                    )
+            hidden_states, encoder_hidden_states = self._process_outputs(
+                hidden_states, encoder_hidden_states
+            )
 
         return hidden_states, encoder_hidden_states
 
@@ -348,13 +401,9 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
                 *args,
                 **kwargs,
             )
-            if not isinstance(hidden_states, torch.Tensor):
-                hidden_states, encoder_hidden_states = hidden_states
-                if not self.forward_pattern.Return_H_First:
-                    hidden_states, encoder_hidden_states = (
-                        encoder_hidden_states,
-                        hidden_states,
-                    )
+            hidden_states, encoder_hidden_states = self._process_outputs(
+                hidden_states, encoder_hidden_states
+            )
 
         # compute hidden_states residual
         hidden_states = hidden_states.contiguous()
@@ -396,12 +445,8 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
                 *args,
                 **kwargs,
             )
-            if not isinstance(hidden_states, torch.Tensor):
-                hidden_states, encoder_hidden_states = hidden_states
-                if not self.forward_pattern.Return_H_First:
-                    hidden_states, encoder_hidden_states = (
-                        encoder_hidden_states,
-                        hidden_states,
-                    )
+            hidden_states, encoder_hidden_states = self._process_outputs(
+                hidden_states, encoder_hidden_states
+            )
 
         return hidden_states, encoder_hidden_states

cache_dit/cache_factory/cache_contexts/__init__.py

@@ -11,4 +11,5 @@ from cache_dit.cache_factory.cache_contexts.cache_context import (
 )
 from cache_dit.cache_factory.cache_contexts.cache_manager import (
     CachedContextManager,
+    CacheNotExistError,
 )

cache_dit/cache_factory/cache_contexts/cache_manager.py

@@ -14,6 +14,10 @@ from cache_dit.logger import init_logger
 logger = init_logger(__name__)
 
 
+class CacheNotExistError(Exception):
+    pass
+
+
 class CachedContextManager:
     # Each Pipeline should have it's own context manager instance.
 
@@ -27,16 +31,19 @@ class CachedContextManager:
         self._cached_context_manager[_context.name] = _context
         return _context
 
-    def set_context(self, cached_context: CachedContext | str):
+    def set_context(self, cached_context: CachedContext | str) -> CachedContext:
         if isinstance(cached_context, CachedContext):
             self._current_context = cached_context
         else:
+            if cached_context not in self._cached_context_manager:
+                raise CacheNotExistError("Context not exist!")
             self._current_context = self._cached_context_manager[cached_context]
+        return self._current_context
 
     def get_context(self, name: str = None) -> CachedContext:
         if name is not None:
             if name not in self._cached_context_manager:
-                raise ValueError("Context not exist!")
+                raise CacheNotExistError("Context not exist!")
             return self._cached_context_manager[name]
         return self._current_context
 

cache_dit/cache_factory/cache_interface.py

@@ -38,23 +38,43 @@ def enable_cache(
     BlockAdapter,
 ]:
     r"""
-    Unified Cache API for  almost Any Diffusion Transformers (with Transformer Blocks
-    that match the specific Input and Output patterns).
-
-    For a good balance between performance and precision, DBCache is configured by default
-    with F8B0, 8 warmup steps, and unlimited cached steps.
+    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` or `BlockAdapter`, *required*):
             The standard Diffusion Pipeline or custom BlockAdapter (from cache-dit or user-defined).
             For example: cache_dit.enable_cache(FluxPipeline(...)). Please check https://github.com/vipshop/cache-dit/blob/main/docs/BlockAdapter.md
             for the usgae of BlockAdapter.
+
         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 diff and delivering more
-                    accurate information to subsequent blocks. Please check https://github.com/vipshop/cache-dit/blob/main/docs/DBCache.md
+                    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
@@ -77,14 +97,18 @@ def enable_cache(
                     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):
-                    Compute cfg forward first or not, default False, namely, 0, 2, 4, ..., -> non-CFG step;
+                    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):
-                    Compute separate diff values for CFG and non-CFG step, default True. If False, we will
-                    use the computed diff from current non-CFG transformer step for current CFG step.
+                    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.
+
         calibrator_config (`CalibratorConfig`, *optional*, defaults to None):
-            Config for calibrator, if calibrator_config is not None, means that user want to use DBCache
-            with specific calibrator, such as taylorseer, foca, and so on.
+            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()):
@@ -93,6 +117,7 @@ def enable_cache(
                     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.
+
         kwargs (`dict`, *optional*, defaults to {})
             Other cache context kwargs, please check https://github.com/vipshop/cache-dit/blob/main/src/cache_dit/cache_factory/cache_contexts/cache_context.py
             for more details.

cache_dit/cache_factory/patch_functors/__init__.py

@@ -10,3 +10,6 @@ from cache_dit.cache_factory.patch_functors.functor_hidream import (
 from cache_dit.cache_factory.patch_functors.functor_hunyuan_dit import (
     HunyuanDiTPatchFunctor,
 )
+from cache_dit.cache_factory.patch_functors.functor_qwen_image_controlnet import (
+    QwenImageControlNetPatchFunctor,
+)

cache_dit/cache_factory/patch_functors/functor_hidream.py

@@ -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,