cache-dit 0.2.28__py3-none-any.whl → 0.2.30__py3-none-any.whl

cache_dit/cache_factory/block_adapters/block_adapters.py

@@ -2,6 +2,7 @@ import torch
 
 import inspect
 import dataclasses
+from collections.abc import Iterable
 
 from typing import Any, Tuple, List, Optional, Union
 
@@ -74,7 +75,7 @@ class BlockAdapter:
         List[List[ParamsModifier]],
     ] = None
 
-    check_num_outputs: bool = True
+    check_num_outputs: bool = False
 
     # Pipeline Level Flags
     # Patch Functor: Flux, etc.
@@ -82,9 +83,6 @@ class BlockAdapter:
     # Flags for separate cfg
     has_separate_cfg: bool = False
 
-    # Other Flags
-    disable_patch: bool = False
-
     # Flags to control auto block adapter
     # NOTE: NOT support for multi-transformers.
     auto: bool = False
@@ -107,15 +105,94 @@ class BlockAdapter:
         default="max", metadata={"allowed_values": ["max", "min"]}
     )
 
+    # Other Flags
+    skip_post_init: bool = False
+
     def __post_init__(self):
-        assert any((self.pipe is not None, self.transformer is not None))
-        self.patchify()
+        if self.skip_post_init:
+            return
+        if any((self.pipe is not None, self.transformer is not None)):
+            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 patchify(self, *args, **kwargs):
+            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)
+                    ) is not 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 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:
@@ -141,7 +218,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,
@@ -164,6 +241,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!")
@@ -185,24 +266,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",
@@ -230,10 +310,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(
@@ -293,6 +373,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,
@@ -373,103 +465,51 @@ class BlockAdapter:
         if getattr(adapter, "_is_normalized", False):
             return adapter
 
-        if not isinstance(adapter.transformer, list):
+        if BlockAdapter.nested_depth(adapter.transformer) == 0:
             adapter.transformer = [adapter.transformer]
 
-        if isinstance(adapter.blocks, torch.nn.ModuleList):
-            # blocks_0 = [[blocks_0,],] -> match [TRN_0,]
-            adapter.blocks = [[adapter.blocks]]
-        elif isinstance(adapter.blocks, list):
-            if isinstance(adapter.blocks[0], torch.nn.ModuleList):
-                # [blocks_0, blocks_1] -> [[blocks_0, blocks_1],] -> match [TRN_0,]
-                if len(adapter.blocks) == len(adapter.transformer):
-                    adapter.blocks = [[blocks] for blocks in adapter.blocks]
-                else:
-                    adapter.blocks = [adapter.blocks]
-            elif isinstance(adapter.blocks[0], list):
-                # [[blocks_0, blocks_1],[blocks_2, blocks_3],] -> match [TRN_0, TRN_1,]
-                pass
-
-        if isinstance(adapter.blocks_name, str):
-            adapter.blocks_name = [[adapter.blocks_name]]
-        elif isinstance(adapter.blocks_name, list):
-            if isinstance(adapter.blocks_name[0], str):
-                if len(adapter.blocks_name) == len(adapter.transformer):
-                    adapter.blocks_name = [
-                        [blocks_name] for blocks_name in adapter.blocks_name
-                    ]
-                else:
-                    adapter.blocks_name = [adapter.blocks_name]
-            elif isinstance(adapter.blocks_name[0], list):
-                pass
-
-        if isinstance(adapter.forward_pattern, ForwardPattern):
-            adapter.forward_pattern = [[adapter.forward_pattern]]
-        elif isinstance(adapter.forward_pattern, list):
-            if isinstance(adapter.forward_pattern[0], ForwardPattern):
-                if len(adapter.forward_pattern) == len(adapter.transformer):
-                    adapter.forward_pattern = [
-                        [forward_pattern]
-                        for forward_pattern in adapter.forward_pattern
-                    ]
-                else:
-                    adapter.forward_pattern = [adapter.forward_pattern]
-            elif isinstance(adapter.forward_pattern[0], list):
-                pass
-
-        if isinstance(adapter.dummy_blocks_names, list):
-            if len(adapter.dummy_blocks_names) > 0:
-                if isinstance(adapter.dummy_blocks_names[0], str):
-                    if len(adapter.dummy_blocks_names) == len(
-                        adapter.transformer
-                    ):
-                        adapter.dummy_blocks_names = [
-                            [dummy_blocks_names]
-                            for dummy_blocks_names in adapter.dummy_blocks_names
-                        ]
-                    else:
-                        adapter.dummy_blocks_names = [
-                            adapter.dummy_blocks_names
-                        ]
-                elif isinstance(adapter.dummy_blocks_names[0], list):
-                    pass
-            else:
-                # Empty dummy_blocks_names
-                adapter.dummy_blocks_names = [
-                    [] for _ in range(len(adapter.transformer))
-                ]
-
-        if adapter.params_modifiers is not None:
-            if isinstance(adapter.params_modifiers, ParamsModifier):
-                adapter.params_modifiers = [[adapter.params_modifiers]]
-            elif isinstance(adapter.params_modifiers, list):
-                if isinstance(adapter.params_modifiers[0], ParamsModifier):
-                    if len(adapter.params_modifiers) == len(
-                        adapter.transformer
-                    ):
-                        adapter.params_modifiers = [
-                            [params_modifiers]
-                            for params_modifiers in adapter.params_modifiers
-                        ]
+        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:
-                        adapter.params_modifiers = [adapter.params_modifiers]
-                elif isinstance(adapter.params_modifiers[0], list):
-                    pass
+                        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)
 
