cache-dit 0.3.2__py3-none-any.whl → 1.0.14__py3-none-any.whl

cache_dit/__init__.py

@@ -4,32 +4,50 @@ except ImportError:
     __version__ = "unknown version"
     version_tuple = (0, 0, "unknown version")
 
-from cache_dit.utils import summary
-from cache_dit.utils import strify
+
 from cache_dit.utils import disable_print
 from cache_dit.logger import init_logger
-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 BasicCacheConfig
-from cache_dit.cache_factory import CalibratorConfig
-from cache_dit.cache_factory import TaylorSeerCalibratorConfig
-from cache_dit.cache_factory import FoCaCalibratorConfig
-from cache_dit.cache_factory import supported_pipelines
-from cache_dit.cache_factory import get_adapter
+from cache_dit.caching import load_options
+from cache_dit.caching import enable_cache
+from cache_dit.caching import disable_cache
+from cache_dit.caching import cache_type
+from cache_dit.caching import block_range
+from cache_dit.caching import CacheType
+from cache_dit.caching import BlockAdapter
+from cache_dit.caching import ParamsModifier
+from cache_dit.caching import ForwardPattern
+from cache_dit.caching import PatchFunctor
+from cache_dit.caching import BasicCacheConfig
+from cache_dit.caching import DBCacheConfig
+from cache_dit.caching import DBPruneConfig
+from cache_dit.caching import CalibratorConfig
+from cache_dit.caching import TaylorSeerCalibratorConfig
+from cache_dit.caching import FoCaCalibratorConfig
+from cache_dit.caching import supported_pipelines
+from cache_dit.caching import get_adapter
 from cache_dit.compile import set_compile_configs
-from cache_dit.quantize import quantize
+from cache_dit.parallelism import ParallelismBackend
+from cache_dit.parallelism import ParallelismConfig
+from cache_dit.summary import supported_matrix
+from cache_dit.summary import summary
+from cache_dit.summary import strify
+
+try:
+    from cache_dit.quantize import quantize
+except ImportError as e:  # noqa: F841
+    err_msg = str(e)
+
+    def quantize(*args, **kwargs):
+        raise ImportError(
+            "Quantization requires additional dependencies. "
+            "Please install cache-dit[quantization] or cache-dit[all] "
+            f"to use this feature. Error message: {err_msg}"
+        )
 
 
 NONE = CacheType.NONE
 DBCache = CacheType.DBCache
+DBPrune = CacheType.DBPrune
 
 Pattern_0 = ForwardPattern.Pattern_0
 Pattern_1 = ForwardPattern.Pattern_1

cache_dit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.2'
-__version_tuple__ = version_tuple = (0, 3, 2)
+__version__ = version = '1.0.14'
+__version_tuple__ = version_tuple = (1, 0, 14)
 
 __commit_id__ = commit_id = None

cache_dit/caching/__init__.py

@@ -0,0 +1,36 @@
+from cache_dit.caching.cache_types import CacheType
+from cache_dit.caching.cache_types import cache_type
+from cache_dit.caching.cache_types import block_range
+
+from cache_dit.caching.forward_pattern import ForwardPattern
+from cache_dit.caching.params_modifier import ParamsModifier
+from cache_dit.caching.patch_functors import PatchFunctor
+
+from cache_dit.caching.block_adapters import BlockAdapter
+from cache_dit.caching.block_adapters import BlockAdapterRegistry
+from cache_dit.caching.block_adapters import FakeDiffusionPipeline
+
+from cache_dit.caching.cache_contexts import BasicCacheConfig
+from cache_dit.caching.cache_contexts import DBCacheConfig
+from cache_dit.caching.cache_contexts import CachedContext
+from cache_dit.caching.cache_contexts import CachedContextManager
+from cache_dit.caching.cache_contexts import DBPruneConfig
+from cache_dit.caching.cache_contexts import PrunedContext
+from cache_dit.caching.cache_contexts import PrunedContextManager
+from cache_dit.caching.cache_contexts import ContextManager
+from cache_dit.caching.cache_contexts import CalibratorConfig
+from cache_dit.caching.cache_contexts import TaylorSeerCalibratorConfig
+from cache_dit.caching.cache_contexts import FoCaCalibratorConfig
+
+from cache_dit.caching.cache_blocks import CachedBlocks
+from cache_dit.caching.cache_blocks import PrunedBlocks
+from cache_dit.caching.cache_blocks import UnifiedBlocks
+
+from cache_dit.caching.cache_adapters import CachedAdapter
+
+from cache_dit.caching.cache_interface import enable_cache
+from cache_dit.caching.cache_interface import disable_cache
+from cache_dit.caching.cache_interface import supported_pipelines
+from cache_dit.caching.cache_interface import get_adapter
+
+from cache_dit.caching.utils import load_options

