cache-dit 1.0.8__py3-none-any.whl → 1.0.10__py3-none-any.whl

cache_dit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '1.0.8'
-__version_tuple__ = version_tuple = (1, 0, 8)
+__version__ = version = '1.0.10'
+__version_tuple__ = version_tuple = (1, 0, 10)
 
 __commit_id__ = commit_id = None

cache_dit/cache_factory/__init__.py

@@ -8,6 +8,7 @@ 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 BlockAdapterRegistry
+from cache_dit.cache_factory.block_adapters import FakeDiffusionPipeline
 
 from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
 from cache_dit.cache_factory.cache_contexts import DBCacheConfig

cache_dit/cache_factory/block_adapters/__init__.py

@@ -1,5 +1,8 @@
 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 (
+    FakeDiffusionPipeline,
+)
 from cache_dit.cache_factory.block_adapters.block_adapters import ParamsModifier
 from cache_dit.cache_factory.block_adapters.block_registers import (
     BlockAdapterRegistry,
@@ -27,6 +30,7 @@ def flux_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_1,
                 ForwardPattern.Pattern_1,
             ],
+            check_forward_pattern=True,
             **kwargs,
         )
     else:
@@ -41,6 +45,7 @@ def flux_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_1,
                 ForwardPattern.Pattern_3,
             ],
+            check_forward_pattern=True,
             **kwargs,
         )
 
@@ -55,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,
     )
 
@@ -69,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,
     )
 
@@ -104,6 +111,7 @@ def wan_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_2,
                 ForwardPattern.Pattern_2,
             ],
+            check_forward_pattern=True,
             has_separate_cfg=True,
             **kwargs,
         )
@@ -114,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,
         )
@@ -135,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,
@@ -159,6 +169,7 @@ def qwenimage_adapter(pipe, **kwargs) -> BlockAdapter:
             blocks=pipe.transformer.transformer_blocks,
             forward_pattern=ForwardPattern.Pattern_1,
             patch_functor=QwenImageControlNetPatchFunctor(),
+            check_forward_pattern=True,
             has_separate_cfg=True,
         )
     else:
@@ -167,6 +178,7 @@ def qwenimage_adapter(pipe, **kwargs) -> BlockAdapter:
             transformer=pipe.transformer,
             blocks=pipe.transformer.transformer_blocks,
             forward_pattern=ForwardPattern.Pattern_1,
+            check_forward_pattern=True,
             has_separate_cfg=True,
             **kwargs,
         )
@@ -182,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,
     )
 
@@ -196,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,
     )
 
@@ -210,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,
     )
 
@@ -224,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,
     )
@@ -239,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,
     )
@@ -254,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,
     )
 
@@ -271,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,
     )
@@ -286,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,
     )
 
@@ -300,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,
     )
 
@@ -316,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,
     )
 
@@ -330,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,
     )
 
@@ -350,6 +373,7 @@ def bria_adapter(pipe, **kwargs) -> BlockAdapter:
             ForwardPattern.Pattern_0,
             ForwardPattern.Pattern_0,
         ],
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -367,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,
     )
 
@@ -381,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,
     )
 
@@ -395,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,
     )
 
@@ -409,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,
     )
 
@@ -423,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,
     )
 
@@ -445,6 +474,7 @@ def visualcloze_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_1,
                 ForwardPattern.Pattern_1,
             ],
+            check_forward_pattern=True,
             **kwargs,
         )
     else:
@@ -459,6 +489,7 @@ def visualcloze_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_1,
                 ForwardPattern.Pattern_3,
             ],
+            check_forward_pattern=True,
             **kwargs,
         )
 
@@ -473,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,
     )
 
@@ -495,6 +527,7 @@ def chroma_adapter(pipe, **kwargs) -> BlockAdapter:
             ForwardPattern.Pattern_3,
         ],
         patch_functor=ChromaPatchFunctor(),
+        check_forward_pattern=True,
         has_separate_cfg=True,
         **kwargs,
     )
@@ -510,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,
     )
 
@@ -559,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,
     )
 
@@ -575,6 +610,7 @@ def hunyuanditpag_adapter(pipe, **kwargs) -> BlockAdapter:
         blocks=pipe.transformer.blocks,
         forward_pattern=ForwardPattern.Pattern_3,
         patch_functor=HunyuanDiTPatchFunctor(),
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -613,6 +649,7 @@ def prx_adapter(pipe, **kwargs) -> BlockAdapter:
             transformer=pipe.transformer,
             blocks=pipe.transformer.blocks,
             forward_pattern=ForwardPattern.Pattern_3,
+            check_forward_pattern=True,
             check_num_outputs=False,
             **kwargs,
         )

cache_dit/cache_factory/block_adapters/block_adapters.py

