cache-dit 0.1.8__py3-none-any.whl → 0.2.1__py3-none-any.whl

cache_dit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.8'
-__version_tuple__ = version_tuple = (0, 1, 8)
+__version__ = version = '0.2.1'
+__version_tuple__ = version_tuple = (0, 2, 1)

cache_dit/cache_factory/dual_block_cache/cache_context.py

@@ -9,6 +9,7 @@ from typing import Any, DefaultDict, Dict, List, Optional, Union
 import torch
 
 import cache_dit.primitives as DP
+from cache_dit.cache_factory.taylorseer import TaylorSeer
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
@@ -60,6 +61,18 @@ class DBCacheContext:
     residual_diffs: DefaultDict[str, float] = dataclasses.field(
         default_factory=lambda: defaultdict(float),
     )
+    # TODO: Support TaylorSeers and SLG in Dual Block Cache
+    # TaylorSeers: 
+    # Title: From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers
+    # Url: https://arxiv.org/pdf/2503.06923
+    taylorseer: Optional[TaylorSeer] = None
+    alter_taylorseer: Optional[TaylorSeer] = None
+
+    # Skip Layer Guidance, SLG
+    # https://github.com/huggingface/candle/issues/2588
+    slg_layers: Optional[List[int]] = None
+    slg_start: float = 0.0
+    slg_end: float = 0.1
 
     def get_incremental_name(self, name=None):
         if name is None:
@@ -700,7 +713,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                 encoder_hidden_states,
                 hidden_states_residual,
                 encoder_hidden_states_residual,
