cache-dit 0.2.28__py3-none-any.whl → 0.2.29__py3-none-any.whl

cache_dit/__init__.py

@@ -6,6 +6,7 @@ except ImportError:
 
 from cache_dit.cache_factory import load_options
 from cache_dit.cache_factory import enable_cache
+from cache_dit.cache_factory import disable_cache
 from cache_dit.cache_factory import cache_type
 from cache_dit.cache_factory import block_range
 from cache_dit.cache_factory import CacheType

cache_dit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.28'
-__version_tuple__ = version_tuple = (0, 2, 28)
+__version__ = version = '0.2.29'
+__version_tuple__ = version_tuple = (0, 2, 29)
 
 __commit_id__ = commit_id = None

cache_dit/cache_factory/__init__.py

@@ -17,6 +17,7 @@ from cache_dit.cache_factory.cache_blocks import CachedBlocks
 from cache_dit.cache_factory.cache_adapters import CachedAdapter
 
 from cache_dit.cache_factory.cache_interface import enable_cache
+from cache_dit.cache_factory.cache_interface import disable_cache
 from cache_dit.cache_factory.cache_interface import supported_pipelines
 from cache_dit.cache_factory.cache_interface import get_adapter
 

cache_dit/cache_factory/block_adapters/__init__.py

@@ -9,23 +9,37 @@ from cache_dit.cache_factory.block_adapters.block_registers import (
 @BlockAdapterRegistry.register("Flux")
 def flux_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import FluxTransformer2DModel
-    from cache_dit.cache_factory.patch_functors import FluxPatchFunctor
+    from cache_dit.utils import is_diffusers_at_least_0_3_5
 
     assert isinstance(pipe.transformer, FluxTransformer2DModel)
-
-    return BlockAdapter(
-        pipe=pipe,
-        transformer=pipe.transformer,
-        blocks=(
-            pipe.transformer.transformer_blocks
-            + pipe.transformer.single_transformer_blocks
-        ),
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=["single_transformer_blocks"],
-        patch_functor=FluxPatchFunctor(),
-        forward_pattern=ForwardPattern.Pattern_1,
-        disable_patch=kwargs.pop("disable_patch", False),
-    )
+    if is_diffusers_at_least_0_3_5():
+        return BlockAdapter(
+            pipe=pipe,
+            transformer=pipe.transformer,
+            blocks=[
+                pipe.transformer.transformer_blocks,
+                pipe.transformer.single_transformer_blocks,
+            ],
+            forward_pattern=[
+                ForwardPattern.Pattern_1,
+                ForwardPattern.Pattern_1,
+            ],
+            **kwargs,
+        )
+    else:
+        return BlockAdapter(
+            pipe=pipe,
+            transformer=pipe.transformer,
+            blocks=[
+                pipe.transformer.transformer_blocks,
+                pipe.transformer.single_transformer_blocks,
+            ],
+            forward_pattern=[
+                ForwardPattern.Pattern_1,
+                ForwardPattern.Pattern_3,
+            ],
+            **kwargs,
+        )
 
 
 @BlockAdapterRegistry.register("Mochi")
@@ -37,9 +51,8 @@ def mochi_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_0,
+        **kwargs,
     )
 
 
@@ -52,9 +65,8 @@ def cogvideox_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_0,
+        **kwargs,
     )
 
 
@@ -85,16 +97,12 @@ def wan_adapter(pipe, **kwargs) -> BlockAdapter:
                 pipe.transformer.blocks,
                 pipe.transformer_2.blocks,
             ],
-            blocks_name=[
-                "blocks",
-                "blocks",
-            ],
             forward_pattern=[
                 ForwardPattern.Pattern_2,
                 ForwardPattern.Pattern_2,
             ],
-            dummy_blocks_names=[],
             has_separate_cfg=True,
+            **kwargs,
         )
     else:
         # Wan 2.1
@@ -102,10 +110,9 @@ def wan_adapter(pipe, **kwargs) -> BlockAdapter:
             pipe=pipe,
             transformer=pipe.transformer,
             blocks=pipe.transformer.blocks,
-            blocks_name="blocks",
-            dummy_blocks_names=[],
             forward_pattern=ForwardPattern.Pattern_2,
             has_separate_cfg=True,
+            **kwargs,
         )
 
 
@@ -116,13 +123,15 @@ def hunyuanvideo_adapter(pipe, **kwargs) -> BlockAdapter:
     assert isinstance(pipe.transformer, HunyuanVideoTransformer3DModel)
     return BlockAdapter(
         pipe=pipe,
-        blocks=(
-            pipe.transformer.transformer_blocks
-            + pipe.transformer.single_transformer_blocks
-        ),
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=["single_transformer_blocks"],
-        forward_pattern=ForwardPattern.Pattern_0,
+        blocks=[
+            pipe.transformer.transformer_blocks,
+            pipe.transformer.single_transformer_blocks,
+        ],
+        forward_pattern=[
+            ForwardPattern.Pattern_0,
+            ForwardPattern.Pattern_0,
+        ],
+        **kwargs,
     )
 
 
@@ -135,10 +144,9 @@ def qwenimage_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_1,
         has_separate_cfg=True,
+        **kwargs,
     )
 
 
@@ -151,9 +159,8 @@ def ltxvideo_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_2,
+        **kwargs,
     )
 
 
@@ -166,9 +173,8 @@ def allegro_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_2,
+        **kwargs,
     )
 
 
@@ -181,9 +187,8 @@ def cogview3plus_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_0,
+        **kwargs,
     )
 
 
