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

cache_dit/__init__.py

@@ -6,10 +6,12 @@ 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
 from cache_dit.cache_factory import BlockAdapter
+from cache_dit.cache_factory import ParamsModifier
 from cache_dit.cache_factory import ForwardPattern
 from cache_dit.cache_factory import PatchFunctor
 from cache_dit.cache_factory import supported_pipelines

cache_dit/_version.py

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

cache_dit/cache_factory/__init__.py

@@ -7,14 +7,17 @@ from cache_dit.cache_factory.forward_pattern import ForwardPattern
 from cache_dit.cache_factory.patch_functors import PatchFunctor
 
 from cache_dit.cache_factory.block_adapters import BlockAdapter
+from cache_dit.cache_factory.block_adapters import ParamsModifier
 from cache_dit.cache_factory.block_adapters import BlockAdapterRegistry
 
 from cache_dit.cache_factory.cache_contexts import CachedContext
+from cache_dit.cache_factory.cache_contexts import CachedContextManager
 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

@@ -1,5 +1,6 @@
 from cache_dit.cache_factory.forward_pattern import ForwardPattern
 from cache_dit.cache_factory.block_adapters.block_adapters import BlockAdapter
+from cache_dit.cache_factory.block_adapters.block_adapters import ParamsModifier
 from cache_dit.cache_factory.block_adapters.block_registers import (
     BlockAdapterRegistry,
 )
@@ -8,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")
@@ -36,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,
     )
 
 
@@ -51,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,
     )
 
 
@@ -69,15 +82,27 @@ def wan_adapter(pipe, **kwargs) -> BlockAdapter:
         (WanTransformer3DModel, WanVACETransformer3DModel),
     )
     if getattr(pipe, "transformer_2", None):
-        # Wan 2.2, cache for low-noise transformer
+        assert isinstance(
+            pipe.transformer_2,
+            (WanTransformer3DModel, WanVACETransformer3DModel),
+        )
+        # Wan 2.2 MoE
         return BlockAdapter(
             pipe=pipe,
-            transformer=pipe.transformer_2,
-            blocks=pipe.transformer_2.blocks,
-            blocks_name="blocks",
-            dummy_blocks_names=[],
-            forward_pattern=ForwardPattern.Pattern_2,
+            transformer=[
+                pipe.transformer,
+                pipe.transformer_2,
+            ],
+            blocks=[
+                pipe.transformer.blocks,
+                pipe.transformer_2.blocks,
+            ],
+            forward_pattern=[
+                ForwardPattern.Pattern_2,
+                ForwardPattern.Pattern_2,
+            ],
             has_separate_cfg=True,
+            **kwargs,
         )
     else:
         # Wan 2.1
@@ -85,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,
         )
 
 
@@ -99,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,
     )
 
 
@@ -118,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,
     )
 
 
@@ -134,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,
     )
 
 
@@ -149,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,
     )
 
 
@@ -164,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,
     )
 
 
@@ -179,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,
     )
 
 
@@ -195,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,
     )
 
 
@@ -211,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,
     )
 
 
@@ -226,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,
     )
 
 
@@ -242,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,
     )
 
 
@@ -257,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,
     )
 
 
@@ -272,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,
     )
 
 
@@ -287,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,
     )
 
 
@@ -301,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,
     )
 
 
@@ -323,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,
     )
 
 
@@ -338,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)
@@ -353,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,
     )
 
 
@@ -368,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,
     )
 
 
@@ -383,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,
     )
 
 
@@ -398,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,
     )
 
 
@@ -413,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,
     )
 
 
@@ -428,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,
     )
 
 
@@ -443,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,
     )
 
 
@@ -477,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,
     )
 
 
@@ -505,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,
     )
 
 
@@ -529,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,
     )