cache-dit 0.2.14__py3-none-any.whl → 0.2.16__py3-none-any.whl

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/__init__.py

@@ -15,6 +15,8 @@ def apply_db_cache_on_transformer(transformer, *args, **kwargs):
         adapter_name = "wan"
     elif transformer_cls_name.startswith("HunyuanVideo"):
         adapter_name = "hunyuan_video"
+    elif transformer_cls_name.startswith("QwenImage"):
+        adapter_name = "qwen_image"
     else:
         raise ValueError(
             f"Unknown transformer class name: {transformer_cls_name}"
@@ -41,6 +43,8 @@ def apply_db_cache_on_pipe(pipe: DiffusionPipeline, *args, **kwargs):
         adapter_name = "wan"
     elif pipe_cls_name.startswith("HunyuanVideo"):
         adapter_name = "hunyuan_video"
+    elif pipe_cls_name.startswith("QwenImage"):
+        adapter_name = "qwen_image"
     else:
         raise ValueError(f"Unknown pipeline class name: {pipe_cls_name}")
 

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/cogvideox.py

@@ -4,7 +4,10 @@ import unittest
 import torch
 from diffusers import CogVideoXTransformer3DModel, DiffusionPipeline
 
-from cache_dit.cache_factory.dual_block_cache import cache_context
+from cache_dit.cache_factory.dual_block_cache import (
+    cache_context,
+    DBCachedTransformerBlocks,
+)
 
 
 def apply_db_cache_on_transformer(
@@ -15,7 +18,7 @@ def apply_db_cache_on_transformer(
 
     cached_transformer_blocks = torch.nn.ModuleList(
         [
-            cache_context.DBCachedTransformerBlocks(
+            DBCachedTransformerBlocks(
                 transformer.transformer_blocks,
                 transformer=transformer,
             )

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/flux.py

@@ -4,7 +4,15 @@ import unittest
 import torch
 from diffusers import DiffusionPipeline, FluxTransformer2DModel
 
-from cache_dit.cache_factory.dual_block_cache import cache_context
+from cache_dit.cache_factory.patch.flux import maybe_patch_flux_transformer
+from cache_dit.cache_factory.dual_block_cache import (
+    cache_context,
+    DBCachedTransformerBlocks,
+)
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
 
 
 def apply_db_cache_on_transformer(
@@ -13,11 +21,13 @@ def apply_db_cache_on_transformer(
     if getattr(transformer, "_is_cached", False):
         return transformer
 
+    transformer = maybe_patch_flux_transformer(transformer)
+
     cached_transformer_blocks = torch.nn.ModuleList(
         [
-            cache_context.DBCachedTransformerBlocks(
-                transformer.transformer_blocks,
-                transformer.single_transformer_blocks,
+            DBCachedTransformerBlocks(
+                transformer.transformer_blocks
+                + transformer.single_transformer_blocks,
                 transformer=transformer,
                 return_hidden_states_first=False,
             )

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/hunyuan_video.py

@@ -11,7 +11,10 @@ from diffusers.utils import (
     USE_PEFT_BACKEND,
 )
 
-from cache_dit.cache_factory.dual_block_cache import cache_context
+from cache_dit.cache_factory.dual_block_cache import (
+    cache_context,
+    DBCachedTransformerBlocks,
+)
 from cache_dit.logger import init_logger
 
 try:
@@ -44,7 +47,7 @@ def apply_db_cache_on_transformer(
 
     cached_transformer_blocks = torch.nn.ModuleList(
         [
-            cache_context.DBCachedTransformerBlocks(
+            DBCachedTransformerBlocks(
                 transformer.transformer_blocks
                 + transformer.single_transformer_blocks,
                 transformer=transformer,

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/mochi.py

@@ -4,7 +4,10 @@ import unittest
 import torch
 from diffusers import DiffusionPipeline, MochiTransformer3DModel
 
-from cache_dit.cache_factory.dual_block_cache import cache_context
+from cache_dit.cache_factory.dual_block_cache import (
+    cache_context,
+    DBCachedTransformerBlocks,
+)
 
 
 def apply_db_cache_on_transformer(
@@ -15,7 +18,7 @@ def apply_db_cache_on_transformer(
 
     cached_transformer_blocks = torch.nn.ModuleList(
         [
-            cache_context.DBCachedTransformerBlocks(
+            DBCachedTransformerBlocks(
                 transformer.transformer_blocks,
                 transformer=transformer,
             )

cache_dit/cache_factory/{first_block_cache/diffusers_adapters/mochi.py → dual_block_cache/diffusers_adapters/qwen_image.py}

@@ -1,25 +1,27 @@
-# Adapted from: https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache/mochi.py
-
 import functools
 import unittest
 
 import torch
-from diffusers import DiffusionPipeline, MochiTransformer3DModel
+from diffusers import QwenImagePipeline, QwenImageTransformer2DModel
 
-from cache_dit.cache_factory.first_block_cache import cache_context
+from cache_dit.cache_factory.dual_block_cache import (
+    cache_context,
+    DBCachedTransformerBlocks,
+)
 
 
-def apply_cache_on_transformer(
-    transformer: MochiTransformer3DModel,
+def apply_db_cache_on_transformer(
+    transformer: QwenImageTransformer2DModel,
 ):
     if getattr(transformer, "_is_cached", False):
         return transformer
 
-    cached_transformer_blocks = torch.nn.ModuleList(
+    transformer_blocks = torch.nn.ModuleList(
         [
-            cache_context.CachedTransformerBlocks(
+            DBCachedTransformerBlocks(
                 transformer.transformer_blocks,
                 transformer=transformer,
+                return_hidden_states_first=False,
             )
         ]
     )
@@ -35,7 +37,7 @@ def apply_cache_on_transformer(
         with unittest.mock.patch.object(
             self,
             "transformer_blocks",
-            cached_transformer_blocks,
+            transformer_blocks,
         ):
             return original_forward(
                 *args,
@@ -49,8 +51,8 @@ def apply_cache_on_transformer(
     return transformer
 
 
-def apply_cache_on_pipe(
-    pipe: DiffusionPipeline,
+def apply_db_cache_on_pipe(
+    pipe: QwenImagePipeline,
     *,
     shallow_patch: bool = False,
     residual_diff_threshold=0.06,
@@ -84,6 +86,6 @@ def apply_cache_on_pipe(
         pipe.__class__._is_cached = True
 
     if not shallow_patch:
-        apply_cache_on_transformer(pipe.transformer, **kwargs)
+        apply_db_cache_on_transformer(pipe.transformer)
 
     return pipe

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/wan.py

@@ -2,9 +2,12 @@ import functools
 import unittest
 
 import torch
-from diffusers import DiffusionPipeline, WanTransformer3DModel
+from diffusers import WanPipeline, WanTransformer3DModel
 
-from cache_dit.cache_factory.dual_block_cache import cache_context
+from cache_dit.cache_factory.dual_block_cache import (
+    cache_context,
+    DBCachedTransformerBlocks,
+)
 
 
 def apply_db_cache_on_transformer(
@@ -15,7 +18,7 @@ def apply_db_cache_on_transformer(
 
     blocks = torch.nn.ModuleList(
         [
-            cache_context.DBCachedTransformerBlocks(
+            DBCachedTransformerBlocks(
                 transformer.blocks,
                 transformer=transformer,
                 return_hidden_states_only=True,
@@ -49,7 +52,7 @@ def apply_db_cache_on_transformer(
 
 
 def apply_db_cache_on_pipe(
-    pipe: DiffusionPipeline,
+    pipe: WanPipeline,
     *,
     shallow_patch: bool = False,
     residual_diff_threshold=0.03,

cache_dit/cache_factory/dynamic_block_prune/__init__.py

@@ -0,0 +1,4 @@
+from cache_dit.cache_factory.dynamic_block_prune.prune_blocks import (
+    prune_context,
+    DBPrunedTransformerBlocks,
+)

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/__init__.py

@@ -15,6 +15,8 @@ def apply_db_prune_on_transformer(transformer, *args, **kwargs):
         adapter_name = "wan"
     elif transformer_cls_name.startswith("HunyuanVideo"):
         adapter_name = "hunyuan_video"
+    elif transformer_cls_name.startswith("QwenImage"):
+        adapter_name = "qwen_image"
     else:
         raise ValueError(
             f"Unknown transformer class name: {transformer_cls_name}"
@@ -41,6 +43,8 @@ def apply_db_prune_on_pipe(pipe: DiffusionPipeline, *args, **kwargs):
         adapter_name = "wan"
     elif pipe_cls_name.startswith("HunyuanVideo"):
         adapter_name = "hunyuan_video"
+    elif pipe_cls_name.startswith("QwenImage"):
+        adapter_name = "qwen_image"
     else:
         raise ValueError(f"Unknown pipeline class name: {pipe_cls_name}")
 

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/cogvideox.py

@@ -4,7 +4,10 @@ import unittest
 import torch
 from diffusers import CogVideoXTransformer3DModel, DiffusionPipeline
 
-from cache_dit.cache_factory.dynamic_block_prune import prune_context
+from cache_dit.cache_factory.dynamic_block_prune import (
+    prune_context,
+    DBPrunedTransformerBlocks,
+)
 
 
 def apply_db_prune_on_transformer(
@@ -15,7 +18,7 @@ def apply_db_prune_on_transformer(
 
     cached_transformer_blocks = torch.nn.ModuleList(
         [
-            prune_context.DBPrunedTransformerBlocks(
+            DBPrunedTransformerBlocks(
                 transformer.transformer_blocks,
                 transformer=transformer,
             )

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/flux.py

@@ -4,7 +4,11 @@ import unittest
 import torch
 from diffusers import DiffusionPipeline, FluxTransformer2DModel
 
-from cache_dit.cache_factory.dynamic_block_prune import prune_context
+from cache_dit.cache_factory.patch.flux import maybe_patch_flux_transformer
+from cache_dit.cache_factory.dynamic_block_prune import (
+    prune_context,
+    DBPrunedTransformerBlocks,
+)
 
 
 def apply_db_prune_on_transformer(
@@ -13,11 +17,13 @@ def apply_db_prune_on_transformer(
     if getattr(transformer, "_is_pruned", False):
         return transformer
 
+    transformer = maybe_patch_flux_transformer(transformer)
+
     cached_transformer_blocks = torch.nn.ModuleList(
         [
-            prune_context.DBPrunedTransformerBlocks(
-                transformer.transformer_blocks,
-                transformer.single_transformer_blocks,
+            DBPrunedTransformerBlocks(
+                transformer.transformer_blocks
+                + transformer.single_transformer_blocks,
                 transformer=transformer,
                 return_hidden_states_first=False,
             )

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/hunyuan_video.py

@@ -11,7 +11,10 @@ from diffusers.utils import (
     USE_PEFT_BACKEND,
 )
 
-from cache_dit.cache_factory.dynamic_block_prune import prune_context
+from cache_dit.cache_factory.dynamic_block_prune import (
+    prune_context,
+    DBPrunedTransformerBlocks,
+)
 from cache_dit.logger import init_logger
 
 try:
@@ -44,7 +47,7 @@ def apply_db_prune_on_transformer(
 
     cached_transformer_blocks = torch.nn.ModuleList(
         [
-            prune_context.DBPrunedTransformerBlocks(
+            DBPrunedTransformerBlocks(
                 transformer.transformer_blocks
                 + transformer.single_transformer_blocks,
                 transformer=transformer,

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/mochi.py

@@ -4,7 +4,10 @@ import unittest
 import torch
 from diffusers import DiffusionPipeline, MochiTransformer3DModel
 
-from cache_dit.cache_factory.dynamic_block_prune import prune_context
+from cache_dit.cache_factory.dynamic_block_prune import (
+    prune_context,
+    DBPrunedTransformerBlocks,
+)
 
 
 def apply_db_prune_on_transformer(
@@ -15,7 +18,7 @@ def apply_db_prune_on_transformer(
 
     cached_transformer_blocks = torch.nn.ModuleList(
         [
-            prune_context.DBPrunedTransformerBlocks(
+            DBPrunedTransformerBlocks(
                 transformer.transformer_blocks,
                 transformer=transformer,
             )

cache_dit/cache_factory/{first_block_cache/diffusers_adapters/cogvideox.py → dynamic_block_prune/diffusers_adapters/qwen_image.py}

@@ -1,25 +1,30 @@
-# Adapted from: https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache/cogvideox.py
-
 import functools
 import unittest
 
 import torch
-from diffusers import CogVideoXTransformer3DModel, DiffusionPipeline
+from diffusers import QwenImagePipeline, QwenImageTransformer2DModel
+
+from cache_dit.cache_factory.dynamic_block_prune import (
+    prune_context,
+    DBPrunedTransformerBlocks,
+)
+from cache_dit.logger import init_logger
 
-from cache_dit.cache_factory.first_block_cache import cache_context
+logger = init_logger(__name__)
 
 
-def apply_cache_on_transformer(
-    transformer: CogVideoXTransformer3DModel,
+def apply_db_prune_on_transformer(
+    transformer: QwenImageTransformer2DModel,
 ):
-    if getattr(transformer, "_is_cached", False):
+    if getattr(transformer, "_is_pruned", False):
         return transformer
 
-    cached_transformer_blocks = torch.nn.ModuleList(
+    transformer_blocks = torch.nn.ModuleList(
         [
-            cache_context.CachedTransformerBlocks(
+            DBPrunedTransformerBlocks(
                 transformer.transformer_blocks,
                 transformer=transformer,
+                return_hidden_states_first=False,
             )
         ]
     )
@@ -35,7 +40,7 @@ def apply_cache_on_transformer(
         with unittest.mock.patch.object(
             self,
             "transformer_blocks",
-            cached_transformer_blocks,
+            transformer_blocks,
         ):
             return original_forward(
                 *args,
@@ -44,46 +49,46 @@ def apply_cache_on_transformer(
 
     transformer.forward = new_forward.__get__(transformer)
 
-    transformer._is_cached = True
+    transformer._is_pruned = True
 
     return transformer
 
 
-def apply_cache_on_pipe(
-    pipe: DiffusionPipeline,
+def apply_db_prune_on_pipe(
+    pipe: QwenImagePipeline,
     *,
     shallow_patch: bool = False,
-    residual_diff_threshold=0.04,
+    residual_diff_threshold=0.06,
     downsample_factor=1,
     warmup_steps=0,
-    max_cached_steps=-1,
+    max_pruned_steps=-1,
     **kwargs,
 ):
-    cache_kwargs, kwargs = cache_context.collect_cache_kwargs(
+    prune_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,
+            "max_pruned_steps": max_pruned_steps,
         },
         **kwargs,
     )
-    if not getattr(pipe, "_is_cached", False):
+    if not getattr(pipe, "_is_pruned", 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,
+            with prune_context.prune_context(
+                prune_context.create_prune_context(
+                    **prune_kwargs,
                 )
             ):
                 return original_call(self, *args, **kwargs)
 
         pipe.__class__.__call__ = new_call
-        pipe.__class__._is_cached = True
+        pipe.__class__._is_pruned = True
 
     if not shallow_patch:
-        apply_cache_on_transformer(pipe.transformer, **kwargs)
+        apply_db_prune_on_transformer(pipe.transformer)
 
     return pipe

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/wan.py

@@ -4,7 +4,10 @@ import unittest
 import torch
 from diffusers import DiffusionPipeline, WanTransformer3DModel
 
-from cache_dit.cache_factory.dynamic_block_prune import prune_context
+from cache_dit.cache_factory.dynamic_block_prune import (
+    prune_context,
+    DBPrunedTransformerBlocks,
+)
 
 
 def apply_db_prune_on_transformer(
@@ -15,7 +18,7 @@ def apply_db_prune_on_transformer(
 
     blocks = torch.nn.ModuleList(
         [
-            prune_context.DBPrunedTransformerBlocks(
+            DBPrunedTransformerBlocks(
                 transformer.blocks,
                 transformer=transformer,
                 return_hidden_states_only=True,