cache-dit 0.2.29__py3-none-any.whl → 0.2.31__py3-none-any.whl

cache_dit/cache_factory/cache_blocks/pattern_3_4_5.py

@@ -1,5 +1,6 @@
 import torch
 
+from typing import Dict, Any
 from cache_dit.cache_factory import ForwardPattern
 from cache_dit.cache_factory.cache_blocks.pattern_base import (
     CachedBlocks_Pattern_Base,
@@ -31,7 +32,7 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
         # Call first `n` blocks to process the hidden states for
         # more stable diff calculation.
         # encoder_hidden_states: None Pattern 3, else 4, 5
-        hidden_states, encoder_hidden_states = self.call_Fn_blocks(
+        hidden_states, new_encoder_hidden_states = self.call_Fn_blocks(
             hidden_states,
             *args,
             **kwargs,
@@ -60,11 +61,10 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
         if can_use_cache:
             self.cache_manager.add_cached_step()
             del Fn_hidden_states_residual
-            hidden_states, encoder_hidden_states = (
+            hidden_states, new_encoder_hidden_states = (
                 self.cache_manager.apply_cache(
                     hidden_states,
-                    # None Pattern 3, else 4, 5
-                    encoder_hidden_states,
+                    new_encoder_hidden_states,  # encoder_hidden_states not use cache
                     prefix=(
                         f"{self.cache_prefix}_Bn_residual"
                         if self.cache_manager.is_cache_residual()
@@ -80,12 +80,12 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
             torch._dynamo.graph_break()
             # Call last `n` blocks to further process the hidden states
             # for higher precision.
-            hidden_states, encoder_hidden_states = self.call_Bn_blocks(
-                hidden_states,
-                encoder_hidden_states,
-                *args,
-                **kwargs,
-            )
+            if self.cache_manager.Bn_compute_blocks() > 0:
+                hidden_states, new_encoder_hidden_states = self.call_Bn_blocks(
+                    hidden_states,
+                    *args,
+                    **kwargs,
+                )
         else:
             self.cache_manager.set_Fn_buffer(
                 Fn_hidden_states_residual,
@@ -99,19 +99,20 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 )
             del Fn_hidden_states_residual
             torch._dynamo.graph_break()
+            old_encoder_hidden_states = new_encoder_hidden_states
             (
                 hidden_states,
-                encoder_hidden_states,
+                new_encoder_hidden_states,
                 hidden_states_residual,
-                # None Pattern 3, else 4, 5
-                encoder_hidden_states_residual,
             ) = self.call_Mn_blocks(  # middle
                 hidden_states,
-                # None Pattern 3, else 4, 5
-                encoder_hidden_states,
                 *args,
                 **kwargs,
             )
+            if new_encoder_hidden_states is not None:
+                new_encoder_hidden_states_residual = (
+                    new_encoder_hidden_states - old_encoder_hidden_states
+                )
             torch._dynamo.graph_break()
             if self.cache_manager.is_cache_residual():
                 self.cache_manager.set_Bn_buffer(
@@ -119,34 +120,32 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                     prefix=f"{self.cache_prefix}_Bn_residual",
                 )
             else:
-                # TaylorSeer
                 self.cache_manager.set_Bn_buffer(
                     hidden_states,
                     prefix=f"{self.cache_prefix}_Bn_hidden_states",
                 )
+
             if self.cache_manager.is_encoder_cache_residual():
-                self.cache_manager.set_Bn_encoder_buffer(
-                    # None Pattern 3, else 4, 5
-                    encoder_hidden_states_residual,
-                    prefix=f"{self.cache_prefix}_Bn_residual",
-                )
+                if new_encoder_hidden_states is not None:
+                    self.cache_manager.set_Bn_encoder_buffer(
+                        new_encoder_hidden_states_residual,
+                        prefix=f"{self.cache_prefix}_Bn_residual",
+                    )
             else:
-                # TaylorSeer
-                self.cache_manager.set_Bn_encoder_buffer(
-                    # None Pattern 3, else 4, 5
-                    encoder_hidden_states,
-                    prefix=f"{self.cache_prefix}_Bn_hidden_states",
-                )
+                if new_encoder_hidden_states is not None:
+                    self.cache_manager.set_Bn_encoder_buffer(
+                        new_encoder_hidden_states_residual,
+                        prefix=f"{self.cache_prefix}_Bn_hidden_states",
+                    )
             torch._dynamo.graph_break()
             # Call last `n` blocks to further process the hidden states
             # for higher precision.
-            hidden_states, encoder_hidden_states = self.call_Bn_blocks(
-                hidden_states,
-                # None Pattern 3, else 4, 5
-                encoder_hidden_states,
-                *args,
-                **kwargs,
-            )
+            if self.cache_manager.Bn_compute_blocks() > 0:
+                hidden_states, new_encoder_hidden_states = self.call_Bn_blocks(
+                    hidden_states,
+                    *args,
+                    **kwargs,
+                )
 
         torch._dynamo.graph_break()
 
@@ -154,12 +153,21 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
             hidden_states
             if self.forward_pattern.Return_H_Only
             else (
-                (hidden_states, encoder_hidden_states)
+                (hidden_states, new_encoder_hidden_states)
                 if self.forward_pattern.Return_H_First
-                else (encoder_hidden_states, hidden_states)
+                else (new_encoder_hidden_states, hidden_states)
             )
         )
 
+    @torch.compiler.disable
+    def maybe_update_kwargs(
+        self, encoder_hidden_states, kwargs: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        # if "encoder_hidden_states" in kwargs:
+        #     kwargs["encoder_hidden_states"] = encoder_hidden_states
+        # return kwargs
+        return kwargs
+
     def call_Fn_blocks(
         self,
         hidden_states: torch.Tensor,
@@ -172,7 +180,7 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
             f"Fn_compute_blocks {self.cache_manager.Fn_compute_blocks()} must be less than "
             f"the number of transformer blocks {len(self.transformer_blocks)}"
         )
-        encoder_hidden_states = None  # Pattern 3
+        new_encoder_hidden_states = None
         for block in self._Fn_blocks():
             hidden_states = block(
                 hidden_states,
@@ -180,25 +188,27 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 **kwargs,
             )
             if not isinstance(hidden_states, torch.Tensor):  # Pattern 4, 5
-                hidden_states, encoder_hidden_states = hidden_states
+                hidden_states, new_encoder_hidden_states = hidden_states
                 if not self.forward_pattern.Return_H_First:
-                    hidden_states, encoder_hidden_states = (
-                        encoder_hidden_states,
+                    hidden_states, new_encoder_hidden_states = (
+                        new_encoder_hidden_states,
                         hidden_states,
                     )
+            kwargs = self.maybe_update_kwargs(
+                new_encoder_hidden_states,
+                kwargs,
+            )
 
-        return hidden_states, encoder_hidden_states
+        return hidden_states, new_encoder_hidden_states
 
     def call_Mn_blocks(
         self,
         hidden_states: torch.Tensor,
-        # None Pattern 3, else 4, 5
-        encoder_hidden_states: torch.Tensor | None,
         *args,
         **kwargs,
     ):
         original_hidden_states = hidden_states
-        original_encoder_hidden_states = encoder_hidden_states
+        new_encoder_hidden_states = None
         for block in self._Mn_blocks():
             hidden_states = block(
                 hidden_states,
@@ -206,44 +216,33 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 **kwargs,
             )
             if not isinstance(hidden_states, torch.Tensor):  # Pattern 4, 5
-                hidden_states, encoder_hidden_states = hidden_states
+                hidden_states, new_encoder_hidden_states = hidden_states
                 if not self.forward_pattern.Return_H_First:
-                    hidden_states, encoder_hidden_states = (
-                        encoder_hidden_states,
+                    hidden_states, new_encoder_hidden_states = (
+                        new_encoder_hidden_states,
                         hidden_states,
                     )
+            kwargs = self.maybe_update_kwargs(
+                new_encoder_hidden_states,
+                kwargs,
+            )
 
         # compute hidden_states residual
         hidden_states = hidden_states.contiguous()
         hidden_states_residual = hidden_states - original_hidden_states
-        if (
-            original_encoder_hidden_states is not None
-            and encoder_hidden_states is not None
-        ):  # Pattern 4, 5
-            encoder_hidden_states_residual = (
-                encoder_hidden_states - original_encoder_hidden_states
-            )
-        else:
-            encoder_hidden_states_residual = None  # Pattern 3
 
         return (
             hidden_states,
-            encoder_hidden_states,
+            new_encoder_hidden_states,
             hidden_states_residual,
-            encoder_hidden_states_residual,
         )
 
     def call_Bn_blocks(
         self,
         hidden_states: torch.Tensor,
-        # None Pattern 3, else 4, 5
-        encoder_hidden_states: torch.Tensor | None,
         *args,
         **kwargs,
     ):
-        if self.cache_manager.Bn_compute_blocks() == 0:
-            return hidden_states, encoder_hidden_states
-
         assert self.cache_manager.Bn_compute_blocks() <= len(
             self.transformer_blocks
         ), (
@@ -264,11 +263,15 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                     **kwargs,
                 )
                 if not isinstance(hidden_states, torch.Tensor):  # Pattern 4,5
-                    hidden_states, encoder_hidden_states = hidden_states
+                    hidden_states, new_encoder_hidden_states = hidden_states
                     if not self.forward_pattern.Return_H_First:
-                        hidden_states, encoder_hidden_states = (
-                            encoder_hidden_states,
+                        hidden_states, new_encoder_hidden_states = (
+                            new_encoder_hidden_states,
                             hidden_states,
                         )
+                kwargs = self.maybe_update_kwargs(
+                    new_encoder_hidden_states,
+                    kwargs,
+                )
 
-        return hidden_states, encoder_hidden_states
+        return hidden_states, new_encoder_hidden_states

cache_dit/cache_factory/cache_blocks/pattern_base.py

@@ -25,6 +25,7 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         transformer_blocks: torch.nn.ModuleList,
         transformer: torch.nn.Module = None,
         forward_pattern: ForwardPattern = ForwardPattern.Pattern_0,
+        check_forward_pattern: bool = True,
         check_num_outputs: bool = True,
         # 1. Cache context configuration
         cache_prefix: str = None,  # maybe un-need.
@@ -38,6 +39,7 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         self.transformer = transformer
         self.transformer_blocks = transformer_blocks
         self.forward_pattern = forward_pattern
+        self.check_forward_pattern = check_forward_pattern
         self.check_num_outputs = check_num_outputs
         # 1. Cache context configuration
         self.cache_prefix = cache_prefix
@@ -52,6 +54,12 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         )
 
     def _check_forward_pattern(self):
+        if not self.check_forward_pattern:
+            logger.warning(
+                f"Skipped Forward Pattern Check: {self.forward_pattern}"
+            )
+            return
+
         assert (
             self.forward_pattern.Supported
             and self.forward_pattern in self._supported_patterns
@@ -59,6 +67,11 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
 
         if self.transformer_blocks is not None:
             for block in self.transformer_blocks:
+                # Special case for HiDreamBlock
+                if hasattr(block, "block"):
+                    if isinstance(block.block, torch.nn.Module):
+                        block = block.block
+
                 forward_parameters = set(
                     inspect.signature(block.forward).parameters.keys()
                 )

cache_dit/cache_factory/cache_contexts/cache_manager.py

@@ -733,17 +733,15 @@ class CachedContextManager:
                     encoder_prefix
                 )
 
-            assert (
-                encoder_hidden_states_prev is not None
-            ), f"{prefix}_encoder_buffer must be set before"
+            if encoder_hidden_states_prev is not None:
 
-            if self.is_encoder_cache_residual():
-                encoder_hidden_states = (
-                    encoder_hidden_states_prev + encoder_hidden_states
-                )
-            else:
-                # If encoder cache is not residual, we use the encoder hidden states directly
-                encoder_hidden_states = encoder_hidden_states_prev
+                if self.is_encoder_cache_residual():
+                    encoder_hidden_states = (
+                        encoder_hidden_states_prev + encoder_hidden_states
+                    )
+                else:
+                    # If encoder cache is not residual, we use the encoder hidden states directly
+                    encoder_hidden_states = encoder_hidden_states_prev
 
             encoder_hidden_states = encoder_hidden_states.contiguous()
 

cache_dit/cache_factory/cache_interface.py

@@ -1,11 +1,9 @@
-import torch
-from typing import Any, Tuple, List
+from typing import Any, Tuple, List, Union
 from diffusers import DiffusionPipeline
 from cache_dit.cache_factory.cache_types import CacheType
 from cache_dit.cache_factory.block_adapters import BlockAdapter
 from cache_dit.cache_factory.block_adapters import BlockAdapterRegistry
 from cache_dit.cache_factory.cache_adapters import CachedAdapter
-from cache_dit.cache_factory.cache_contexts import CachedContextManager
 
 from cache_dit.logger import init_logger
 
@@ -14,7 +12,10 @@ logger = init_logger(__name__)
 
 def enable_cache(
     # DiffusionPipeline or BlockAdapter
-    pipe_or_adapter: DiffusionPipeline | BlockAdapter,
+    pipe_or_adapter: Union[
+        DiffusionPipeline,
+        BlockAdapter,
+    ],
     # Cache context kwargs
     Fn_compute_blocks: int = 8,
     Bn_compute_blocks: int = 0,
@@ -32,7 +33,10 @@ def enable_cache(
     taylorseer_cache_type: str = "residual",
     taylorseer_order: int = 2,
     **other_cache_context_kwargs,
-) -> BlockAdapter:
+) -> Union[
+    DiffusionPipeline,
+    BlockAdapter,
+]:
     r"""
     Unified Cache API for  almost Any Diffusion Transformers (with Transformer Blocks
     that match the specific Input and Output patterns).
@@ -102,11 +106,11 @@ def enable_cache(
     >>> 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.
     """
-
     # Collect cache context kwargs
     cache_context_kwargs = other_cache_context_kwargs.copy()
-    if cache_type := cache_context_kwargs.get("cache_type", None):
+    if (cache_type := cache_context_kwargs.get("cache_type", None)) is not None:
         if cache_type == CacheType.NONE:
             return pipe_or_adapter
 
@@ -145,79 +149,17 @@ def enable_cache(
 
 
 def disable_cache(
-    # DiffusionPipeline or BlockAdapter
-    pipe_or_adapter: DiffusionPipeline | BlockAdapter,
+    pipe_or_adapter: Union[
+        DiffusionPipeline,
+        BlockAdapter,
+    ],
 ):
-    from cache_dit.cache_factory.cache_blocks.utils import (
-        remove_cached_stats,
+    CachedAdapter.maybe_release_hooks(pipe_or_adapter)
+    logger.warning(
+        f"Cache Acceleration is disabled for: "
+        f"{pipe_or_adapter.__class__.__name__}."
     )
 
-    def _disable_blocks(blocks: torch.nn.ModuleList):
-        if blocks is None:
-            return
-        if hasattr(blocks, "_forward_pattern"):
-            del blocks._forward_pattern
-        if hasattr(blocks, "_cache_context_kwargs"):
-            del blocks._cache_context_kwargs
-        remove_cached_stats(blocks)
-
-    def _disable_transformer(transformer: torch.nn.Module):
-        if transformer is None or not BlockAdapter.is_cached(transformer):
-            return
-        if original_forward := getattr(transformer, "_original_forward"):
-            transformer.forward = original_forward.__get__(transformer)
-            del transformer._original_forward
-        if hasattr(transformer, "_is_cached"):
-            del transformer._is_cached
-        if hasattr(transformer, "_forward_pattern"):
-            del transformer._forward_pattern
-        if hasattr(transformer, "_has_separate_cfg"):
-            del transformer._has_separate_cfg
-        if hasattr(transformer, "_cache_context_kwargs"):
-            del transformer._cache_context_kwargs
-        remove_cached_stats(transformer)
-        for blocks in BlockAdapter.find_blocks(transformer):
-            _disable_blocks(blocks)
-
-    def _disable_pipe(pipe: DiffusionPipeline):
-        if pipe is None or not BlockAdapter.is_cached(pipe):
-            return
-        if original_call := getattr(pipe, "_original_call"):
-            pipe.__class__.__call__ = original_call
-            del pipe.__class__._original_call
-        if cache_manager := getattr(pipe, "_cache_manager"):
-            assert isinstance(cache_manager, CachedContextManager)
-            cache_manager.clear_contexts()
-            del pipe._cache_manager
-        if hasattr(pipe, "_is_cached"):
-            del pipe.__class__._is_cached
-        if hasattr(pipe, "_cache_context_kwargs"):
-            del pipe._cache_context_kwargs
-        remove_cached_stats(pipe)
-
-    if isinstance(pipe_or_adapter, DiffusionPipeline):
-        pipe = pipe_or_adapter
-        _disable_pipe(pipe)
-        if hasattr(pipe, "transformer"):
-            _disable_transformer(pipe.transformer)
-        if hasattr(pipe, "transformer_2"):  # Wan 2.2
-            _disable_transformer(pipe.transformer_2)
-        pipe_cls_name = pipe.__class__.__name__
-        logger.warning(f"Cache Acceleration is disabled for: {pipe_cls_name}")
-    elif isinstance(pipe_or_adapter, BlockAdapter):
-        # BlockAdapter
-        adapter = pipe_or_adapter
-        BlockAdapter.assert_normalized(adapter)
-        _disable_pipe(adapter.pipe)
-        for transformer in BlockAdapter.flatten(adapter.transformer):
-            _disable_transformer(transformer)
-        for blocks in BlockAdapter.flatten(adapter.blocks):
-            _disable_blocks(blocks)
-        pipe_cls_name = adapter.pipe.__class__.__name__
-        logger.warning(f"Cache Acceleration is disabled for: {pipe_cls_name}")
-    else:
-        pass  # do nothing
-
 
 def supported_pipelines(
     **kwargs,

cache_dit/cache_factory/cache_types.py

@@ -22,11 +22,11 @@ def cache_type(type_hint: "CacheType | str") -> "CacheType":
     if isinstance(type_hint, CacheType):
         return type_hint
 
-    elif type_hint.lower() in (
-        "dual_block_cache",
-        "db_cache",
-        "dbcache",
-        "db",
+    elif type_hint.upper() in (
+        "DUAL_BLOCK_CACHE",
+        "DB_CACHE",
+        "DBCACHE",
+        "DB",
     ):
         return CacheType.DBCache
     return CacheType.NONE

cache_dit/cache_factory/patch_functors/__init__.py

@@ -3,3 +3,9 @@ from cache_dit.cache_factory.patch_functors.functor_flux import FluxPatchFunctor
 from cache_dit.cache_factory.patch_functors.functor_chroma import (
     ChromaPatchFunctor,
 )
+from cache_dit.cache_factory.patch_functors.functor_hidream import (
+    HiDreamPatchFunctor,
+)
+from cache_dit.cache_factory.patch_functors.functor_hunyuan_dit import (
+    HunyuanDiTPatchFunctor,
+)

cache_dit/cache_factory/patch_functors/functor_chroma.py

@@ -46,8 +46,10 @@ class ChromaPatchFunctor(PatchFunctor):
                     block.forward = __patch_single_forward__.__get__(block)
                     is_patched = True
 
+        cls_name = transformer.__class__.__name__
+
         if is_patched:
-            logger.warning("Patched Chroma for cache-dit.")
+            logger.warning(f"Patched {cls_name} for cache-dit.")
             assert not getattr(transformer, "_is_parallelized", False), (
                 "Please call `cache_dit.enable_cache` before Parallelize, "
                 "the __patch_transformer_forward__ will overwrite the "
@@ -56,9 +58,9 @@ class ChromaPatchFunctor(PatchFunctor):
             transformer.forward = __patch_transformer_forward__.__get__(
                 transformer
             )
-            transformer._is_patched = True
 
-        cls_name = transformer.__class__.__name__
+        transformer._is_patched = is_patched  # True or False
+
         logger.info(
             f"Applied {self.__class__.__name__} for {cls_name}, "
             f"Patch: {is_patched}."

cache_dit/cache_factory/patch_functors/functor_flux.py

@@ -47,8 +47,10 @@ class FluxPatchFunctor(PatchFunctor):
                     block.forward = __patch_single_forward__.__get__(block)
                     is_patched = True
 
+        cls_name = transformer.__class__.__name__
+
         if is_patched:
-            logger.warning("Patched Flux for cache-dit.")
+            logger.warning(f"Patched {cls_name} for cache-dit.")
             assert not getattr(transformer, "_is_parallelized", False), (
                 "Please call `cache_dit.enable_cache` before Parallelize, "
                 "the __patch_transformer_forward__ will overwrite the "
@@ -57,9 +59,9 @@ class FluxPatchFunctor(PatchFunctor):
             transformer.forward = __patch_transformer_forward__.__get__(
                 transformer
             )
-            transformer._is_patched = True
 
-        cls_name = transformer.__class__.__name__
+        transformer._is_patched = is_patched  # True or False
+
         logger.info(
             f"Applied {self.__class__.__name__} for {cls_name}, "
             f"Patch: {is_patched}."