@@ -6,7 +6,7 @@ from collections.abc import Iterable
 
 from typing import Any, Tuple, List, Optional, Union
 
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, ModelMixin
 from cache_dit.cache_factory.patch_functors import PatchFunctor
 from cache_dit.cache_factory.forward_pattern import ForwardPattern
 from cache_dit.cache_factory.params_modifier import ParamsModifier
@@ -16,12 +16,22 @@ from cache_dit.logger import init_logger
 logger = init_logger(__name__)
 
 
+class FakeDiffusionPipeline:
+    # A placeholder for pipelines when pipe is None.
+    def __init__(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+    ):
+        self.transformer = transformer  # Reference only
+
+
 @dataclasses.dataclass
 class BlockAdapter:
 
     # Transformer configurations.
     pipe: Union[
         DiffusionPipeline,
+        FakeDiffusionPipeline,
         Any,
     ] = None
 
@@ -73,7 +83,7 @@ class BlockAdapter:
         ]
     ] = None
 
-    check_forward_pattern: bool = True
+    check_forward_pattern: Optional[bool] = None
     check_num_outputs: bool = False
 
     # Pipeline Level Flags
@@ -110,12 +120,43 @@ class BlockAdapter:
     def __post_init__(self):
         if self.skip_post_init:
             return
+
+        self.maybe_fake_pipe()
         if any((self.pipe is not None, self.transformer is not None)):
             self.maybe_fill_attrs()
             self.maybe_patchify()
             self.maybe_skip_checks()
 
+    def maybe_fake_pipe(self):
+        if self.pipe is None:
+            self.pipe = FakeDiffusionPipeline()
+            logger.warning("pipe is None, use FakeDiffusionPipeline instead.")
+
     def maybe_skip_checks(self):
+        if self.check_forward_pattern is None:
+            if self.transformer is not None:
+                if self.nested_depth(self.transformer) == 0:
+                    transformer = self.transformer
+                elif self.nested_depth(self.transformer) == 1:
+                    transformer = self.transformer[0]
+                else:
+                    raise ValueError(
+                        "transformer nested depth can't more than 1, "
+                        f"current is: {self.nested_depth(self.transformer)}"
+                    )
+                if transformer.__module__.startswith("diffusers"):
+                    self.check_forward_pattern = True
+                    logger.info(
+                        f"Found transformer from diffusers: {transformer.__module__} "
+                        "enable check_forward_pattern by default."
+                    )
+                else:
+                    self.check_forward_pattern = False
+                    logger.info(
+                        f"Found transformer NOT from diffusers: {transformer.__module__} "
+                        "disable check_forward_pattern by default."
+                    )
+
         if getattr(self.transformer, "_hf_hook", None) is not None:
             logger.warning("_hf_hook is not None, force skip pattern check!")
             self.check_forward_pattern = False
@@ -208,7 +249,10 @@ class BlockAdapter:
             if self.transformer is not None:
                 self.patch_functor.apply(self.transformer, *args, **kwargs)
             else:
-                assert hasattr(self.pipe, "transformer")
+                assert hasattr(self.pipe, "transformer"), (
+                    "pipe.transformer can not be None when patch_functor "
+                    "is provided and transformer is None."
+                )
                 self.patch_functor.apply(self.pipe.transformer, *args, **kwargs)
 
     @staticmethod
@@ -224,6 +268,10 @@ class BlockAdapter:
             adapter.forward_pattern is not None
         ), "adapter.forward_pattern can not be None."
         pipe = adapter.pipe
+        if isinstance(pipe, FakeDiffusionPipeline):
+            raise ValueError(
+                "Can not auto block adapter for FakeDiffusionPipeline."
+            )
 
         assert hasattr(pipe, "transformer"), "pipe.transformer can not be None."
 
@@ -489,6 +537,7 @@ class BlockAdapter:
     @staticmethod
     def normalize(
         adapter: "BlockAdapter",
+        unique: bool = True,
     ) -> "BlockAdapter":
 
         if getattr(adapter, "_is_normalized", False):
@@ -523,7 +572,10 @@ class BlockAdapter:
         adapter.forward_pattern = _normalize_attr(adapter.forward_pattern)
         adapter.dummy_blocks_names = _normalize_attr(adapter.dummy_blocks_names)
         adapter.params_modifiers = _normalize_attr(adapter.params_modifiers)
-        BlockAdapter.unique(adapter)
+        # Some times, the cache_config will be None.
+        # So we do not perform unique check here.
+        if unique:
+            BlockAdapter.unique(adapter)
 
         adapter._is_normalized = True
 
@@ -571,6 +623,10 @@ class BlockAdapter:
         if not getattr(adapter, "_is_normalized", False):
             raise RuntimeError("block_adapter must be normailzed.")
 
+    @classmethod
+    def is_normalized(cls, adapter: "BlockAdapter") -> bool:
+        return getattr(adapter, "_is_normalized", False)
+
     @classmethod
     def is_cached(cls, adapter: Any) -> bool:
         if isinstance(adapter, cls):
