cache-dit 0.2.25__py3-none-any.whl → 0.2.26__py3-none-any.whl

cache_dit/cache_factory/cache_blocks/pattern_3_4_5.py

@@ -0,0 +1,270 @@
+import torch
+
+from cache_dit.cache_factory import cache_context
+from cache_dit.cache_factory import ForwardPattern
+from cache_dit.cache_factory.cache_blocks.utils import (
+    patch_cached_stats,
+)
+from cache_dit.cache_factory.cache_blocks.pattern_base import (
+    DBCachedBlocks_Pattern_Base,
+)
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class DBCachedBlocks_Pattern_3_4_5(DBCachedBlocks_Pattern_Base):
+    _supported_patterns = [
+        ForwardPattern.Pattern_3,
+        ForwardPattern.Pattern_4,
+        ForwardPattern.Pattern_5,
+    ]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        original_hidden_states = hidden_states
+        # 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,
+            *args,
+            **kwargs,
+        )
+
+        Fn_hidden_states_residual = hidden_states - original_hidden_states
+        del original_hidden_states
+
+        cache_context.mark_step_begin()
+        # Residual L1 diff or Hidden States L1 diff
+        can_use_cache = cache_context.get_can_use_cache(
+            (
+                Fn_hidden_states_residual
+                if not cache_context.is_l1_diff_enabled()
+                else hidden_states
+            ),
+            parallelized=self._is_parallelized(),
+            prefix=(
+                "Fn_residual"
+                if not cache_context.is_l1_diff_enabled()
+                else "Fn_hidden_states"
+            ),
+        )
+
+        torch._dynamo.graph_break()
+        if can_use_cache:
+            cache_context.add_cached_step()
+            del Fn_hidden_states_residual
+            hidden_states, encoder_hidden_states = (
+                cache_context.apply_hidden_states_residual(
+                    hidden_states,
+                    # None Pattern 3, else 4, 5
+                    encoder_hidden_states,
+                    prefix=(
+                        "Bn_residual"
+                        if cache_context.is_cache_residual()
+                        else "Bn_hidden_states"
+                    ),
+                    encoder_prefix=(
+                        "Bn_residual"
+                        if cache_context.is_encoder_cache_residual()
+                        else "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,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+        else:
+            cache_context.set_Fn_buffer(
+                Fn_hidden_states_residual, prefix="Fn_residual"
+            )
+            if cache_context.is_l1_diff_enabled():
+                # for hidden states L1 diff
+                cache_context.set_Fn_buffer(hidden_states, "Fn_hidden_states")
+            del Fn_hidden_states_residual
+            torch._dynamo.graph_break()
+            (
+                hidden_states,
+                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,
+            )
+            torch._dynamo.graph_break()
+            if cache_context.is_cache_residual():
+                cache_context.set_Bn_buffer(
+                    hidden_states_residual,
+                    prefix="Bn_residual",
+                )
+            else:
+                # TaylorSeer
+                cache_context.set_Bn_buffer(
+                    hidden_states,
+                    prefix="Bn_hidden_states",
+                )
+            if cache_context.is_encoder_cache_residual():
+                cache_context.set_Bn_encoder_buffer(
+                    # None Pattern 3, else 4, 5
+                    encoder_hidden_states_residual,
+                    prefix="Bn_residual",
+                )
+            else:
+                # TaylorSeer
+                cache_context.set_Bn_encoder_buffer(
+                    # None Pattern 3, else 4, 5
+                    encoder_hidden_states,
+                    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,
+            )
+
+        patch_cached_stats(self.transformer)
+        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)
+            )
+        )
+
+    def call_Fn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        assert cache_context.Fn_compute_blocks() <= len(
+            self.transformer_blocks
+        ), (
+            f"Fn_compute_blocks {cache_context.Fn_compute_blocks()} must be less than "
+            f"the number of transformer blocks {len(self.transformer_blocks)}"
+        )
+        encoder_hidden_states = None  # Pattern 3
+        for block in self._Fn_blocks():
+            hidden_states = block(
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+            if not isinstance(hidden_states, torch.Tensor):  # Pattern 4, 5
+                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
+
+    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
+        for block in self._Mn_blocks():
+            hidden_states = block(
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+            if not isinstance(hidden_states, torch.Tensor):  # Pattern 4, 5
+                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,
+                    )
+
+        # 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,
+            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 cache_context.Bn_compute_blocks() == 0:
+            return hidden_states, encoder_hidden_states
+
+        assert cache_context.Bn_compute_blocks() <= len(
+            self.transformer_blocks
+        ), (
+            f"Bn_compute_blocks {cache_context.Bn_compute_blocks()} must be less than "
+            f"the number of transformer blocks {len(self.transformer_blocks)}"
+        )
+        if len(cache_context.Bn_compute_blocks_ids()) > 0:
+            raise ValueError(
+                f"Bn_compute_blocks_ids is not support for "
+                f"patterns: {self._supported_patterns}."
+            )
+        else:
+            # Compute all Bn blocks if no specific Bn compute blocks ids are set.
+            for block in self._Bn_blocks():
+                hidden_states = block(
+                    hidden_states,
+                    *args,
+                    **kwargs,
+                )
+                if not isinstance(hidden_states, torch.Tensor):  # Pattern 4,5
+                    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

cache_dit/cache_factory/{cache_blocks.py → cache_blocks/pattern_base.py}

@@ -4,12 +4,15 @@ import torch.distributed as dist
 
 from cache_dit.cache_factory import cache_context
 from cache_dit.cache_factory import ForwardPattern
+from cache_dit.cache_factory.cache_blocks.utils import (
+    patch_cached_stats,
+)
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
 
 
-class DBCachedTransformerBlocks(torch.nn.Module):
+class DBCachedBlocks_Pattern_Base(torch.nn.Module):
     _supported_patterns = [
         ForwardPattern.Pattern_0,
         ForwardPattern.Pattern_1,
@@ -29,18 +32,30 @@ class DBCachedTransformerBlocks(torch.nn.Module):
         self.transformer_blocks = transformer_blocks
         self.forward_pattern = forward_pattern
         self._check_forward_pattern()
+        logger.info(f"Match Cached Blocks: {self.__class__.__name__}")
 
     def _check_forward_pattern(self):
         assert (
             self.forward_pattern.Supported
             and self.forward_pattern in self._supported_patterns
-        ), f"Pattern {self.forward_pattern} is not support for DBCache now!"
+        ), f"Pattern {self.forward_pattern} is not supported now!"
 
         if self.transformer_blocks is not None:
             for block in self.transformer_blocks:
                 forward_parameters = set(
                     inspect.signature(block.forward).parameters.keys()
                 )
+                num_outputs = str(
+                    inspect.signature(block.forward).return_annotation
+                ).count("torch.Tensor")
+
+                if num_outputs > 0:
+                    assert len(self.forward_pattern.Out) == num_outputs, (
+                        f"The number of block's outputs is {num_outputs} don't not "
+                        f"match the number of the pattern: {self.forward_pattern}, "
+                        f"Out: {len(self.forward_pattern.Out)}."
+                    )
+
                 for required_param in self.forward_pattern.In:
                     assert (
                         required_param in forward_parameters
@@ -479,19 +494,3 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                         )
 
         return hidden_states, encoder_hidden_states
-
-
-@torch.compiler.disable
-def patch_cached_stats(
-    transformer,
-):
-    # Patch the cached stats to the transformer, the cached stats
-    # will be reset for each calling of pipe.__call__(**kwargs).
-    if transformer is None:
-        return
-
-    # TODO: Patch more cached stats to the transformer
-    transformer._cached_steps = cache_context.get_cached_steps()
-    transformer._residual_diffs = cache_context.get_residual_diffs()
-    transformer._cfg_cached_steps = cache_context.get_cfg_cached_steps()
-    transformer._cfg_residual_diffs = cache_context.get_cfg_residual_diffs()

cache_dit/cache_factory/cache_blocks/utils.py

@@ -0,0 +1,19 @@
+import torch
+
+from cache_dit.cache_factory import cache_context
+
+
+@torch.compiler.disable
+def patch_cached_stats(
+    transformer,
+):
+    # Patch the cached stats to the transformer, the cached stats
+    # will be reset for each calling of pipe.__call__(**kwargs).
+    if transformer is None:
+        return
+
+    # TODO: Patch more cached stats to the transformer
+    transformer._cached_steps = cache_context.get_cached_steps()
+    transformer._residual_diffs = cache_context.get_residual_diffs()
+    transformer._cfg_cached_steps = cache_context.get_cfg_cached_steps()
+    transformer._cfg_residual_diffs = cache_context.get_cfg_residual_diffs()

cache_dit/cache_factory/cache_context.py

@@ -941,7 +941,11 @@ def get_Bn_buffer(prefix: str = "Bn"):
 
 
 @torch.compiler.disable
-def set_Bn_encoder_buffer(buffer: torch.Tensor, prefix: str = "Bn"):
+def set_Bn_encoder_buffer(buffer: torch.Tensor | None, prefix: str = "Bn"):
+    # DON'T set None Buffer
+    if buffer is None:
+        return
+
     # This buffer is use for encoder hidden states approximation.
     if is_encoder_taylorseer_enabled():
         # taylorseer, encoder_taylorseer

cache_dit/cache_factory/cache_interface.py

@@ -1,3 +1,4 @@
+from typing import Any, Tuple, List
 from diffusers import DiffusionPipeline
 from cache_dit.cache_factory.forward_pattern import ForwardPattern
 from cache_dit.cache_factory.cache_types import CacheType
@@ -9,9 +10,13 @@ from cache_dit.logger import init_logger
 logger = init_logger(__name__)
 
 
+def supported_pipelines() -> Tuple[int, List[str]]:
+    return UnifiedCacheAdapter.supported_pipelines()
+
+
 def enable_cache(
     # BlockAdapter & forward pattern
-    pipe_or_adapter: DiffusionPipeline | BlockAdapter,
+    pipe_or_adapter: DiffusionPipeline | BlockAdapter | Any,
     forward_pattern: ForwardPattern = ForwardPattern.Pattern_0,
     # Cache context kwargs
     Fn_compute_blocks: int = 8,
@@ -30,7 +35,7 @@ def enable_cache(
     taylorseer_cache_type: str = "residual",
     taylorseer_order: int = 2,
     **other_cache_kwargs,
-) -> DiffusionPipeline:
+) -> DiffusionPipeline | Any:
     r"""
     Unified Cache API for  almost Any Diffusion Transformers (with Transformer Blocks
     that match the specific Input and Output patterns).

cache_dit/cache_factory/forward_pattern.py

@@ -19,39 +19,57 @@ class ForwardPattern(Enum):
         self.Supported = Supported
 
     Pattern_0 = (
-        True,
-        False,
-        False,
-        ("hidden_states", "encoder_hidden_states"),
-        ("hidden_states", "encoder_hidden_states"),
-        True,
+        True,  # Return_H_First
+        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,
-        False,
-        False,
-        ("hidden_states", "encoder_hidden_states"),
-        ("encoder_hidden_states", "hidden_states"),
-        True,
+        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,
-        True,
-        False,
-        ("hidden_states", "encoder_hidden_states"),
-        ("hidden_states",),
-        True,
+        False, # Return_H_First
+        True,  # Return_H_Only
+        False, # Forward_H_only
+        ("hidden_states", "encoder_hidden_states"),  # In
+        ("hidden_states",),                          # Out
+        True,  # Supported
     )
 
     Pattern_3 = (
-        False,
-        True,
-        False,
-        ("hidden_states",),
-        ("hidden_states",),
-        False,
+        False, # Return_H_First
+        True,  # Return_H_Only
+        True,  # Forward_H_only
+        ("hidden_states",),  # In
+        ("hidden_states",),  # Out
+        True,  # Supported
+    )
+
+    Pattern_4 = (
+        True,  # Return_H_First
+        False, # Return_H_Only
+        True,  # Forward_H_only
+        ("hidden_states",),                          # In
+        ("hidden_states", "encoder_hidden_states"),  # Out
+        True,  # Supported
+    )
+
+    Pattern_5 = (
+        False, # Return_H_First
+        False, # Return_H_Only
+        True,  # Forward_H_only
+        ("hidden_states",),                          # In
+        ("encoder_hidden_states", "hidden_states"),  # Out
+        True,  # Supported
     )
 
     @staticmethod
@@ -60,4 +78,7 @@ class ForwardPattern(Enum):
             ForwardPattern.Pattern_0,
             ForwardPattern.Pattern_1,
             ForwardPattern.Pattern_2,
+            ForwardPattern.Pattern_3,
+            ForwardPattern.Pattern_4,
+            ForwardPattern.Pattern_5,
         ]

cache_dit/cache_factory/patch_functors/__init__.py

@@ -0,0 +1,5 @@
+from cache_dit.cache_factory.patch_functors.functor_base import PatchFunctor
+from cache_dit.cache_factory.patch_functors.functor_flux import FluxPatchFunctor
+from cache_dit.cache_factory.patch_functors.functor_chroma import (
+    ChromaPatchFunctor,
+)

cache_dit/cache_factory/patch_functors/functor_base.py

@@ -0,0 +1,18 @@
+import torch
+from abc import abstractmethod
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class PatchFunctor:
+
+    @abstractmethod
+    def apply(
+        self,
+        transformer: torch.nn.Module,
+        *args,
+        **kwargs,
+    ) -> torch.nn.Module:
+        raise NotImplementedError("apply method is not implemented.")