cache_dit/{cache_factory → caching}/block_adapters/__init__.py

@@ -1,7 +1,10 @@
-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 (
+from cache_dit.caching.forward_pattern import ForwardPattern
+from cache_dit.caching.block_adapters.block_adapters import BlockAdapter
+from cache_dit.caching.block_adapters.block_adapters import (
+    FakeDiffusionPipeline,
+)
+from cache_dit.caching.block_adapters.block_adapters import ParamsModifier
+from cache_dit.caching.block_adapters.block_registers import (
     BlockAdapterRegistry,
 )
 
@@ -12,7 +15,10 @@ def flux_adapter(pipe, **kwargs) -> BlockAdapter:
     from cache_dit.utils import is_diffusers_at_least_0_3_5
 
     assert isinstance(pipe.transformer, FluxTransformer2DModel)
-    if is_diffusers_at_least_0_3_5():
+    transformer_cls_name: str = pipe.transformer.__class__.__name__
+    if is_diffusers_at_least_0_3_5() and not transformer_cls_name.startswith(
+        "Nunchaku"
+    ):
         return BlockAdapter(
             pipe=pipe,
             transformer=pipe.transformer,
@@ -24,6 +30,7 @@ def flux_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_1,
                 ForwardPattern.Pattern_1,
             ],
+            check_forward_pattern=True,
             **kwargs,
         )
     else:
@@ -38,6 +45,7 @@ def flux_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_1,
                 ForwardPattern.Pattern_3,
             ],
+            check_forward_pattern=True,
             **kwargs,
         )
 
@@ -52,6 +60,7 @@ def mochi_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_0,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -66,6 +75,7 @@ def cogvideox_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_0,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -101,6 +111,7 @@ def wan_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_2,
                 ForwardPattern.Pattern_2,
             ],
+            check_forward_pattern=True,
             has_separate_cfg=True,
             **kwargs,
         )
@@ -111,6 +122,7 @@ def wan_adapter(pipe, **kwargs) -> BlockAdapter:
             transformer=pipe.transformer,
             blocks=pipe.transformer.blocks,
             forward_pattern=ForwardPattern.Pattern_2,
+            check_forward_pattern=True,
             has_separate_cfg=True,
             **kwargs,
         )
@@ -132,6 +144,7 @@ def hunyuanvideo_adapter(pipe, **kwargs) -> BlockAdapter:
             ForwardPattern.Pattern_0,
             ForwardPattern.Pattern_0,
         ],
+        check_forward_pattern=True,
         # The type hint in diffusers is wrong
         check_num_outputs=False,
         **kwargs,
@@ -143,14 +156,32 @@ def qwenimage_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import QwenImageTransformer2DModel
 
     assert isinstance(pipe.transformer, QwenImageTransformer2DModel)
-    return BlockAdapter(
-        pipe=pipe,
-        transformer=pipe.transformer,
-        blocks=pipe.transformer.transformer_blocks,
-        forward_pattern=ForwardPattern.Pattern_1,
-        has_separate_cfg=True,
-        **kwargs,
-    )
+
+    pipe_cls_name: str = pipe.__class__.__name__
+    if pipe_cls_name.startswith("QwenImageControlNet"):
+        from cache_dit.caching.patch_functors import (
+            QwenImageControlNetPatchFunctor,
+        )
+
+        return BlockAdapter(
+            pipe=pipe,
+            transformer=pipe.transformer,
+            blocks=pipe.transformer.transformer_blocks,
+            forward_pattern=ForwardPattern.Pattern_1,
+            patch_functor=QwenImageControlNetPatchFunctor(),
+            check_forward_pattern=True,
+            has_separate_cfg=True,
+        )
+    else:
+        return BlockAdapter(
+            pipe=pipe,
+            transformer=pipe.transformer,
+            blocks=pipe.transformer.transformer_blocks,
+            forward_pattern=ForwardPattern.Pattern_1,
+            check_forward_pattern=True,
+            has_separate_cfg=True,
+            **kwargs,
+        )
 
 
 @BlockAdapterRegistry.register("LTX")