@@ -592,6 +648,21 @@ class BlockAdapter:
         else:
             return getattr(adapter, "_is_cached", False)
 
+    @classmethod
+    def is_parallelized(cls, adapter: Any) -> bool:
+        if isinstance(adapter, cls):
+            cls.assert_normalized(adapter)
+            return getattr(adapter.transformer[0], "_is_parallelized", False)
+        elif isinstance(adapter, DiffusionPipeline):
+            return getattr(adapter.transformer, "_is_parallelized", False)
+        elif isinstance(adapter, torch.nn.Module):
+            return getattr(adapter, "_is_parallelized", False)
+        elif isinstance(adapter, list):  # [TRN_0,...]
+            assert isinstance(adapter[0], torch.nn.Module)
+            return getattr(adapter[0], "_is_parallelized", False)
+        else:
+            return getattr(adapter, "_is_parallelized", False)
+
     @classmethod
     def nested_depth(cls, obj: Any):
         # str: 0; List[str]: 1; List[List[str]]: 2

cache_dit/cache_factory/block_adapters/block_registers.py

@@ -1,7 +1,11 @@
-from typing import Any, Tuple, List, Dict, Callable
+import torch
+from typing import Any, Tuple, List, Dict, Callable, Union
 
 from diffusers import DiffusionPipeline
-from cache_dit.cache_factory.block_adapters.block_adapters import BlockAdapter
+from cache_dit.cache_factory.block_adapters.block_adapters import (
+    BlockAdapter,
+    FakeDiffusionPipeline,
+)
 
 from cache_dit.logger import init_logger
 
@@ -35,24 +39,42 @@ class BlockAdapterRegistry:
     @classmethod
     def get_adapter(
         cls,
-        pipe: DiffusionPipeline | str | Any,
+        pipe_or_module: DiffusionPipeline | torch.nn.Module | str | Any,
         **kwargs,
-    ) -> BlockAdapter:
-        if not isinstance(pipe, str):
-            pipe_cls_name: str = pipe.__class__.__name__
+    ) -> BlockAdapter | None:
+        if not isinstance(pipe_or_module, str):
+            cls_name: str = pipe_or_module.__class__.__name__
         else:
-            pipe_cls_name = pipe
+            cls_name = pipe_or_module
 
         for name in cls._adapters:
-            if pipe_cls_name.startswith(name):
-                return cls._adapters[name](pipe, **kwargs)
-
-        return BlockAdapter()
+            if cls_name.startswith(name):
+                if not isinstance(pipe_or_module, DiffusionPipeline):
+                    assert isinstance(pipe_or_module, torch.nn.Module)
+                    # NOTE: Make pre-registered adapters support Transformer-only case.
+                    # WARN: This branch is not officially supported and only for testing
+                    # purpose. We construct a fake diffusion pipeline that contains the
+                    # given transformer module. Currently, only works for DiT models which
+                    # only have one transformer module. Case like multiple transformers
+                    # is not supported, e.g, Wan2.2. Please use BlockAdapter directly for
+                    # such cases.
+                    return cls._adapters[name](
+                        FakeDiffusionPipeline(pipe_or_module), **kwargs
+                    )
+                else:
+                    return cls._adapters[name](pipe_or_module, **kwargs)
+
+        return None
 
     @classmethod
     def has_separate_cfg(
         cls,
-        pipe_or_adapter: DiffusionPipeline | BlockAdapter | Any,
+        pipe_or_adapter: Union[
+            DiffusionPipeline,
+            FakeDiffusionPipeline,
+            BlockAdapter,
+            Any,
+        ],
     ) -> bool:
 
         # Prefer custom setting from block adapter.
@@ -60,11 +82,16 @@ class BlockAdapterRegistry:
             return pipe_or_adapter.has_separate_cfg
 
         has_separate_cfg = False
+        if isinstance(pipe_or_adapter, FakeDiffusionPipeline):
+            return False
+
         if isinstance(pipe_or_adapter, DiffusionPipeline):
-            has_separate_cfg = cls.get_adapter(
+            adapter = cls.get_adapter(
                 pipe_or_adapter,
                 skip_post_init=True,  # check cfg setting only
-            ).has_separate_cfg
+            )
+            if adapter is not None:
+                has_separate_cfg = adapter.has_separate_cfg
 
         if has_separate_cfg:
             return True
@@ -77,11 +104,11 @@ class BlockAdapterRegistry:
         return False
 
     @classmethod
-    def is_supported(cls, pipe) -> bool:
-        pipe_cls_name: str = pipe.__class__.__name__
+    def is_supported(cls, pipe_or_module) -> bool:
+        cls_name: str = pipe_or_module.__class__.__name__
 
         for name in cls._adapters:
-            if pipe_cls_name.startswith(name):
+            if cls_name.startswith(name):
                 return True
         return False