-        assert len(adapter.transformer) == len(adapter.blocks)
-        assert len(adapter.transformer) == len(adapter.blocks_name)
-        assert len(adapter.transformer) == len(adapter.forward_pattern)
-        assert len(adapter.transformer) == len(adapter.dummy_blocks_names)
-        if adapter.params_modifiers is not None:
-            assert len(adapter.transformer) == len(adapter.params_modifiers)
+        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)):
-                # Generate unique blocks names
                 adapter.unique_blocks_name.append(
                     [
                         f"{name}_{hash(id(blocks))}"
@@ -479,10 +519,10 @@ class BlockAdapter:
                         )
                     ]
                 )
+        else:
+            assert len(adapter.transformer) == len(adapter.unique_blocks_name)
 
-        assert len(adapter.transformer) == len(adapter.unique_blocks_name)
-
-        # Match Forward Pattern
+        # 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]
@@ -496,10 +536,6 @@ class BlockAdapter:
                     f"supported lists: {ForwardPattern.supported_patterns()}"
                 )
 
-        adapter._is_normalized = True
-
-        return adapter
-
     @classmethod
     def assert_normalized(cls, adapter: "BlockAdapter"):
         if not getattr(adapter, "_is_normalized", False):
@@ -524,15 +560,49 @@ class BlockAdapter:
             assert isinstance(adapter[0], torch.nn.Module)
             return getattr(adapter[0], "_is_cached", False)
         else:
-            raise TypeError(f"Can't check this type: {adapter}!")
+            raise TypeError(f"Can't check this type: {type(adapter)}!")
 
     @classmethod
-    def flatten(cls, attr: List[List[Any]]):
-        if isinstance(attr, list):
-            if not isinstance(attr[0], list):
-                return attr
-            flatten_attr = []
-            for i in range(len(attr)):
-                flatten_attr.extend(attr[i])
-            return flatten_attr
-        return attr
+    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

@@ -1,4 +1,4 @@
-from typing import Any, Tuple, List, Dict
+from typing import Any, Tuple, List, Dict, Callable
 
 from diffusers import DiffusionPipeline
 from cache_dit.cache_factory.block_adapters.block_adapters import BlockAdapter
@@ -9,20 +9,23 @@ logger = init_logger(__name__)
 
 
 class BlockAdapterRegistry:
-    _adapters: Dict[str, BlockAdapter] = {}
-    _predefined_adapters_has_spearate_cfg: List[str] = {
+    _adapters: Dict[str, Callable[..., BlockAdapter]] = {}
+    _predefined_adapters_has_spearate_cfg: List[str] = [
         "QwenImage",
         "Wan",
         "CogView4",
         "Cosmos",
         "SkyReelsV2",
         "Chroma",
-    }
+    ]
 
     @classmethod
-    def register(cls, name):
-        def decorator(func):
-            cls._adapters[name] = func
+    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
@@ -47,15 +50,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