cache-dit 0.2.28__py3-none-any.whl → 0.2.29__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
 
@@ -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,92 @@ class BlockAdapter:
         default="max", metadata={"allowed_values": ["max", "min"]}
     )
 
+    # 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 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):
+                        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 +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,
@@ -164,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!")
@@ -185,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",
@@ -230,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(
@@ -293,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,
@@ -373,103 +463,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 +517,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 +534,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):
@@ -527,12 +561,46 @@ class BlockAdapter:
             raise TypeError(f"Can't check this 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

@@ -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

cache_dit/cache_factory/cache_adapters.py

@@ -29,36 +29,39 @@ class CachedAdapter:
     @classmethod
     def apply(
         cls,
-        pipe: DiffusionPipeline = None,
-        block_adapter: BlockAdapter = None,
+        pipe_or_adapter: DiffusionPipeline | BlockAdapter,
         **cache_context_kwargs,
-    ) -> DiffusionPipeline:
+    ) -> BlockAdapter:
         assert (
-            pipe is not None or block_adapter is not None
+            pipe_or_adapter is not None
         ), "pipe or block_adapter can not both None!"
 
-        if pipe is not None:
-            if BlockAdapterRegistry.is_supported(pipe):
+        if isinstance(pipe_or_adapter, DiffusionPipeline):
+            if BlockAdapterRegistry.is_supported(pipe_or_adapter):
                 logger.info(
-                    f"{pipe.__class__.__name__} is officially supported by cache-dit. "
-                    "Use it's pre-defined BlockAdapter directly!"
+                    f"{pipe_or_adapter.__class__.__name__} is officially "
+                    "supported by cache-dit. Use it's pre-defined BlockAdapter "
+                    "directly!"
+                )
+                block_adapter = BlockAdapterRegistry.get_adapter(
+                    pipe_or_adapter
                 )
-                block_adapter = BlockAdapterRegistry.get_adapter(pipe)
                 return cls.cachify(
                     block_adapter,
                     **cache_context_kwargs,
                 )
             else:
                 raise ValueError(
-                    f"{pipe.__class__.__name__} is not officially supported "
+                    f"{pipe_or_adapter.__class__.__name__} is not officially supported "
                     "by cache-dit, please set BlockAdapter instead!"
                 )
         else:
+            assert isinstance(pipe_or_adapter, BlockAdapter)
             logger.info(
-                "Adapting cache acceleration using custom BlockAdapter!"
+                "Adapting Cache Acceleration using custom BlockAdapter!"
             )
             return cls.cachify(
-                block_adapter,
+                pipe_or_adapter,
                 **cache_context_kwargs,
             )
 
@@ -67,7 +70,7 @@ class CachedAdapter:
         cls,
         block_adapter: BlockAdapter,
         **cache_context_kwargs,
-    ) -> DiffusionPipeline:
+    ) -> BlockAdapter:
 
         if block_adapter.auto:
             block_adapter = BlockAdapter.auto_block_adapter(
@@ -93,7 +96,7 @@ class CachedAdapter:
                 block_adapter,
             )
 
-        return block_adapter.pipe
+        return block_adapter
 
     @classmethod
     def patch_params(
@@ -126,18 +129,29 @@ class CachedAdapter:
             params_shift += len(blocks)
 
     @classmethod
-    def check_context_kwargs(cls, pipe, **cache_context_kwargs):
+    def check_context_kwargs(
+        cls,
+        block_adapter: BlockAdapter,
+        **cache_context_kwargs,
+    ):
         # Check cache_context_kwargs
         if not cache_context_kwargs["enable_spearate_cfg"]:
             # Check cfg for some specific case if users don't set it as True
-            cache_context_kwargs["enable_spearate_cfg"] = (
-                BlockAdapterRegistry.has_separate_cfg(pipe)
-            )
-            logger.info(
-                f"Use default 'enable_spearate_cfg': "
-                f"{cache_context_kwargs['enable_spearate_cfg']}, "
-                f"Pipeline: {pipe.__class__.__name__}."
-            )
+            if BlockAdapterRegistry.has_separate_cfg(block_adapter):
+                cache_context_kwargs["enable_spearate_cfg"] = True
+                logger.info(
+                    f"Use custom 'enable_spearate_cfg' from BlockAdapter: True. "
+                    f"Pipeline: {block_adapter.pipe.__class__.__name__}."
+                )
+            else:
+                cache_context_kwargs["enable_spearate_cfg"] = (
+                    BlockAdapterRegistry.has_separate_cfg(block_adapter.pipe)
+                )
+                logger.info(
+                    f"Use default 'enable_spearate_cfg' from block adapter "
+                    f"register: {cache_context_kwargs['enable_spearate_cfg']}, "
+                    f"Pipeline: {block_adapter.pipe.__class__.__name__}."
+                )
 
         if cache_type := cache_context_kwargs.pop("cache_type", None):
             assert (
@@ -160,8 +174,7 @@ class CachedAdapter:
 
         # Check cache_context_kwargs
         cache_context_kwargs = cls.check_context_kwargs(
-            block_adapter.pipe,
-            **cache_context_kwargs,
+            block_adapter, **cache_context_kwargs
         )
         # Apply cache on pipeline: wrap cache context
         pipe_cls_name = block_adapter.pipe.__class__.__name__

cache_dit/cache_factory/cache_blocks/utils.py

@@ -23,3 +23,19 @@ def patch_cached_stats(
     module._residual_diffs = cache_manager.get_residual_diffs()
     module._cfg_cached_steps = cache_manager.get_cfg_cached_steps()
     module._cfg_residual_diffs = cache_manager.get_cfg_residual_diffs()
+
+
+def remove_cached_stats(
+    module: torch.nn.Module | Any,
+):
+    if module is None:
+        return
+
+    if hasattr(module, "_cached_steps"):
+        del module._cached_steps
+    if hasattr(module, "_residual_diffs"):
+        del module._residual_diffs
+    if hasattr(module, "_cfg_cached_steps"):
+        del module._cfg_cached_steps
+    if hasattr(module, "_cfg_residual_diffs"):
+        del module._cfg_residual_diffs

cache_dit/cache_factory/cache_contexts/cache_manager.py

@@ -63,6 +63,20 @@ class CachedContextManager:
             _context = self.new_context(*args, **kwargs)
         return _context
 
+    def remove_context(self, cached_context: CachedContext | str):
+        if isinstance(cached_context, CachedContext):
+            cached_context.clear_buffers()
+            if cached_context.name in self._cached_context_manager:
+                del self._cached_context_manager[cached_context.name]
+        else:
+            if cached_context in self._cached_context_manager:
+                self._cached_context_manager[cached_context].clear_buffers()
+                del self._cached_context_manager[cached_context]
+
+    def clear_contexts(self):
+        for cached_context in self._cached_context_manager:
+            self.remove_context(cached_context)
+
     @contextlib.contextmanager
     def enter_context(self, cached_context: CachedContext | str):
         old_cached_context = self._current_context