@@ -196,10 +201,9 @@ def cogview4_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_0,
         has_separate_cfg=True,
+        **kwargs,
     )
 
 
@@ -212,10 +216,9 @@ def cosmos_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_2,
         has_separate_cfg=True,
+        **kwargs,
     )
 
 
@@ -228,9 +231,8 @@ def easyanimate_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_0,
+        **kwargs,
     )
 
 
@@ -243,10 +245,9 @@ def skyreelsv2_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.blocks,
-        blocks_name="blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_2,
         has_separate_cfg=True,
+        **kwargs,
     )
 
 
@@ -259,9 +260,8 @@ def sd3_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_1,
+        **kwargs,
     )
 
 
@@ -274,9 +274,8 @@ def consisid_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_0,
+        **kwargs,
     )
 
 
@@ -289,9 +288,8 @@ def dit_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
     )
 
 
@@ -304,9 +302,8 @@ def amused_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_layers,
-        blocks_name="transformer_layers",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
     )
 
 
@@ -318,13 +315,15 @@ def bria_adapter(pipe, **kwargs) -> BlockAdapter:
     return BlockAdapter(
         pipe=pipe,
         transformer=pipe.transformer,
-        blocks=(
-            pipe.transformer.transformer_blocks
-            + pipe.transformer.single_transformer_blocks
-        ),
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=["single_transformer_blocks"],
-        forward_pattern=ForwardPattern.Pattern_0,
+        blocks=[
+            pipe.transformer.transformer_blocks,
+            pipe.transformer.single_transformer_blocks,
+        ],
+        forward_pattern=[
+            ForwardPattern.Pattern_0,
+            ForwardPattern.Pattern_0,
+        ],
+        **kwargs,
     )
 
 
@@ -340,9 +339,8 @@ def hunyuandit_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.blocks,
-        blocks_name="blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
     )
 
 
@@ -355,14 +353,13 @@ def hunyuanditpag_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.blocks,
-        blocks_name="blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
     )
 
 
 @BlockAdapterRegistry.register("Lumina")
-def lumina_adapter(pipe) -> BlockAdapter:
+def lumina_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import LuminaNextDiT2DModel
 
     assert isinstance(pipe.transformer, LuminaNextDiT2DModel)
@@ -370,9 +367,8 @@ def lumina_adapter(pipe) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.layers,
-        blocks_name="layers",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
     )
 
 
@@ -385,9 +381,8 @@ def lumina2_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.layers,
-        blocks_name="layers",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
     )
 
 
@@ -400,9 +395,8 @@ def omnigen_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.layers,
-        blocks_name="layers",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
     )
 
 
@@ -415,9 +409,8 @@ def pixart_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
     )
 
 
@@ -430,9 +423,8 @@ def sana_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
     )
 
 
@@ -445,9 +437,8 @@ def shape_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.prior,
         blocks=pipe.prior.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
     )
 
 
@@ -460,29 +451,28 @@ def stabledudio_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
         forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
     )
 
 
 @BlockAdapterRegistry.register("VisualCloze")
 def visualcloze_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import FluxTransformer2DModel
-    from cache_dit.cache_factory.patch_functors import FluxPatchFunctor
 
     assert isinstance(pipe.transformer, FluxTransformer2DModel)
     return BlockAdapter(
         pipe=pipe,
         transformer=pipe.transformer,
-        blocks=(
-            pipe.transformer.transformer_blocks
-            + pipe.transformer.single_transformer_blocks
-        ),
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=["single_transformer_blocks"],
-        patch_functor=FluxPatchFunctor(),
-        forward_pattern=ForwardPattern.Pattern_1,
+        blocks=[
+            pipe.transformer.transformer_blocks,
+            pipe.transformer.single_transformer_blocks,
+        ],
+        forward_pattern=[
+            ForwardPattern.Pattern_1,
+            ForwardPattern.Pattern_3,
+        ],
+        **kwargs,
     )
 
 
@@ -494,19 +484,9 @@ def auraflow_adapter(pipe, **kwargs) -> BlockAdapter:
     return BlockAdapter(
         pipe=pipe,
         transformer=pipe.transformer,
-        blocks=[
-            # Only 4 joint blocks, apply no-cache
-            # pipe.transformer.joint_transformer_blocks,
-            pipe.transformer.single_transformer_blocks,
-        ],
-        blocks_name=[
-            # "joint_transformer_blocks",
-            "single_transformer_blocks",
-        ],
-        forward_pattern=[
-            # ForwardPattern.Pattern_1,
-            ForwardPattern.Pattern_3,
-        ],
+        blocks=pipe.transformer.single_transformer_blocks,
+        forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
     )
 
 
@@ -522,15 +502,12 @@ def chroma_adapter(pipe, **kwargs) -> BlockAdapter:
             pipe.transformer.transformer_blocks,
             pipe.transformer.single_transformer_blocks,
         ],
-        blocks_name=[
-            "transformer_blocks",
-            "single_transformer_blocks",
-        ],
         forward_pattern=[
             ForwardPattern.Pattern_1,
             ForwardPattern.Pattern_3,
         ],
         has_separate_cfg=True,
+        **kwargs,
     )
 
 
@@ -546,10 +523,10 @@ def hidream_adapter(pipe, **kwargs) -> BlockAdapter:
             pipe.transformer.double_stream_blocks,
             pipe.transformer.single_stream_blocks,
         ],
-        dummy_blocks_names=[],
         forward_pattern=[
             ForwardPattern.Pattern_4,
             ForwardPattern.Pattern_3,
         ],
         check_num_outputs=False,
+        **kwargs,
     )