-            ) = self.call_MN2n_transformer_blocks(  # middle
+            ) = self.call_Mn_transformer_blocks(  # middle
                 hidden_states,
                 encoder_hidden_states,
                 *args,
@@ -772,32 +785,32 @@ class DBCachedTransformerBlocks(torch.nn.Module):
         return selected_Fn_transformer_blocks
 
     @torch.compiler.disable
-    def _MN2n_single_transformer_blocks(self):  # middle
+    def _Mn_single_transformer_blocks(self):  # middle blocks
         # M(N-2n): transformer_blocks [n,...] + single_transformer_blocks [0,...,N-n]
-        selected_MN2n_single_transformer_blocks = []
+        selected_Mn_single_transformer_blocks = []
         if self.single_transformer_blocks is not None:
             if Bn_compute_blocks() == 0:  # WARN: x[:-0] = []
-                selected_MN2n_single_transformer_blocks = (
+                selected_Mn_single_transformer_blocks = (
                     self.single_transformer_blocks
                 )
             else:
-                selected_MN2n_single_transformer_blocks = (
+                selected_Mn_single_transformer_blocks = (
                     self.single_transformer_blocks[: -Bn_compute_blocks()]
                 )
-        return selected_MN2n_single_transformer_blocks
+        return selected_Mn_single_transformer_blocks
 
     @torch.compiler.disable
-    def _MN2n_transformer_blocks(self):
+    def _Mn_transformer_blocks(self): # middle blocks
         # M(N-2n): only transformer_blocks [n,...,N-n], middle
         if Bn_compute_blocks() == 0:  # WARN: x[:-0] = []
-            selected_MN2n_transformer_blocks = self.transformer_blocks[
+            selected_Mn_transformer_blocks = self.transformer_blocks[
                 Fn_compute_blocks() :
             ]
         else:
-            selected_MN2n_transformer_blocks = self.transformer_blocks[
+            selected_Mn_transformer_blocks = self.transformer_blocks[
                 Fn_compute_blocks() : -Bn_compute_blocks()
             ]
-        return selected_MN2n_transformer_blocks
+        return selected_Mn_transformer_blocks
 
     @torch.compiler.disable
     def _Bn_single_transformer_blocks(self):
@@ -845,7 +858,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
 
         return hidden_states, encoder_hidden_states
 
-    def call_MN2n_transformer_blocks(
+    def call_Mn_transformer_blocks(
         self,
         hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
@@ -873,7 +886,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
             hidden_states = torch.cat(
                 [encoder_hidden_states, hidden_states], dim=1
             )
-            for block in self._MN2n_single_transformer_blocks():
+            for block in self._Mn_single_transformer_blocks():
                 hidden_states = block(
                     hidden_states,
                     *args,
@@ -887,7 +900,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                 dim=1,
             )
         else:
-            for block in self._MN2n_transformer_blocks():
+            for block in self._Mn_transformer_blocks():
                 hidden_states = block(
                     hidden_states,
                     encoder_hidden_states,
@@ -1016,7 +1029,9 @@ class DBCachedTransformerBlocks(torch.nn.Module):
 
     def _compute_and_cache_single_transformer_block(
         self,
-        i: int,  # Block index in the transformer blocks
+        # Block index in the transformer blocks
+        # Bn: 8, block_id should be in [0, 8)
+        block_id: int,
         # Helper inputs for hidden states split and reshape
         original_hidden_states: torch.Tensor,
         original_encoder_hidden_states: torch.Tensor,
@@ -1042,7 +1057,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
             )
             # Cache residuals for the non-compute Bn blocks for
             # subsequent cache steps.
-            if i not in Bn_compute_blocks_ids():
+            if block_id not in Bn_compute_blocks_ids():
                 Bn_i_hidden_states = hidden_states
                 (
                     Bn_i_hidden_states_residual,
@@ -1057,16 +1072,16 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                 # Save original_hidden_states for diff calculation.
                 set_Bn_buffer(
                     Bn_i_original_hidden_states,
-                    prefix=f"Bn_{i}_single_original",
+                    prefix=f"Bn_{block_id}_single_original",
                 )
 
                 set_Bn_buffer(
                     Bn_i_hidden_states_residual,
-                    prefix=f"Bn_{i}_single_residual",
+                    prefix=f"Bn_{block_id}_single_residual",
                 )
                 set_Bn_encoder_buffer(
                     Bn_i_encoder_hidden_states_residual,
-                    prefix=f"Bn_{i}_single_residual",
+                    prefix=f"Bn_{block_id}_single_residual",
                 )
                 del Bn_i_hidden_states
                 del Bn_i_hidden_states_residual
@@ -1077,7 +1092,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
         else:
             # Cache steps: Reuse the cached residuals.
             # Check if the block is in the Bn_compute_blocks_ids.
-            if i in Bn_compute_blocks_ids():
+            if block_id in Bn_compute_blocks_ids():
                 hidden_states = block(
                     hidden_states,
                     *args,
@@ -1091,7 +1106,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                     hidden_states,  # curr step
                     parallelized=self._is_parallelized(),
                     threshold=non_compute_blocks_diff_threshold(),
-                    prefix=f"Bn_{i}_single_original",  # prev step
+                    prefix=f"Bn_{block_id}_single_original",  # prev step
                 ):
                     Bn_i_original_hidden_states = hidden_states
                     (
@@ -1106,7 +1121,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                         apply_hidden_states_residual(
                             Bn_i_original_hidden_states,
                             Bn_i_original_encoder_hidden_states,
-                            prefix=f"Bn_{i}_single_residual",
+                            prefix=f"Bn_{block_id}_single_residual",
                         )
                     )
                     hidden_states = torch.cat(
@@ -1125,7 +1140,9 @@ class DBCachedTransformerBlocks(torch.nn.Module):
 
     def _compute_and_cache_transformer_block(
         self,
-        i: int,  # Block index in the transformer blocks
+        # Block index in the transformer blocks
+        # Bn: 8, block_id should be in [0, 8)
+        block_id: int,  
         # Below are the inputs to the block
         block,  # The transformer block to be executed
         hidden_states: torch.Tensor,
@@ -1158,7 +1175,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                     )
             # Cache residuals for the non-compute Bn blocks for
             # subsequent cache steps.
-            if i not in Bn_compute_blocks_ids():
+            if block_id not in Bn_compute_blocks_ids():
                 Bn_i_hidden_states_residual = (
                     hidden_states - Bn_i_original_hidden_states
                 )
@@ -1169,16 +1186,16 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                 # Save original_hidden_states for diff calculation.
                 set_Bn_buffer(
                     Bn_i_original_hidden_states,
-                    prefix=f"Bn_{i}_original",
+                    prefix=f"Bn_{block_id}_original",
                 )
 
                 set_Bn_buffer(
                     Bn_i_hidden_states_residual,
-                    prefix=f"Bn_{i}_residual",
+                    prefix=f"Bn_{block_id}_residual",
                 )
                 set_Bn_encoder_buffer(
                     Bn_i_encoder_hidden_states_residual,
-                    prefix=f"Bn_{i}_residual",
+                    prefix=f"Bn_{block_id}_residual",
                 )
                 del Bn_i_hidden_states_residual
                 del Bn_i_encoder_hidden_states_residual
@@ -1189,7 +1206,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
         else:
             # Cache steps: Reuse the cached residuals.
             # Check if the block is in the Bn_compute_blocks_ids.
-            if i in Bn_compute_blocks_ids():
+            if block_id in Bn_compute_blocks_ids():
                 hidden_states = block(
                     hidden_states,
                     encoder_hidden_states,
@@ -1211,13 +1228,13 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                     hidden_states,  # curr step
                     parallelized=self._is_parallelized(),
                     threshold=non_compute_blocks_diff_threshold(),
-                    prefix=f"Bn_{i}_original",  # prev step
+                    prefix=f"Bn_{block_id}_original",  # prev step
                 ):
                     hidden_states, encoder_hidden_states = (
                         apply_hidden_states_residual(
                             hidden_states,
                             encoder_hidden_states,
-                            prefix=f"Bn_{i}_residual",
+                            prefix=f"Bn_{block_id}_residual",
                         )
                     )
                 else:

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/__init__.py

@@ -13,6 +13,10 @@ def apply_db_cache_on_transformer(transformer, *args, **kwargs):
         adapter_name = "mochi"
     elif transformer_cls_name.startswith("CogVideoX"):
         adapter_name = "cogvideox"
+    elif transformer_cls_name.startswith("Wan"):
+        adapter_name = "wan"
+    elif transformer_cls_name.startswith("HunyuanVideo"):
+        adapter_name = "hunyuan_video"
     else:
         raise ValueError(
             f"Unknown transformer class name: {transformer_cls_name}"
@@ -35,6 +39,10 @@ def apply_db_cache_on_pipe(pipe: DiffusionPipeline, *args, **kwargs):
         adapter_name = "mochi"
     elif pipe_cls_name.startswith("CogVideoX"):
         adapter_name = "cogvideox"
+    elif pipe_cls_name.startswith("Wan"):
+        adapter_name = "wan"
+    elif pipe_cls_name.startswith("HunyuanVideo"):
+        adapter_name = "hunyuan_video"
     else:
         raise ValueError(f"Unknown pipeline class name: {pipe_cls_name}")
 

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/hunyuan_video.py

@@ -0,0 +1,295 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/hunyuan_video.py
+import functools
+import unittest
+from typing import Any, Dict, Optional, Union
+
+import torch
+from diffusers import DiffusionPipeline, HunyuanVideoTransformer3DModel
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.utils import (
+    scale_lora_layers,
+    unscale_lora_layers,
+    USE_PEFT_BACKEND,
+)
+
+from cache_dit.cache_factory.dual_block_cache import cache_context
+from cache_dit.logger import init_logger
+
+try:
+    from para_attn.para_attn_interface import SparseKVAttnMode
+
+    def is_sparse_kv_attn_available():
+        return True
+
+except ImportError:
+
+    class SparseKVAttnMode:
+        def __enter__(self):
+            pass
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            pass
+
+    def is_sparse_kv_attn_available():
+        return False
+
+
+logger = init_logger(__name__)  # pylint: disable=invalid-name
+
+
+def apply_db_cache_on_transformer(
+    transformer: HunyuanVideoTransformer3DModel,
+):
+    if getattr(transformer, "_is_cached", False):
+        return transformer
+
+    cached_transformer_blocks = torch.nn.ModuleList(
+        [
+            cache_context.DBCachedTransformerBlocks(
+                transformer.transformer_blocks
+                + transformer.single_transformer_blocks,
+                transformer=transformer,
+            )
+        ]
+    )
+    dummy_single_transformer_blocks = torch.nn.ModuleList()
+
+    original_forward = transformer.forward
+
+    @functools.wraps(transformer.__class__.forward)
+    def new_forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.LongTensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_attention_mask: torch.Tensor,
+        pooled_projections: torch.Tensor,
+        guidance: torch.Tensor = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = True,
+        **kwargs,
+    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+        with (
+            unittest.mock.patch.object(
+                self,
+                "transformer_blocks",
+                cached_transformer_blocks,
+            ),
+            unittest.mock.patch.object(
+                self,
+                "single_transformer_blocks",
+                dummy_single_transformer_blocks,
+            ),
+        ):
+            if getattr(self, "_is_parallelized", False):
+                return original_forward(
+                    hidden_states,
+                    timestep,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    pooled_projections,
+                    guidance=guidance,
+                    attention_kwargs=attention_kwargs,
+                    return_dict=return_dict,
+                    **kwargs,
+                )
+            else:
+                if attention_kwargs is not None:
+                    attention_kwargs = attention_kwargs.copy()
+                    lora_scale = attention_kwargs.pop("scale", 1.0)
+                else:
+                    lora_scale = 1.0
+
+                if USE_PEFT_BACKEND:
+                    # weight the lora layers by setting `lora_scale` for each PEFT layer
+                    scale_lora_layers(self, lora_scale)
+                else:
+                    if (
+                        attention_kwargs is not None
+                        and attention_kwargs.get("scale", None) is not None
+                    ):
+                        logger.warning(
+                            "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
+                        )
+
+                batch_size, num_channels, num_frames, height, width = (
+                    hidden_states.shape
+                )
+                p, p_t = self.config.patch_size, self.config.patch_size_t
+                post_patch_num_frames = num_frames // p_t
+                post_patch_height = height // p
+                post_patch_width = width // p
+
+                # 1. RoPE
+                image_rotary_emb = self.rope(hidden_states)
+
+                # 2. Conditional embeddings
+                temb = self.time_text_embed(
+                    timestep, guidance, pooled_projections
+                )
+                hidden_states = self.x_embedder(hidden_states)
+                encoder_hidden_states = self.context_embedder(
+                    encoder_hidden_states, timestep, encoder_attention_mask
+                )
+
+                # 3. Attention mask preparation
+                latent_sequence_length = hidden_states.shape[1]
+                latent_attention_mask = torch.ones(
+                    batch_size,
+                    1,
+                    latent_sequence_length,
+                    device=hidden_states.device,
+                    dtype=torch.bool,
+                )  # [B, 1, N]
+                attention_mask = torch.cat(
+                    [
+                        latent_attention_mask,
+                        encoder_attention_mask.unsqueeze(1).to(torch.bool),
+                    ],
+                    dim=-1,
+                )  # [B, 1, N + M]
+
+                with SparseKVAttnMode():
+                    # 4. Transformer blocks
+                    hidden_states, encoder_hidden_states = (
+                        self.call_transformer_blocks(
+                            hidden_states,
+                            encoder_hidden_states,
+                            temb,
+                            attention_mask,
+                            image_rotary_emb,
+                        )
+                    )
+
+                # 5. Output projection
+                hidden_states = self.norm_out(hidden_states, temb)
+                hidden_states = self.proj_out(hidden_states)
+
+                hidden_states = hidden_states.reshape(
+                    batch_size,
+                    post_patch_num_frames,
+                    post_patch_height,
+                    post_patch_width,
+                    -1,
+                    p_t,
+                    p,
+                    p,
+                )
+                hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7)
+                hidden_states = (
+                    hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+                )
+
+                hidden_states = hidden_states.to(timestep.dtype)
+
+                if USE_PEFT_BACKEND:
+                    # remove `lora_scale` from each PEFT layer
+                    unscale_lora_layers(self, lora_scale)
+
+                if not return_dict:
+                    return (hidden_states,)
+
+                return Transformer2DModelOutput(sample=hidden_states)
+
+    transformer.forward = new_forward.__get__(transformer)
+
+    def call_transformer_blocks(
+        self, hidden_states, encoder_hidden_states, *args, **kwargs
+    ):
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+            def create_custom_forward(module, return_dict=None):
+                def custom_forward(*inputs):
+                    if return_dict is not None:
+                        return module(*inputs, return_dict=return_dict)
+                    else:
+                        return module(*inputs)
+
+                return custom_forward
+
+            ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False}
+
+            for block in self.transformer_blocks:
+                hidden_states, encoder_hidden_states = (
+                    torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(block),
+                        hidden_states,
+                        encoder_hidden_states,
+                        *args,
+                        **kwargs,
+                        **ckpt_kwargs,
+                    )
+                )
+
+            for block in self.single_transformer_blocks:
+                hidden_states, encoder_hidden_states = (
+                    torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(block),
+                        hidden_states,
+                        encoder_hidden_states,
+                        *args,
+                        **kwargs,
+                        **ckpt_kwargs,
+                    )
+                )
+
+        else:
+            for block in self.transformer_blocks:
+                hidden_states, encoder_hidden_states = block(
+                    hidden_states, encoder_hidden_states, *args, **kwargs
+                )
+
+            for block in self.single_transformer_blocks:
+                hidden_states, encoder_hidden_states = block(
+                    hidden_states, encoder_hidden_states, *args, **kwargs
+                )
+
+        return hidden_states, encoder_hidden_states
+
+    transformer.call_transformer_blocks = call_transformer_blocks.__get__(
+        transformer
+    )
+
+    transformer._is_cached = True
+
+    return transformer
+
+
+def apply_db_cache_on_pipe(
+    pipe: DiffusionPipeline,
+    *,
+    shallow_patch: bool = False,
+    residual_diff_threshold=0.06,
+    downsample_factor=1,
+    warmup_steps=0,
+    max_cached_steps=-1,
+    **kwargs,
+):
+    cache_kwargs, kwargs = cache_context.collect_cache_kwargs(
+        default_attrs={
+            "residual_diff_threshold": residual_diff_threshold,
+            "downsample_factor": downsample_factor,
+            "warmup_steps": warmup_steps,
+            "max_cached_steps": max_cached_steps,
+        },
+        **kwargs,
+    )
+    if not getattr(pipe, "_is_cached", False):
+        original_call = pipe.__class__.__call__
+
+        @functools.wraps(original_call)
+        def new_call(self, *args, **kwargs):
+            with cache_context.cache_context(
+                cache_context.create_cache_context(
+                    **cache_kwargs,
+                )
+            ):
+                return original_call(self, *args, **kwargs)
+
+        pipe.__class__.__call__ = new_call
+        pipe.__class__._is_cached = True
+
+    if not shallow_patch:
+        apply_db_cache_on_transformer(pipe.transformer, **kwargs)
+
+    return pipe

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/wan.py

@@ -0,0 +1,99 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache/wan.py
+
+import functools
+import unittest
+
+import torch
+from diffusers import DiffusionPipeline, WanTransformer3DModel
+
+from cache_dit.cache_factory.dual_block_cache import cache_context
+
+
+def apply_db_cache_on_transformer(
+    transformer: WanTransformer3DModel,
+):
+    if getattr(transformer, "_is_cached", False):
+        return transformer
+
+    blocks = torch.nn.ModuleList(
+        [
+            cache_context.DBCachedTransformerBlocks(
+                transformer.blocks,
+                transformer=transformer,
+                return_hidden_states_only=True,
+            )
+        ]
+    )
+
+    original_forward = transformer.forward
+
+    @functools.wraps(transformer.__class__.forward)
+    def new_forward(
+        self,
+        *args,
+        **kwargs,
+    ):
+        with unittest.mock.patch.object(
+            self,
+            "blocks",
+            blocks,
+        ):
+            return original_forward(
+                *args,
+                **kwargs,
+            )
+
+    transformer.forward = new_forward.__get__(transformer)
+
+    transformer._is_cached = True
+
+    return transformer
+
+
+def apply_db_cache_on_pipe(
+    pipe: DiffusionPipeline,
+    *,
+    shallow_patch: bool = False,
+    residual_diff_threshold=0.03,
+    downsample_factor=1,
+    # SLG is not supported in WAN with DBCache yet
+    # slg_layers=None,
+    # slg_start: float = 0.0,
+    # slg_end: float = 0.1,
+    warmup_steps=0,
+    max_cached_steps=-1,
+    **kwargs,
+):
+    cache_kwargs, kwargs = cache_context.collect_cache_kwargs(
+        default_attrs={
+            "residual_diff_threshold": residual_diff_threshold,
+            "downsample_factor": downsample_factor,
+            # "enable_alter_cache": True,
+            # "slg_layers": slg_layers,
+            # "slg_start": slg_start,
+            # "slg_end": slg_end,
+            "num_inference_steps": kwargs.get("num_inference_steps", 50),
+            "warmup_steps": warmup_steps,
+            "max_cached_steps": max_cached_steps,
+        },
+        **kwargs,
+    )
+    if not getattr(pipe, "_is_cached", False):
+        original_call = pipe.__class__.__call__
+
+        @functools.wraps(original_call)
+        def new_call(self, *args, **kwargs):
+            with cache_context.cache_context(
+                cache_context.create_cache_context(
+                    **cache_kwargs,
+                )
+            ):
+                return original_call(self, *args, **kwargs)
+
+        pipe.__class__.__call__ = new_call
+        pipe.__class__._is_cached = True
+
+    if not shallow_patch:
+        apply_db_cache_on_transformer(pipe.transformer, **kwargs)
+
+    return pipe