@@ -163,6 +194,7 @@ def ltxvideo_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_2,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -177,6 +209,7 @@ def allegro_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_2,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -191,6 +224,7 @@ def cogview3plus_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_0,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -205,6 +239,7 @@ def cogview4_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_0,
+        check_forward_pattern=True,
         has_separate_cfg=True,
         **kwargs,
     )
@@ -220,6 +255,7 @@ def cosmos_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_2,
+        check_forward_pattern=True,
         has_separate_cfg=True,
         **kwargs,
     )
@@ -235,6 +271,7 @@ def easyanimate_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_0,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -252,6 +289,7 @@ def skyreelsv2_adapter(pipe, **kwargs) -> BlockAdapter:
         # encoder_hidden_states will never change in the blocks
         # forward loop.
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         has_separate_cfg=True,
         **kwargs,
     )
@@ -267,6 +305,7 @@ def sd3_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_1,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -281,6 +320,7 @@ def consisid_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_0,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -288,7 +328,7 @@ def consisid_adapter(pipe, **kwargs) -> BlockAdapter:
 @BlockAdapterRegistry.register("DiT")
 def dit_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import DiTTransformer2DModel
-    from cache_dit.cache_factory.patch_functors import DiTPatchFunctor
+    from cache_dit.caching.patch_functors import DiTPatchFunctor
 
     assert isinstance(pipe.transformer, DiTTransformer2DModel)
     return BlockAdapter(
@@ -297,6 +337,7 @@ def dit_adapter(pipe, **kwargs) -> BlockAdapter:
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_3,
         patch_functor=DiTPatchFunctor(),
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -311,6 +352,7 @@ def amused_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_layers,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -331,6 +373,7 @@ def bria_adapter(pipe, **kwargs) -> BlockAdapter:
             ForwardPattern.Pattern_0,
             ForwardPattern.Pattern_0,
         ],
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -348,6 +391,7 @@ def lumina2_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.layers,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -362,6 +406,7 @@ def omnigen_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.layers,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -376,6 +421,7 @@ def pixart_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -390,6 +436,7 @@ def sana_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -404,6 +451,7 @@ def stabledudio_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -426,6 +474,7 @@ def visualcloze_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_1,
                 ForwardPattern.Pattern_1,
             ],
+            check_forward_pattern=True,
             **kwargs,
         )
     else:
@@ -440,6 +489,7 @@ def visualcloze_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_1,
                 ForwardPattern.Pattern_3,
             ],
+            check_forward_pattern=True,
             **kwargs,
         )
 
@@ -454,6 +504,7 @@ def auraflow_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.single_transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -461,7 +512,7 @@ def auraflow_adapter(pipe, **kwargs) -> BlockAdapter:
 @BlockAdapterRegistry.register("Chroma")
 def chroma_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import ChromaTransformer2DModel
-    from cache_dit.cache_factory.patch_functors import ChromaPatchFunctor
+    from cache_dit.caching.patch_functors import ChromaPatchFunctor
 
     assert isinstance(pipe.transformer, ChromaTransformer2DModel)
     return BlockAdapter(
@@ -476,6 +527,7 @@ def chroma_adapter(pipe, **kwargs) -> BlockAdapter:
             ForwardPattern.Pattern_3,
         ],
         patch_functor=ChromaPatchFunctor(),
+        check_forward_pattern=True,
         has_separate_cfg=True,
         **kwargs,
     )
@@ -491,6 +543,7 @@ def shape_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.prior,
         blocks=pipe.prior.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -503,7 +556,7 @@ def hidream_adapter(pipe, **kwargs) -> BlockAdapter:
     # 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
