cache-dit 0.2.1__py3-none-any.whl → 0.2.3__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.2.1'
-__version_tuple__ = version_tuple = (0, 2, 1)
+__version__ = version = '0.2.3'
+__version_tuple__ = version_tuple = (0, 2, 3)

cache_dit/cache_factory/dual_block_cache/cache_context.py

@@ -61,19 +61,55 @@ class DBCacheContext:
     residual_diffs: DefaultDict[str, float] = dataclasses.field(
         default_factory=lambda: defaultdict(float),
     )
-    # TODO: Support TaylorSeers and SLG in Dual Block Cache
-    # TaylorSeers: 
+    # Support TaylorSeers in Dual Block Cache
     # Title: From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers
     # Url: https://arxiv.org/pdf/2503.06923
+    enable_taylorseer: bool = False
+    enable_encoder_taylorseer: bool = False
+    # NOTE: use residual cache for taylorseer may incur precision loss
+    taylorseer_cache_type: str = "hidden_states"  # residual or hidden_states
+    taylorseer_kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict)
     taylorseer: Optional[TaylorSeer] = None
+    encoder_tarlorseer: Optional[TaylorSeer] = None
     alter_taylorseer: Optional[TaylorSeer] = None
+    alter_encoder_taylorseer: Optional[TaylorSeer] = None
 
+    # TODO: Support SLG in Dual Block Cache
     # 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
 
+    @torch.compiler.disable
+    def __post_init__(self):
+
+        if "warmup_steps" not in self.taylorseer_kwargs:
+            # If warmup_steps is not set in taylorseer_kwargs,
+            # set the same as warmup_steps for DBCache
+            self.taylorseer_kwargs["warmup_steps"] = (
+                self.warmup_steps if self.warmup_steps > 0 else 1
+            )
+
+        # Only set n_derivatives as 2 or 3, which is enough for most cases.
+        if "n_derivatives" not in self.taylorseer_kwargs:
+            self.taylorseer_kwargs["n_derivatives"] = max(
+                2, min(3, self.taylorseer_kwargs["warmup_steps"])
+            )
+
+        if self.enable_taylorseer:
+            self.taylorseer = TaylorSeer(**self.taylorseer_kwargs)
+            if self.enable_alter_cache:
+                self.alter_taylorseer = TaylorSeer(**self.taylorseer_kwargs)
+
+        if self.enable_encoder_taylorseer:
+            self.encoder_tarlorseer = TaylorSeer(**self.taylorseer_kwargs)
+            if self.enable_alter_cache:
+                self.alter_encoder_taylorseer = TaylorSeer(
+                    **self.taylorseer_kwargs
+                )
+
+    @torch.compiler.disable
     def get_incremental_name(self, name=None):
         if name is None:
             name = "default"
@@ -81,9 +117,11 @@ class DBCacheContext:
         self.incremental_name_counters[name] += 1
         return f"{name}_{idx}"
 
+    @torch.compiler.disable
     def reset_incremental_names(self):
         self.incremental_name_counters.clear()
 
+    @torch.compiler.disable
     def get_residual_diff_threshold(self):
         if self.enable_alter_cache:
             residual_diff_threshold = self.alter_residual_diff_threshold
@@ -96,25 +134,30 @@ class DBCacheContext:
             residual_diff_threshold = residual_diff_threshold.item()
         return residual_diff_threshold
 
+    @torch.compiler.disable
     def get_buffer(self, name):
         if self.enable_alter_cache and self.is_alter_cache:
             name = f"{name}_alter"
         return self.buffers.get(name)
 
+    @torch.compiler.disable
     def set_buffer(self, name, buffer):
         if self.enable_alter_cache and self.is_alter_cache:
             name = f"{name}_alter"
         self.buffers[name] = buffer
 
+    @torch.compiler.disable
     def remove_buffer(self, name):
         if self.enable_alter_cache and self.is_alter_cache:
             name = f"{name}_alter"
         if name in self.buffers:
             del self.buffers[name]
 
+    @torch.compiler.disable
     def clear_buffers(self):
         self.buffers.clear()
 
+    @torch.compiler.disable
     def mark_step_begin(self):
         if not self.enable_alter_cache:
             self.executed_steps += 1
@@ -129,25 +172,53 @@ class DBCacheContext:
             self.cached_steps.clear()
             self.residual_diffs.clear()
             self.reset_incremental_names()
