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

cache_dit/cache_factory/block_adapters/block_adapters.py

@@ -2,8 +2,9 @@ import torch
 
 import inspect
 import dataclasses
+from collections.abc import Iterable
 
-from typing import Any, Tuple, List, Optional
+from typing import Any, Tuple, List, Optional, Union
 
 from diffusers import DiffusionPipeline
 from cache_dit.cache_factory.forward_pattern import ForwardPattern
@@ -14,21 +15,76 @@ from cache_dit.logger import init_logger
 logger = init_logger(__name__)
 
 
+class ParamsModifier:
+    def __init__(self, **kwargs):
+        self._context_kwargs = kwargs.copy()
+
+
 @dataclasses.dataclass
 class BlockAdapter:
+
     # Transformer configurations.
-    pipe: DiffusionPipeline | Any = None
-    transformer: torch.nn.Module = None
+    pipe: Union[
+        DiffusionPipeline,
+        Any,
+    ] = None
+
+    # single transformer (most cases) or list of transformers (Wan2.2, etc)
+    transformer: Union[
+        torch.nn.Module,
+        List[torch.nn.Module],
+    ] = None
+
+    # Block Level Flags
+    # Each transformer contains a list of blocks-list,
+    # blocks_name-list, dummy_blocks_names-list, etc.
+    blocks: Union[
+        torch.nn.ModuleList,
+        List[torch.nn.ModuleList],
+        List[List[torch.nn.ModuleList]],
+    ] = None
 
-    # ------------ Block Level Flags ------------
-    blocks: torch.nn.ModuleList | List[torch.nn.ModuleList] = None
     # transformer_blocks, blocks, etc.
-    blocks_name: str | List[str] = None
-    dummy_blocks_names: List[str] = dataclasses.field(default_factory=list)
-    forward_pattern: ForwardPattern | List[ForwardPattern] = None
+    blocks_name: Union[
+        str,
+        List[str],
+        List[List[str]],
+    ] = None
+
+    unique_blocks_name: Union[
+        str,
+        List[str],
+        List[List[str]],
+    ] = dataclasses.field(default_factory=list)
+
+    dummy_blocks_names: Union[
+        List[str],
+        List[List[str]],
+    ] = dataclasses.field(default_factory=list)
+
+    forward_pattern: Union[
+        ForwardPattern,
+        List[ForwardPattern],
+        List[List[ForwardPattern]],
+    ] = None
+
+    # modify cache context params for specific blocks.
+    params_modifiers: Union[
+        ParamsModifier,
+        List[ParamsModifier],
+        List[List[ParamsModifier]],
+    ] = None
+
     check_num_outputs: bool = True
 
+    # Pipeline Level Flags
+    # Patch Functor: Flux, etc.
+    patch_functor: Optional[PatchFunctor] = None
+    # Flags for separate cfg
+    has_separate_cfg: bool = False
+
     # Flags to control auto block adapter