-    from cache_dit.cache_factory.patch_functors import HiDreamPatchFunctor
+    from cache_dit.caching.patch_functors import HiDreamPatchFunctor
 
     assert isinstance(pipe.transformer, HiDreamImageTransformer2DModel)
     return BlockAdapter(
@@ -528,7 +581,7 @@ def hidream_adapter(pipe, **kwargs) -> BlockAdapter:
 @BlockAdapterRegistry.register("HunyuanDiT")
 def hunyuandit_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import HunyuanDiT2DModel, HunyuanDiT2DControlNetModel
-    from cache_dit.cache_factory.patch_functors import HunyuanDiTPatchFunctor
+    from cache_dit.caching.patch_functors import HunyuanDiTPatchFunctor
 
     assert isinstance(
         pipe.transformer,
@@ -540,6 +593,7 @@ def hunyuandit_adapter(pipe, **kwargs) -> BlockAdapter:
         blocks=pipe.transformer.blocks,
         forward_pattern=ForwardPattern.Pattern_3,
         patch_functor=HunyuanDiTPatchFunctor(),
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -547,7 +601,7 @@ def hunyuandit_adapter(pipe, **kwargs) -> BlockAdapter:
 @BlockAdapterRegistry.register("HunyuanDiTPAG")
 def hunyuanditpag_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import HunyuanDiT2DModel
-    from cache_dit.cache_factory.patch_functors import HunyuanDiTPatchFunctor
+    from cache_dit.caching.patch_functors import HunyuanDiTPatchFunctor
 
     assert isinstance(pipe.transformer, HunyuanDiT2DModel)
     return BlockAdapter(
@@ -556,5 +610,82 @@ def hunyuanditpag_adapter(pipe, **kwargs) -> BlockAdapter:
         blocks=pipe.transformer.blocks,
         forward_pattern=ForwardPattern.Pattern_3,
         patch_functor=HunyuanDiTPatchFunctor(),
+        check_forward_pattern=True,
         **kwargs,
     )
+
+
+@BlockAdapterRegistry.register("Kandinsky5")
+def kandinsky5_adapter(pipe, **kwargs) -> BlockAdapter:
+    try:
+        from diffusers import Kandinsky5Transformer3DModel
+
+        assert isinstance(pipe.transformer, Kandinsky5Transformer3DModel)
+        return BlockAdapter(
+            pipe=pipe,
+            transformer=pipe.transformer,
+            blocks=pipe.transformer.visual_transformer_blocks,
+            forward_pattern=ForwardPattern.Pattern_3,  # or Pattern_2
+            has_separate_cfg=True,
+            check_forward_pattern=False,
+            check_num_outputs=False,
+            **kwargs,
+        )
+    except ImportError:
+        raise ImportError(
+            "Kandinsky5Transformer3DModel is not available in the current diffusers version. "
+            "Please upgrade diffusers>=0.36.dev0 to use this adapter."
+        )
+
+
+@BlockAdapterRegistry.register("PRX")
+def prx_adapter(pipe, **kwargs) -> BlockAdapter:
+    try:
+        from diffusers import PRXTransformer2DModel
+
+        assert isinstance(pipe.transformer, PRXTransformer2DModel)
+        return BlockAdapter(
+            pipe=pipe,
+            transformer=pipe.transformer,
+            blocks=pipe.transformer.blocks,
+            forward_pattern=ForwardPattern.Pattern_3,
+            check_forward_pattern=True,
+            check_num_outputs=False,
+            **kwargs,
+        )
+    except ImportError:
+        raise ImportError(
+            "PRXTransformer2DModel is not available in the current diffusers version. "
+            "Please upgrade diffusers>=0.36.dev0 to use this adapter."
+        )
+
+
+@BlockAdapterRegistry.register("HunyuanImage")
+def hunyuan_image_adapter(pipe, **kwargs) -> BlockAdapter:
+    try:
+        from diffusers import HunyuanImageTransformer2DModel
+
+        assert isinstance(pipe.transformer, HunyuanImageTransformer2DModel)
+        return BlockAdapter(
+            pipe=pipe,
+            transformer=pipe.transformer,
+            blocks=[
+                pipe.transformer.transformer_blocks,
+                pipe.transformer.single_transformer_blocks,
+            ],
+            forward_pattern=[
+                ForwardPattern.Pattern_0,
+                ForwardPattern.Pattern_0,
+            ],
+            # set `has_separate_cfg` as True to enable separate cfg caching
+            # since in hyimage-2.1 the `guider_state` contains 2 input batches.
+            # The cfg is `enabled` by default in AdaptiveProjectedMixGuidance.
+            has_separate_cfg=True,
+            check_forward_pattern=True,
+            **kwargs,
+        )
+    except ImportError:
+        raise ImportError(
+            "HunyuanImageTransformer2DModel is not available in the current diffusers version. "
+            "Please upgrade diffusers>=0.36.dev0 to use this adapter."
+        )