+            # Reset the TaylorSeers cache at the beginning of each inference.
+            # reset_cache will set the current step to -1 for TaylorSeer,
+            if self.enable_taylorseer or self.enable_encoder_taylorseer:
+                taylorseer, encoder_taylorseer = self.get_taylorseers()
+                if taylorseer is not None:
+                    taylorseer.reset_cache()
+                if encoder_taylorseer is not None:
+                    encoder_taylorseer.reset_cache()
+
+        # mark_step_begin of TaylorSeer must be called after the cache is reset.
+        if self.enable_taylorseer or self.enable_encoder_taylorseer:
+            taylorseer, encoder_taylorseer = self.get_taylorseers()
+            if taylorseer is not None:
+                taylorseer.mark_step_begin()
+            if encoder_taylorseer is not None:
+                encoder_taylorseer.mark_step_begin()
 
+    @torch.compiler.disable
+    def get_taylorseers(self):
+        if self.enable_alter_cache and self.is_alter_cache:
+            return self.alter_taylorseer, self.alter_encoder_taylorseer
+        return self.taylorseer, self.encoder_tarlorseer
+
+    @torch.compiler.disable
     def add_residual_diff(self, diff):
         step = str(self.get_current_step())
         if step not in self.residual_diffs:
             # Only add the diff if it is not already recorded for this step
             self.residual_diffs[step] = diff
 
+    @torch.compiler.disable
     def get_residual_diffs(self):
         return self.residual_diffs.copy()
 
+    @torch.compiler.disable
     def add_cached_step(self):
         self.cached_steps.append(self.get_current_step())
 
+    @torch.compiler.disable
     def get_cached_steps(self):
         return self.cached_steps.copy()
 
+    @torch.compiler.disable
     def get_current_step(self):
         return self.executed_steps - 1
 
+    @torch.compiler.disable
     def is_in_warmup(self):
         return self.get_current_step() < self.warmup_steps
 
@@ -229,6 +300,50 @@ def get_residual_diffs():
     return cache_context.get_residual_diffs()
 
 
+@torch.compiler.disable
+def is_taylorseer_enabled():
+    cache_context = get_current_cache_context()
+    assert cache_context is not None, "cache_context must be set before"
+    return cache_context.enable_taylorseer
+
+
+@torch.compiler.disable
+def is_encoder_taylorseer_enabled():
+    cache_context = get_current_cache_context()
+    assert cache_context is not None, "cache_context must be set before"
+    return cache_context.enable_encoder_taylorseer
+
+
+@torch.compiler.disable
+def get_taylorseers():
+    cache_context = get_current_cache_context()
+    assert cache_context is not None, "cache_context must be set before"
+    return cache_context.get_taylorseers()
+
+
+@torch.compiler.disable
+def is_taylorseer_cache_residual():
+    cache_context = get_current_cache_context()
+    assert cache_context is not None, "cache_context must be set before"
+    return cache_context.taylorseer_cache_type == "residual"
+
+
+@torch.compiler.disable
+def is_cache_residual():
+    if is_taylorseer_enabled():
+        # residual or hidden_states
+        return is_taylorseer_cache_residual()
+    return True
+
+
+@torch.compiler.disable
+def is_encoder_cache_residual():
+    if is_encoder_taylorseer_enabled():
+        # residual or hidden_states
+        return is_taylorseer_cache_residual()
+    return True
+
+
 @torch.compiler.disable
 def is_alter_cache_enabled():
     cache_context = get_current_cache_context()
@@ -380,16 +495,21 @@ def collect_cache_kwargs(default_attrs: dict, **kwargs):
         for attr in cache_attrs
     }
 
+    def _safe_set_sequence_field(
+        field_name: str,
+        default_value: Any = None,
+    ):
+        if field_name not in cache_kwargs:
+            cache_kwargs[field_name] = kwargs.pop(
+                field_name,
+                default_value,
+            )
+
     # Manually set sequence fields, namely, Fn_compute_blocks_ids
     # and Bn_compute_blocks_ids, which are lists or sets.
-    cache_kwargs["Fn_compute_blocks_ids"] = kwargs.pop(
-        "Fn_compute_blocks_ids",
-        [],
-    )
-    cache_kwargs["Bn_compute_blocks_ids"] = kwargs.pop(
-        "Bn_compute_blocks_ids",
-        [],
-    )
+    _safe_set_sequence_field("Fn_compute_blocks_ids", [])
+    _safe_set_sequence_field("Bn_compute_blocks_ids", [])
+    _safe_set_sequence_field("taylorseer_kwargs", {})
 
     assert default_attrs is not None, "default_attrs must be set before"
     for attr in cache_attrs:
