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

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

@@ -6,22 +6,32 @@ from collections.abc import Iterable
 
 from typing import Any, Tuple, List, Optional, Union
 
-from diffusers import DiffusionPipeline
-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
+from diffusers import DiffusionPipeline, ModelMixin
+from cache_dit.caching.patch_functors import PatchFunctor
+from cache_dit.caching.forward_pattern import ForwardPattern
+from cache_dit.caching.params_modifier import ParamsModifier
 
 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,9 +120,53 @@ 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
+            self.check_num_outputs = False
+        elif getattr(self.transformer, "_diffusers_hook", None) is not None:
+            logger.warning(
+                "_diffusers_hook is not None, force skip pattern check!"
+            )
+            self.check_forward_pattern = False
+            self.check_num_outputs = False
 
     def maybe_fill_attrs(self):
         # NOTE: This func should be call before normalize.
@@ -195,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
@@ -211,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."
 
@@ -476,6 +537,7 @@ class BlockAdapter:
     @staticmethod
     def normalize(
         adapter: "BlockAdapter",
+        unique: bool = True,
     ) -> "BlockAdapter":
 
         if getattr(adapter, "_is_normalized", False):
@@ -510,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
 
@@ -558,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):
@@ -579,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/caching/block_adapters/block_registers.py

@@ -0,0 +1,118 @@
+import torch
+from typing import Any, Tuple, List, Dict, Callable, Union
+
+from diffusers import DiffusionPipeline
+from cache_dit.caching.block_adapters.block_adapters import (
+    BlockAdapter,
+    FakeDiffusionPipeline,
+)
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class BlockAdapterRegistry:
+    _adapters: Dict[str, Callable[..., BlockAdapter]] = {}
+    _predefined_adapters_has_separate_cfg: List[str] = [
+        "QwenImage",
+        "Wan",
+        "CogView4",
+        "Cosmos",
+        "SkyReelsV2",
+        "Chroma",
+        "Lumina2",
+        "Kandinsky5",
+    ]
+
+    @classmethod
+    def register(cls, name: str, supported: bool = True):
+        def decorator(
+            func: Callable[..., BlockAdapter]
+        ) -> Callable[..., BlockAdapter]:
+            if supported:
+                cls._adapters[name] = func
+            return func
+
+        return decorator
+
+    @classmethod
+    def get_adapter(
+        cls,
+        pipe_or_module: DiffusionPipeline | torch.nn.Module | str | Any,
+        **kwargs,
+    ) -> BlockAdapter | None:
+        if not isinstance(pipe_or_module, str):
+            cls_name: str = pipe_or_module.__class__.__name__
+        else:
+            cls_name = pipe_or_module
+
+        for name in cls._adapters:
+            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: Union[
+            DiffusionPipeline,
+            FakeDiffusionPipeline,
+            BlockAdapter,
+            Any,
+        ],
+    ) -> bool:
+
+        # Prefer custom setting from block adapter.
+        if isinstance(pipe_or_adapter, BlockAdapter):
+            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):
+            adapter = cls.get_adapter(
+                pipe_or_adapter,
+                skip_post_init=True,  # check cfg setting only
+            )
+            if adapter is not None:
+                has_separate_cfg = adapter.has_separate_cfg
+
+        if has_separate_cfg:
+            return True
+
+        pipe_cls_name = pipe_or_adapter.__class__.__name__
+        for name in cls._predefined_adapters_has_separate_cfg:
+            if pipe_cls_name.startswith(name):
+                return True
+
+        return False
+
+    @classmethod
+    def is_supported(cls, pipe_or_module) -> bool:
+        cls_name: str = pipe_or_module.__class__.__name__
+
+        for name in cls._adapters:
+            if cls_name.startswith(name):
+                return True
+        return False
+
+    @classmethod
+    def supported_pipelines(cls, **kwargs) -> Tuple[int, List[str]]:
+        val_pipelines = cls._adapters.keys()
+        return len(val_pipelines), [p for p in val_pipelines]

cache_dit/caching/cache_adapters/__init__.py

@@ -0,0 +1 @@
+from cache_dit.caching.cache_adapters.cache_adapter import CachedAdapter