+    # NOTE: NOT support for multi-transformers.
     auto: bool = False
     allow_prefixes: List[str] = dataclasses.field(
         default_factory=lambda: [
@@ -49,24 +105,92 @@ class BlockAdapter:
         default="max", metadata={"allowed_values": ["max", "min"]}
     )
 
-    # NOTE: Other flags.
-    disable_patch: bool = False
-
-    # ------------ Pipeline Level Flags ------------
-    # Patch Functor: Flux, etc.
-    patch_functor: Optional[PatchFunctor] = None
-    # Flags for separate cfg
-    has_separate_cfg: bool = False
+    # Other Flags
+    skip_post_init: bool = False
 
     def __post_init__(self):
+        if self.skip_post_init:
+            return
         assert any((self.pipe is not None, self.transformer is not None))
-        self.patchify()
+        self.maybe_fill_attrs()
+        self.maybe_patchify()
+
+    def maybe_fill_attrs(self):
+        # NOTE: This func should be call before normalize.
+        # Allow empty `blocks_names`, we will auto fill it.
+        # TODO: preprocess more empty attrs.
+        if (
+            self.transformer is not None
+            and self.blocks is not None
+            and self.blocks_name is None
+        ):
+
+            def _find(transformer, blocks):
+                attr_names = dir(transformer)
+                assert isinstance(blocks, torch.nn.ModuleList)
+                blocks_name = None
+                for attr_name in attr_names:
+                    if attr := getattr(transformer, attr_name, None):
+                        if isinstance(attr, torch.nn.ModuleList) and id(
+                            attr
+                        ) == id(blocks):
+                            blocks_name = attr_name
+                            break
+                assert (
+                    blocks_name is not None
+                ), "No blocks_name match, please set it manually!"
+                return blocks_name
+
+            if self.nested_depth(self.transformer) == 0:
+                if self.nested_depth(self.blocks) == 0:  # str
+                    self.blocks_name = _find(self.transformer, self.blocks)
+                elif self.nested_depth(self.blocks) == 1:
+                    self.blocks_name = [
+                        _find(self.transformer, blocks)
+                        for blocks in self.blocks
+                    ]
+                else:
+                    raise ValueError(
+                        "Blocks nested depth can't more than 1 if transformer "
+                        f"is not a list, current is: {self.nested_depth(self.blocks)}"
+                    )
+            elif self.nested_depth(self.transformer) == 1:  # List[str]
+                if self.nested_depth(self.blocks) == 1:  # List[str]
+                    assert len(self.transformer) == len(self.blocks)
+                    self.blocks_name = [
+                        _find(transformer, blocks)
+                        for transformer, blocks in zip(
+                            self.transformer, self.blocks
+                        )
+                    ]
+                elif self.nested_depth(self.blocks) == 2:  # List[List[str]]
+                    assert len(self.transformer) == len(self.blocks)
+                    self.blocks_name = []
+                    for i in range(len(self.blocks)):
+                        self.blocks_name.append(
+                            [
+                                _find(self.transformer[i], blocks)
+                                for blocks in self.blocks[i]
+                            ]
+                        )
+                else:
+                    raise ValueError(
+                        "Blocks nested depth can only be 1 or 2 "
+                        "if transformer is a list, current is: "
+                        f"{self.nested_depth(self.blocks)}"
+                    )
+            else:
+                raise ValueError(
+                    "transformer nested depth can't more than 1, "
+                    f"current is: {self.nested_depth(self.transformer)}"
+                )
+            logger.info(f"Auto fill blocks_name: {self.blocks_name}.")
 
-    def patchify(self, *args, **kwargs):
+    def maybe_patchify(self, *args, **kwargs):
         # Process some specificial cases, specific for transformers
         # that has different forward patterns between single_transformer_blocks
         # and transformer_blocks , such as Flux (diffusers < 0.35.0).
-        if self.patch_functor is not None and not self.disable_patch:
+        if self.patch_functor is not None:
             if self.transformer is not None:
                 self.patch_functor.apply(self.transformer, *args, **kwargs)
             else:
@@ -92,7 +216,7 @@ class BlockAdapter:
         transformer = pipe.transformer
 
         # "transformer_blocks", "blocks", "single_transformer_blocks", "layers"
-        blocks, blocks_name = BlockAdapter.find_blocks(
+        blocks, blocks_name = BlockAdapter.find_match_blocks(
             transformer=transformer,
             allow_prefixes=adapter.allow_prefixes,
             allow_suffixes=adapter.allow_suffixes,
@@ -115,6 +239,10 @@ class BlockAdapter:
     def check_block_adapter(
         adapter: "BlockAdapter",
     ) -> bool:
+
+        if getattr(adapter, "_is_normlized", False):
+            return True
+
         def _check_warning(attr: str):
             if getattr(adapter, attr, None) is None:
                 logger.warning(f"{attr} is None!")
@@ -136,24 +264,23 @@ class BlockAdapter:
         if not _check_warning("forward_pattern"):
             return False
 
-        if isinstance(adapter.blocks, list):
-            for i, blocks in enumerate(adapter.blocks):
-                if not isinstance(blocks, torch.nn.ModuleList):
-                    logger.warning(f"blocks[{i}] is not ModuleList.")
-                    return False
+        if BlockAdapter.nested_depth(adapter.blocks) == 0:
+            blocks = adapter.blocks
         else:
-            if not isinstance(adapter.blocks, torch.nn.ModuleList):
-                logger.warning("blocks is not ModuleList.")
-                return False
+            blocks = BlockAdapter.flatten(adapter.blocks)[0]
+
+        if not isinstance(blocks, torch.nn.ModuleList):
+            logger.warning("blocks is not ModuleList.")
+            return False
 
         return True
 
     @staticmethod
-    def find_blocks(
+    def find_match_blocks(
         transformer: torch.nn.Module,
         allow_prefixes: List[str] = [
-            "transformer",
-            "single_transformer",
+            "transformer_blocks",
+            "single_transformer_blocks",
             "blocks",
             "layers",
             "single_stream_blocks",
@@ -181,10 +308,10 @@ class BlockAdapter:
         valid_count = []
         forward_pattern = kwargs.pop("forward_pattern", None)
         for blocks_name in blocks_names:
-            if blocks := getattr(transformer, blocks_name, None):
+            if (blocks := getattr(transformer, blocks_name, None)) is not None:
                 if isinstance(blocks, torch.nn.ModuleList):
                     block = blocks[0]
-                    block_cls_name = block.__class__.__name__
+                    block_cls_name: str = block.__class__.__name__
                     # Check suffixes
                     if isinstance(block, torch.nn.Module) and (
                         any(
@@ -244,6 +371,18 @@ class BlockAdapter:
 
         return final_blocks, final_name
 
+    @staticmethod
+    def find_blocks(
+        transformer: torch.nn.Module,
+    ) -> List[torch.nn.ModuleList]:
+        total_blocks = []
+        for attr in dir(transformer):
+            if (blocks := getattr(transformer, attr, None)) is not None:
+                if isinstance(blocks, torch.nn.ModuleList):
+                    if isinstance(blocks[0], torch.nn.Module):
+                        total_blocks.append(blocks)
+        return total_blocks
+
     @staticmethod
     def match_block_pattern(
         block: torch.nn.Module,
@@ -320,14 +459,148 @@ class BlockAdapter:
     def normalize(
         adapter: "BlockAdapter",
     ) -> "BlockAdapter":
-        if not isinstance(adapter.blocks, list):
-            adapter.blocks = [adapter.blocks]
-        if not isinstance(adapter.blocks_name, list):
-            adapter.blocks_name = [adapter.blocks_name]
-        if not isinstance(adapter.forward_pattern, list):
-            adapter.forward_pattern = [adapter.forward_pattern]
 
-        assert len(adapter.blocks) == len(adapter.blocks_name)
-        assert len(adapter.blocks) == len(adapter.forward_pattern)
+        if getattr(adapter, "_is_normalized", False):
+            return adapter
+
+        if BlockAdapter.nested_depth(adapter.transformer) == 0:
+            adapter.transformer = [adapter.transformer]
+
+        def _normalize_attr(attr: Any):
+            normalized_attr = attr
+            if attr is None:
+                return normalized_attr
+
+            if BlockAdapter.nested_depth(attr) == 0:
+                normalized_attr = [[attr]]
+            elif BlockAdapter.nested_depth(attr) == 1:  # List
+                if attr:  # not-empty
+                    if len(attr) == len(adapter.transformer):
+                        normalized_attr = [[a] for a in attr]
+                    else:
+                        normalized_attr = [attr]
+                else:  # [] empty
+                    normalized_attr = [
+                        [] for _ in range(len(adapter.transformer))
+                    ]
+
+            assert len(adapter.transformer) == len(normalized_attr)
+            return normalized_attr
+
+        adapter.blocks = _normalize_attr(adapter.blocks)
+        adapter.blocks_name = _normalize_attr(adapter.blocks_name)
+        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)
+
+        adapter._is_normalized = True
 
         return adapter
+
+    @classmethod
+    def unique(cls, adapter: "BlockAdapter"):
+        # NOTE: Users should never call this function
+        for i in range(len(adapter.blocks)):
+            assert len(adapter.blocks[i]) == len(adapter.blocks_name[i])
+            assert len(adapter.blocks[i]) == len(adapter.forward_pattern[i])
+
+        # Generate unique blocks names
+        if len(adapter.unique_blocks_name) == 0:
+            for i in range(len(adapter.transformer)):
+                adapter.unique_blocks_name.append(
+                    [
+                        f"{name}_{hash(id(blocks))}"
+                        for blocks, name in zip(
+                            adapter.blocks[i],
+                            adapter.blocks_name[i],
+                        )
+                    ]
+                )
+        else:
+            assert len(adapter.transformer) == len(adapter.unique_blocks_name)
+
+        # Also check Match Forward Pattern
+        for i in range(len(adapter.transformer)):
+            for forward_pattern, blocks in zip(
+                adapter.forward_pattern[i], adapter.blocks[i]
+            ):
+                assert BlockAdapter.match_blocks_pattern(
+                    blocks,
+                    forward_pattern=forward_pattern,
+                    check_num_outputs=adapter.check_num_outputs,
+                ), (
+                    "No block forward pattern matched, "
+                    f"supported lists: {ForwardPattern.supported_patterns()}"
+                )
+
+    @classmethod
+    def assert_normalized(cls, adapter: "BlockAdapter"):
+        if not getattr(adapter, "_is_normalized", False):
+            raise RuntimeError("block_adapter must be normailzed.")
+
+    @classmethod
+    def is_cached(cls, adapter: Any) -> bool:
+        if isinstance(adapter, cls):
+            cls.assert_normalized(adapter)
+            return all(
+                (
+                    getattr(adapter.pipe, "_is_cached", False),
+                    getattr(adapter.transformer[0], "_is_cached", False),
+                )
+            )
+        elif isinstance(
+            adapter,
+            (DiffusionPipeline, torch.nn.Module),
+        ):
+            return getattr(adapter, "_is_cached", False)
+        elif isinstance(adapter, list):  # [TRN_0,...]
+            assert isinstance(adapter[0], torch.nn.Module)
+            return getattr(adapter[0], "_is_cached", False)
+        else:
+            raise TypeError(f"Can't check this type: {adapter}!")
+
+    @classmethod
+    def nested_depth(cls, obj: Any):
+        # str: 0; List[str]: 1; List[List[str]]: 2
+        atom_types = (
+            str,
+            bytes,
+            torch.nn.ModuleList,
+            torch.nn.Module,
+            torch.Tensor,
+        )
+        if isinstance(obj, atom_types):
+            return 0
+        if not isinstance(obj, Iterable):
+            return 0
+        if isinstance(obj, dict):
+            items = obj.values()
+        else:
+            items = obj
+
+        max_depth = 0
+        for item in items:
+            current_depth = cls.nested_depth(item)
+            if current_depth > max_depth:
+                max_depth = current_depth
+        return 1 + max_depth
+
+    @classmethod
+    def flatten(cls, attr: List[Any]) -> List[Any]:
+        atom_types = (
+            str,
+            bytes,
+            torch.nn.ModuleList,
+            torch.nn.Module,
+            torch.Tensor,
+        )
+        if not isinstance(attr, list):
+            return attr
+        flattened = []
+        for item in attr:
+            if isinstance(item, list) and not isinstance(item, atom_types):
+                flattened.extend(cls.flatten(item))
+            else:
+                flattened.append(item)
+        return flattened

cache_dit/cache_factory/block_adapters/block_registers.py

@@ -47,15 +47,24 @@ class BlockAdapterRegistry:
     @classmethod
     def has_separate_cfg(
         cls,
-        pipe: DiffusionPipeline | str | Any,
+        pipe_or_adapter: DiffusionPipeline | BlockAdapter | Any,
     ) -> bool:
-        if cls.get_adapter(
-            pipe,
-            disable_patch=True,
-        ).has_separate_cfg:
+
+        # 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, DiffusionPipeline):
+            has_separate_cfg = cls.get_adapter(
+                pipe_or_adapter,
+                skip_post_init=True,  # check cfg setting only
+            ).has_separate_cfg
+
+        if has_separate_cfg:
             return True
 
-        pipe_cls_name = pipe.__class__.__name__
+        pipe_cls_name = pipe_or_adapter.__class__.__name__
         for name in cls._predefined_adapters_has_spearate_cfg:
             if pipe_cls_name.startswith(name):
                 return True