cache-dit 0.1.7__py3-none-any.whl → 0.2.0__py3-none-any.whl

cache_dit/cache_factory/dynamic_block_prune/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.dynamic_block_prune import prune_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_prune_on_transformer(
+    transformer: HunyuanVideoTransformer3DModel,
+):
+    if getattr(transformer, "_is_pruned", False):
+        return transformer
+
+    cached_transformer_blocks = torch.nn.ModuleList(
+        [
+            prune_context.DBPrunedTransformerBlocks(
+                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_pruned = True
+
+    return transformer
+
+
+def apply_db_prune_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 = prune_context.collect_prune_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_pruned", False):
+        original_call = pipe.__class__.__call__
+
+        @functools.wraps(original_call)
+        def new_call(self, *args, **kwargs):
+            with prune_context.prune_context(
+                prune_context.create_prune_context(
+                    **cache_kwargs,
+                )
+            ):
+                return original_call(self, *args, **kwargs)
+
+        pipe.__class__.__call__ = new_call
+        pipe.__class__._is_pruned = True
+
+    if not shallow_patch:
+        apply_db_prune_on_transformer(pipe.transformer, **kwargs)
+
+    return pipe

cache_dit/cache_factory/dynamic_block_prune/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.dynamic_block_prune import prune_context
+
+
+def apply_db_prune_on_transformer(
+    transformer: WanTransformer3DModel,
+):
+    if getattr(transformer, "_is_pruned", False):
+        return transformer
+
+    blocks = torch.nn.ModuleList(
+        [
+            prune_context.DBPrunedTransformerBlocks(
+                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_pruned = True
+
+    return transformer
+
+
+def apply_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 = prune_context.collect_prune_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_pruned", False):
+        original_call = pipe.__class__.__call__
+
+        @functools.wraps(original_call)
+        def new_call(self, *args, **kwargs):
+            with prune_context.prune_context(
+                prune_context.create_prune_context(
+                    **cache_kwargs,
+                )
+            ):
+                return original_call(self, *args, **kwargs)
+
+        pipe.__class__.__call__ = new_call
+        pipe.__class__._is_pruned = True
+
+    if not shallow_patch:
+        apply_db_prune_on_transformer(pipe.transformer, **kwargs)
+
+    return pipe

cache_dit/cache_factory/dynamic_block_prune/prune_context.py

@@ -628,7 +628,7 @@ class DBPrunedTransformerBlocks(torch.nn.Module):
         return sorted(non_prune_blocks_ids)
 
     # @torch.compile(dynamic=True)
-    # mark this function as compile with dynamic=True will 
+    # mark this function as compile with dynamic=True will
     # cause precision degradate, so, we choose to disable it
     # now, until we find a better solution or fixed the bug.
     @torch.compiler.disable
@@ -668,7 +668,7 @@ class DBPrunedTransformerBlocks(torch.nn.Module):
         )
 
     # @torch.compile(dynamic=True)
-    # mark this function as compile with dynamic=True will 
+    # mark this function as compile with dynamic=True will
     # cause precision degradate, so, we choose to disable it
     # now, until we find a better solution or fixed the bug.
     @torch.compiler.disable

cache_dit/cache_factory/first_block_cache/diffusers_adapters/__init__.py

@@ -16,6 +16,8 @@ def apply_fb_cache_on_transformer(transformer, *args, **kwargs):
         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}"
@@ -40,6 +42,8 @@ def apply_fb_cache_on_pipe(pipe: DiffusionPipeline, *args, **kwargs):
         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}")