cache-dit 0.2.28__py3-none-any.whl → 0.2.30__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.30'
+__version_tuple__ = version_tuple = (0, 2, 30)
 
 __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,18 @@ 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,
+        transformer=pipe.transformer,
+        blocks=[
+            pipe.transformer.transformer_blocks,
+            pipe.transformer.single_transformer_blocks,
+        ],
+        forward_pattern=[
+            ForwardPattern.Pattern_0,
+            ForwardPattern.Pattern_0,
+        ],
+        # The type hint in diffusers is wrong
+        check_num_outputs=False,
+        **kwargs,
     )
 
 
@@ -135,10 +147,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 +162,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 +176,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 +190,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 +204,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 +219,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 +234,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 +248,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 +263,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 +277,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 +291,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 +305,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,51 +318,20 @@ 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,
-    )
-
-
-@BlockAdapterRegistry.register("HunyuanDiT")
-def hunyuandit_adapter(pipe, **kwargs) -> BlockAdapter:
-    from diffusers import HunyuanDiT2DModel, HunyuanDiT2DControlNetModel
-
-    assert isinstance(
-        pipe.transformer,
-        (HunyuanDiT2DModel, HunyuanDiT2DControlNetModel),
-    )
-    return BlockAdapter(
-        pipe=pipe,
-        transformer=pipe.transformer,
-        blocks=pipe.transformer.blocks,
-        blocks_name="blocks",
-        dummy_blocks_names=[],
-        forward_pattern=ForwardPattern.Pattern_3,
-    )
-
-
-@BlockAdapterRegistry.register("HunyuanDiTPAG")
-def hunyuanditpag_adapter(pipe, **kwargs) -> BlockAdapter:
-    from diffusers import HunyuanDiT2DModel
-
-    assert isinstance(pipe.transformer, HunyuanDiT2DModel)
-    return BlockAdapter(
-        pipe=pipe,
-        transformer=pipe.transformer,
-        blocks=pipe.transformer.blocks,
-        blocks_name="blocks",
-        dummy_blocks_names=[],
-        forward_pattern=ForwardPattern.Pattern_3,
+        blocks=[
+            pipe.transformer.transformer_blocks,
+            pipe.transformer.single_transformer_blocks,
+        ],
+        forward_pattern=[
+            ForwardPattern.Pattern_0,
+            ForwardPattern.Pattern_0,
+        ],
+        **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 +339,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 +353,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 +367,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,39 +381,24 @@ 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,
     )
 
 
-@BlockAdapterRegistry.register("Sana")
+@BlockAdapterRegistry.register("Sana", supported=False)
 def sana_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import SanaTransformer2DModel
 
+    # TODO: fix -> got multiple values for argument 'encoder_hidden_states'
+
     assert isinstance(pipe.transformer, SanaTransformer2DModel)
     return BlockAdapter(
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
-        blocks_name="transformer_blocks",
-        dummy_blocks_names=[],
-        forward_pattern=ForwardPattern.Pattern_3,
-    )
-
-
-@BlockAdapterRegistry.register("ShapE")
-def shape_adapter(pipe, **kwargs) -> BlockAdapter:
-    from diffusers import PriorTransformer
-
-    assert isinstance(pipe.prior, PriorTransformer)
-    return 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,30 +411,45 @@ 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
+    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,
-    )
+    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("AuraFlow")
@@ -494,19 +460,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,20 +478,36 @@ 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,
+    )
+
+
+@BlockAdapterRegistry.register("ShapE")
+def shape_adapter(pipe, **kwargs) -> BlockAdapter:
+    from diffusers import PriorTransformer
+
+    assert isinstance(pipe.prior, PriorTransformer)
+    return BlockAdapter(
+        pipe=pipe,
+        transformer=pipe.prior,
+        blocks=pipe.prior.transformer_blocks,
+        forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
     )
 
 
-@BlockAdapterRegistry.register("HiDream")
+@BlockAdapterRegistry.register("HiDream", supported=True)
 def hidream_adapter(pipe, **kwargs) -> BlockAdapter:
+    # NOTE: Need to patch Transformer forward to fully support
+    # double_stream_blocks and single_stream_blocks, namely, need
+    # to remove the logics inside the blocks forward loop:
+    # https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_hidream_image.py#L893
+    # https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_hidream_image.py#L927
     from diffusers import HiDreamImageTransformer2DModel
 
     assert isinstance(pipe.transformer, HiDreamImageTransformer2DModel)
@@ -543,13 +515,47 @@ def hidream_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=[
-            pipe.transformer.double_stream_blocks,
+            # pipe.transformer.double_stream_blocks,
             pipe.transformer.single_stream_blocks,
         ],
-        dummy_blocks_names=[],
         forward_pattern=[
-            ForwardPattern.Pattern_4,
+            # ForwardPattern.Pattern_4,
             ForwardPattern.Pattern_3,
         ],
+        # The type hint in diffusers is wrong
         check_num_outputs=False,
+        **kwargs,
+    )
+
+
+@BlockAdapterRegistry.register("HunyuanDiT", supported=False)
+def hunyuandit_adapter(pipe, **kwargs) -> BlockAdapter:
+    # TODO: Patch Transformer forward
+    from diffusers import HunyuanDiT2DModel, HunyuanDiT2DControlNetModel
+
+    assert isinstance(
+        pipe.transformer,
+        (HunyuanDiT2DModel, HunyuanDiT2DControlNetModel),
+    )
+    return BlockAdapter(
+        pipe=pipe,
+        transformer=pipe.transformer,
+        blocks=pipe.transformer.blocks,
+        forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
+    )
+
+
+@BlockAdapterRegistry.register("HunyuanDiTPAG", supported=False)
+def hunyuanditpag_adapter(pipe, **kwargs) -> BlockAdapter:
+    # TODO: Patch Transformer forward
+    from diffusers import HunyuanDiT2DModel
+
+    assert isinstance(pipe.transformer, HunyuanDiT2DModel)
+    return BlockAdapter(
+        pipe=pipe,
+        transformer=pipe.transformer,
+        blocks=pipe.transformer.blocks,
+        forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
     )