@@ -484,6 +604,7 @@ def are_two_tensors_similar(
 @torch.compiler.disable
 def set_Fn_buffer(buffer: torch.Tensor, prefix: str = "Fn"):
     # Set hidden_states or residual for Fn blocks.
+    # This buffer is only use for L1 diff calculation.
     downsample_factor = get_downsample_factor()
     if downsample_factor > 1:
         buffer = buffer[..., ::downsample_factor]
@@ -510,22 +631,79 @@ def get_Fn_encoder_buffer(prefix: str = "Fn"):
 @torch.compiler.disable
 def set_Bn_buffer(buffer: torch.Tensor, prefix: str = "Bn"):
     # Set hidden_states or residual for Bn blocks.
-    set_buffer(f"{prefix}_buffer", buffer)
+    # This buffer is use for hidden states approximation.
+    if is_taylorseer_enabled():
+        # taylorseer, encoder_taylorseer
+        taylorseer, _ = get_taylorseers()
+        if taylorseer is not None:
+            # Use TaylorSeer to update the buffer
+            taylorseer.update(buffer)
+        else:
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug(
+                    "TaylorSeer is enabled but not set in the cache context. "
+                    "Falling back to default buffer retrieval."
+                )
+            set_buffer(f"{prefix}_buffer", buffer)
+    else:
+        set_buffer(f"{prefix}_buffer", buffer)
 
 
 @torch.compiler.disable
 def get_Bn_buffer(prefix: str = "Bn"):
-    return get_buffer(f"{prefix}_buffer")
+    if is_taylorseer_enabled():
+        taylorseer, _ = get_taylorseers()
+        if taylorseer is not None:
+            return taylorseer.approximate_value()
+        else:
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug(
+                    "TaylorSeer is enabled but not set in the cache context. "
+                    "Falling back to default buffer retrieval."
+                )
+            # Fallback to default buffer retrieval
+            return get_buffer(f"{prefix}_buffer")
+    else:
+        return get_buffer(f"{prefix}_buffer")
 
 
 @torch.compiler.disable
 def set_Bn_encoder_buffer(buffer: torch.Tensor, prefix: str = "Bn"):
-    set_buffer(f"{prefix}_encoder_buffer", buffer)
+    # This buffer is use for encoder hidden states approximation.
+    if is_encoder_taylorseer_enabled():
+        # taylorseer, encoder_taylorseer
+        _, encoder_taylorseer = get_taylorseers()
+        if encoder_taylorseer is not None:
+            # Use TaylorSeer to update the buffer
+            encoder_taylorseer.update(buffer)
+        else:
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug(
+                    "TaylorSeer is enabled but not set in the cache context. "
+                    "Falling back to default buffer retrieval."
+                )
+            set_buffer(f"{prefix}_encoder_buffer", buffer)
+    else:
+        set_buffer(f"{prefix}_encoder_buffer", buffer)
 
 
 @torch.compiler.disable
 def get_Bn_encoder_buffer(prefix: str = "Bn"):
-    return get_buffer(f"{prefix}_encoder_buffer")
+    if is_encoder_taylorseer_enabled():
+        _, encoder_taylorseer = get_taylorseers()
+        if encoder_taylorseer is not None:
+            # Use TaylorSeer to approximate the value
+            return encoder_taylorseer.approximate_value()
+        else:
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug(
+                    "TaylorSeer is enabled but not set in the cache context. "
+                    "Falling back to default buffer retrieval."
+                )
+            # Fallback to default buffer retrieval
+            return get_buffer(f"{prefix}_encoder_buffer")
+    else:
+        return get_buffer(f"{prefix}_encoder_buffer")
 
 
 @torch.compiler.disable
@@ -533,29 +711,38 @@ def apply_hidden_states_residual(
     hidden_states: torch.Tensor,
     encoder_hidden_states: torch.Tensor,
     prefix: str = "Bn",
+    encoder_prefix: str = "Bn_encoder",
 ):
     # Allow Bn and Fn prefix to be used for residual cache.
     if "Bn" in prefix:
-        hidden_states_residual = get_Bn_buffer(prefix)
+        hidden_states_prev = get_Bn_buffer(prefix)
     else:
-        hidden_states_residual = get_Fn_buffer(prefix)
+        hidden_states_prev = get_Fn_buffer(prefix)
 
-    assert (
-        hidden_states_residual is not None
-    ), f"{prefix}_buffer must be set before"
-    hidden_states = hidden_states_residual + hidden_states
+    assert hidden_states_prev is not None, f"{prefix}_buffer must be set before"
 
-    if "Bn" in prefix:
-        encoder_hidden_states_residual = get_Bn_encoder_buffer(prefix)
+    if is_cache_residual():
+        hidden_states = hidden_states_prev + hidden_states
+    else:
+        # If cache is not residual, we use the hidden states directly
+        hidden_states = hidden_states_prev
+
+    if "Bn" in encoder_prefix:
+        encoder_hidden_states_prev = get_Bn_encoder_buffer(encoder_prefix)
     else:
-        encoder_hidden_states_residual = get_Fn_encoder_buffer(prefix)
+        encoder_hidden_states_prev = get_Fn_encoder_buffer(encoder_prefix)
 
     assert (
-        encoder_hidden_states_residual is not None
+        encoder_hidden_states_prev is not None
     ), f"{prefix}_encoder_buffer must be set before"
-    encoder_hidden_states = (
-        encoder_hidden_states_residual + encoder_hidden_states
-    )
+
+    if 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
 
     hidden_states = hidden_states.contiguous()
     encoder_hidden_states = encoder_hidden_states.contiguous()
@@ -687,11 +874,22 @@ class DBCachedTransformerBlocks(torch.nn.Module):
 
         torch._dynamo.graph_break()
         if can_use_cache:
+            torch._dynamo.graph_break()
             add_cached_step()
             del Fn_hidden_states_residual
             hidden_states, encoder_hidden_states = apply_hidden_states_residual(
-                hidden_states, encoder_hidden_states, prefix="Bn_residual"
+                hidden_states,
+                encoder_hidden_states,
+                prefix=(
+                    "Bn_residual" if is_cache_residual() else "Bn_hidden_states"
+                ),
+                encoder_prefix=(
+                    "Bn_residual"
+                    if 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 = (
@@ -703,11 +901,13 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                 )
             )
         else:
+            torch._dynamo.graph_break()
             set_Fn_buffer(Fn_hidden_states_residual, prefix="Fn_residual")
             if is_l1_diff_enabled():
                 # for hidden states L1 diff
                 set_Fn_buffer(hidden_states, "Fn_hidden_states")
             del Fn_hidden_states_residual
+            torch._dynamo.graph_break()
             (
                 hidden_states,
                 encoder_hidden_states,
@@ -719,10 +919,30 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                 *args,
                 **kwargs,
             )
-            set_Bn_buffer(hidden_states_residual, prefix="Bn_residual")
-            set_Bn_encoder_buffer(
-                encoder_hidden_states_residual, prefix="Bn_residual"
-            )
+            torch._dynamo.graph_break()
+            if is_cache_residual():
+                set_Bn_buffer(
+                    hidden_states_residual,
+                    prefix="Bn_residual",
+                )
+            else:
+                # TaylorSeer
+                set_Bn_buffer(
+                    hidden_states,
+                    prefix="Bn_hidden_states",
+                )
+            if is_encoder_cache_residual():
+                set_Bn_encoder_buffer(
+                    encoder_hidden_states_residual,
+                    prefix="Bn_residual",
+                )
+            else:
+                # TaylorSeer
+                set_Bn_encoder_buffer(
+                    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 = (
@@ -772,16 +992,6 @@ class DBCachedTransformerBlocks(torch.nn.Module):
         selected_Fn_transformer_blocks = self.transformer_blocks[
             : Fn_compute_blocks()
         ]
-        # Skip the blocks if they are not in the Fn_compute_blocks_ids.
-        # WARN: DON'T set len(Fn_compute_blocks_ids) > 0 NOW, still have
-        # some precision issues. We don't know whether a step should be
-        # cached or not before the first Fn blocks are processed.
-        if len(Fn_compute_blocks_ids()) > 0:
-            selected_Fn_transformer_blocks = [
-                selected_Fn_transformer_blocks[i]
-                for i in Fn_compute_blocks_ids()
-                if i < len(selected_Fn_transformer_blocks)
-            ]
         return selected_Fn_transformer_blocks
 
     @torch.compiler.disable
@@ -800,7 +1010,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
         return selected_Mn_single_transformer_blocks
 
     @torch.compiler.disable
-    def _Mn_transformer_blocks(self): # middle blocks
+    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_Mn_transformer_blocks = self.transformer_blocks[
@@ -1074,6 +1284,10 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                     Bn_i_original_hidden_states,
                     prefix=f"Bn_{block_id}_single_original",
                 )
+                set_Bn_encoder_buffer(
+                    Bn_i_original_hidden_states,
+                    prefix=f"Bn_{block_id}_single_original",
+                )
 
                 set_Bn_buffer(
                     Bn_i_hidden_states_residual,
@@ -1121,7 +1335,16 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                         apply_hidden_states_residual(
                             Bn_i_original_hidden_states,
                             Bn_i_original_encoder_hidden_states,
-                            prefix=f"Bn_{block_id}_single_residual",
+                            prefix=(
+                                f"Bn_{block_id}_single_residual"
+                                if is_cache_residual()
+                                else f"Bn_{block_id}_single_original"
+                            ),
+                            encoder_prefix=(
+                                f"Bn_{block_id}_single_residual"
+                                if is_encoder_cache_residual()
+                                else f"Bn_{block_id}_single_original"
+                            ),
                         )
                     )
                     hidden_states = torch.cat(
@@ -1142,7 +1365,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
         self,
         # Block index in the transformer blocks
         # Bn: 8, block_id should be in [0, 8)
-        block_id: int,  
+        block_id: int,
         # Below are the inputs to the block
         block,  # The transformer block to be executed
         hidden_states: torch.Tensor,
@@ -1188,6 +1411,10 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                     Bn_i_original_hidden_states,
                     prefix=f"Bn_{block_id}_original",
                 )
+                set_Bn_encoder_buffer(
+                    Bn_i_original_encoder_hidden_states,
+                    prefix=f"Bn_{block_id}_original",
+                )
 
                 set_Bn_buffer(
                     Bn_i_hidden_states_residual,
@@ -1234,7 +1461,16 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                         apply_hidden_states_residual(
                             hidden_states,
                             encoder_hidden_states,
-                            prefix=f"Bn_{block_id}_residual",
+                            prefix=(
+                                f"Bn_{block_id}_residual"
+                                if is_cache_residual()
+                                else f"Bn_{block_id}_original"
+                            ),
+                            encoder_prefix=(
+                                f"Bn_{block_id}_residual"
+                                if is_encoder_cache_residual()
+                                else f"Bn_{block_id}_original"
+                            ),
                         )
                     )
                 else:
@@ -1362,17 +1598,6 @@ def patch_cached_stats(
     if transformer is None:
         return
 
-    cached_transformer_blocks = getattr(transformer, "transformer_blocks", None)
-    if cached_transformer_blocks is None:
-        return
-
-    if isinstance(cached_transformer_blocks, torch.nn.ModuleList):
-        cached_transformer_blocks = cached_transformer_blocks[0]
-    if not isinstance(
-        cached_transformer_blocks, DBCachedTransformerBlocks
-    ) or not isinstance(transformer, torch.nn.Module):
-        return
-
     # TODO: Patch more cached stats to the transformer
     transformer._cached_steps = get_cached_steps()
     transformer._residual_diffs = get_residual_diffs()

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/__init__.py

@@ -1,5 +1,3 @@
-# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/__init__.py
-
 import importlib
 
 from diffusers import DiffusionPipeline

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/cogvideox.py

@@ -1,5 +1,3 @@
-# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/cogvideox.py
-
 import functools
 import unittest
 

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/flux.py

@@ -1,5 +1,3 @@
-# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/flux.py
-
 import functools
 import unittest
 

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/hunyuan_video.py

@@ -1,4 +1,3 @@
-# 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

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/mochi.py

@@ -1,4 +1,3 @@
-# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/mochi.py
 import functools
 import unittest
 

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/wan.py

@@ -1,5 +1,3 @@
-# Adapted from: https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache/wan.py
-
 import functools
 import unittest
 
@@ -94,6 +92,6 @@ def apply_db_cache_on_pipe(
         pipe.__class__._is_cached = True
 
     if not shallow_patch:
-        apply_db_cache_on_transformer(pipe.transformer, **kwargs)
+        apply_db_cache_on_transformer(pipe.transformer)
 
     return pipe

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/__init__.py

@@ -1,5 +1,3 @@
-# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/__init__.py
-
 import importlib
 
 from diffusers import DiffusionPipeline

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/cogvideox.py

@@ -1,5 +1,3 @@
-# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/cogvideox.py
-
 import functools
 import unittest
 

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/flux.py

@@ -1,5 +1,3 @@
-# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/flux.py
-
 import functools
 import unittest
 

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/hunyuan_video.py

@@ -1,4 +1,3 @@
-# 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

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/mochi.py

@@ -1,5 +1,3 @@
-# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/mochi.py
-
 import functools
 import unittest
 

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/wan.py

@@ -1,5 +1,3 @@
-# Adapted from: https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache/wan.py
-
 import functools
 import unittest