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

cache_dit/{cache_factory → caching}/cache_adapters/cache_adapter.py

@@ -1,25 +1,21 @@
+import copy
 import torch
-
 import unittest
 import functools
-
 from contextlib import ExitStack
-from typing import Dict, List, Tuple, Any, Union, Callable
-
-from diffusers import DiffusionPipeline
-
-from cache_dit.cache_factory.cache_types import CacheType
-from cache_dit.cache_factory.block_adapters import BlockAdapter
-from cache_dit.cache_factory.block_adapters import ParamsModifier
-from cache_dit.cache_factory.block_adapters import BlockAdapterRegistry
-from cache_dit.cache_factory.cache_contexts import CachedContextManager
-from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
-from cache_dit.cache_factory.cache_contexts import CalibratorConfig
-from cache_dit.cache_factory.cache_blocks import CachedBlocks
-from cache_dit.cache_factory.cache_blocks.utils import (
-    patch_cached_stats,
-    remove_cached_stats,
-)
+from typing import Dict, List, Tuple, Any, Union, Callable, Optional
+
+from diffusers import DiffusionPipeline, ModelMixin
+
+from cache_dit.caching.cache_types import CacheType
+from cache_dit.caching.block_adapters import BlockAdapter
+from cache_dit.caching.block_adapters import FakeDiffusionPipeline
+from cache_dit.caching.block_adapters import ParamsModifier
+from cache_dit.caching.block_adapters import BlockAdapterRegistry
+from cache_dit.caching.cache_contexts import ContextManager
+from cache_dit.caching.cache_contexts import BasicCacheConfig
+from cache_dit.caching.cache_contexts import CalibratorConfig
+from cache_dit.caching.cache_blocks import UnifiedBlocks
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
@@ -37,8 +33,11 @@ class CachedAdapter:
         pipe_or_adapter: Union[
             DiffusionPipeline,
             BlockAdapter,
+            # Transformer-only
+            torch.nn.Module,
+            ModelMixin,
         ],
-        **cache_context_kwargs,
+        **context_kwargs,
     ) -> Union[
         DiffusionPipeline,
         BlockAdapter,
@@ -47,7 +46,9 @@ class CachedAdapter:
             pipe_or_adapter is not None
         ), "pipe or block_adapter can not both None!"
 
-        if isinstance(pipe_or_adapter, DiffusionPipeline):
+        if isinstance(
+            pipe_or_adapter, (DiffusionPipeline, torch.nn.Module, ModelMixin)
+        ):
             if BlockAdapterRegistry.is_supported(pipe_or_adapter):
                 logger.info(
                     f"{pipe_or_adapter.__class__.__name__} is officially "
@@ -57,16 +58,22 @@ class CachedAdapter:
                 block_adapter = BlockAdapterRegistry.get_adapter(
                     pipe_or_adapter
                 )
-                if params_modifiers := cache_context_kwargs.pop(
+                assert block_adapter is not None, (
+                    f"BlockAdapter for {pipe_or_adapter.__class__.__name__} "
+                    "should not be None!"
+                )
+                if params_modifiers := context_kwargs.pop(
                     "params_modifiers",
                     None,
                 ):
                     block_adapter.params_modifiers = params_modifiers
 
-                return cls.cachify(
-                    block_adapter,
-                    **cache_context_kwargs,
-                ).pipe
+                block_adapter = cls.cachify(block_adapter, **context_kwargs)
+                if isinstance(pipe_or_adapter, DiffusionPipeline):
+                    return block_adapter.pipe
+
+                return block_adapter.transformer
+
             else:
                 raise ValueError(
                     f"{pipe_or_adapter.__class__.__name__} is not officially supported "
@@ -78,21 +85,21 @@ class CachedAdapter:
                 "Adapting Cache Acceleration using custom BlockAdapter!"
             )
             if pipe_or_adapter.params_modifiers is None:
-                if params_modifiers := cache_context_kwargs.pop(
+                if params_modifiers := context_kwargs.pop(
                     "params_modifiers", None
                 ):
                     pipe_or_adapter.params_modifiers = params_modifiers
 
             return cls.cachify(
                 pipe_or_adapter,
-                **cache_context_kwargs,
+                **context_kwargs,
             )
 
     @classmethod
     def cachify(
         cls,
         block_adapter: BlockAdapter,
-        **cache_context_kwargs,
+        **context_kwargs,
     ) -> BlockAdapter:
 
         if block_adapter.auto:
@@ -109,14 +116,15 @@ class CachedAdapter:
 
             # 1. Apply cache on pipeline: wrap cache context, must
             # call create_context before mock_blocks.
-            cls.create_context(
+            _, contexts_kwargs = cls.create_context(
                 block_adapter,
-                **cache_context_kwargs,
+                **context_kwargs,
             )
 
             # 2. Apply cache on transformer: mock cached blocks
             cls.mock_blocks(
                 block_adapter,
+                contexts_kwargs,
             )
 
         return block_adapter
@@ -125,12 +133,10 @@ class CachedAdapter:
     def check_context_kwargs(
         cls,
         block_adapter: BlockAdapter,
-        **cache_context_kwargs,
+        **context_kwargs,
     ):
-        # Check cache_context_kwargs
-        cache_config: BasicCacheConfig = cache_context_kwargs[
-            "cache_config"
-        ]  # ref
+        # Check context_kwargs
+        cache_config: BasicCacheConfig = context_kwargs["cache_config"]  # ref
         assert cache_config is not None, "cache_config can not be None."
         if cache_config.enable_separate_cfg is None:
             # Check cfg for some specific case if users don't set it as True
@@ -156,87 +162,113 @@ class CachedAdapter:
                 f"Pipeline: {block_adapter.pipe.__class__.__name__}."
             )
 
-        cache_type = cache_context_kwargs.pop("cache_type", None)
+        cache_type = context_kwargs.pop("cache_type", None)
         if cache_type is not None:
-            assert (
-                cache_type == CacheType.DBCache
-            ), "Custom cache setting only support for DBCache now!"
+            assert isinstance(
+                cache_type, CacheType
+            ), f"cache_type must be CacheType Enum, but got {type(cache_type)}."
+            assert cache_type == cache_config.cache_type, (
+                f"cache_type from context_kwargs ({cache_type}) must be the same "
+                f"as that from cache_config ({cache_config.cache_type})."
+            )
 
-        return cache_context_kwargs
+        return context_kwargs
 
     @classmethod
     def create_context(
         cls,
         block_adapter: BlockAdapter,
-        **cache_context_kwargs,
-    ) -> DiffusionPipeline:
+        **context_kwargs,
+    ) -> Tuple[List[str], List[Dict[str, Any]]]:
 
         BlockAdapter.assert_normalized(block_adapter)
 
         if BlockAdapter.is_cached(block_adapter.pipe):
             return block_adapter.pipe
 
-        # Check cache_context_kwargs
-        cache_context_kwargs = cls.check_context_kwargs(
-            block_adapter, **cache_context_kwargs
+        # Check context_kwargs
+        context_kwargs = cls.check_context_kwargs(
+            block_adapter, **context_kwargs
         )
-        # Apply cache on pipeline: wrap cache context
-        pipe_cls_name = block_adapter.pipe.__class__.__name__
 
         # Each Pipeline should have it's own context manager instance.
         # Different transformers (Wan2.2, etc) should shared the same
         # cache manager but with different cache context (according
         # to their unique instance id).
-        cache_manager = CachedContextManager(
+        cache_config: BasicCacheConfig = context_kwargs.get(
+            "cache_config", None
+        )
+        assert cache_config is not None, "cache_config can not be None."
+        # Apply cache on pipeline: wrap cache context
+        pipe_cls_name = block_adapter.pipe.__class__.__name__
+        context_manager = ContextManager(
             name=f"{pipe_cls_name}_{hash(id(block_adapter.pipe))}",
+            cache_type=cache_config.cache_type,
+            # Force use persistent_context for FakeDiffusionPipeline
+            persistent_context=isinstance(
+                block_adapter.pipe, FakeDiffusionPipeline
+            ),
         )
-        block_adapter.pipe._cache_manager = cache_manager  # instance level
-
         flatten_contexts, contexts_kwargs = cls.modify_context_params(
-            block_adapter, **cache_context_kwargs
+            block_adapter, **context_kwargs
         )
 
-        original_call = block_adapter.pipe.__class__.__call__
-
-        @functools.wraps(original_call)
-        def new_call(self, *args, **kwargs):
-            with ExitStack() as stack:
-                # cache context will be reset for each pipe inference
-                for context_name, context_kwargs in zip(
-                    flatten_contexts, contexts_kwargs
-                ):
-                    stack.enter_context(
-                        cache_manager.enter_context(
-                            cache_manager.reset_context(
-                                context_name,
-                                **context_kwargs,
-                            ),
+        block_adapter.pipe._context_manager = context_manager  # instance level
+
+        if not context_manager.persistent_context:
+
+            original_call = block_adapter.pipe.__class__.__call__
+
+            @functools.wraps(original_call)
+            def new_call(self, *args, **kwargs):
+                with ExitStack() as stack:
+                    # cache context will be reset for each pipe inference
+                    for context_name, context_kwargs in zip(
+                        flatten_contexts, contexts_kwargs
+                    ):
+                        stack.enter_context(
+                            context_manager.enter_context(
+                                context_manager.reset_context(
+                                    context_name,
+                                    **context_kwargs,
+                                ),
+                            )
                         )
-                    )
-                outputs = original_call(self, *args, **kwargs)
-                cls.apply_stats_hooks(block_adapter)
-                return outputs
+                    outputs = original_call(self, *args, **kwargs)
+                    cls.apply_stats_hooks(block_adapter)
+                    return outputs
+
+            block_adapter.pipe.__class__.__call__ = new_call
+            block_adapter.pipe.__class__._original_call = original_call
+
+        else:
+            # Init persistent cache context for transformer
+            for context_name, context_kwargs in zip(
+                flatten_contexts, contexts_kwargs
+            ):
+                context_manager.reset_context(
+                    context_name,
+                    **context_kwargs,
+                )
 
-        block_adapter.pipe.__class__.__call__ = new_call
-        block_adapter.pipe.__class__._original_call = original_call
         block_adapter.pipe.__class__._is_cached = True
 
         cls.apply_params_hooks(block_adapter, contexts_kwargs)
 
-        return block_adapter.pipe
+        return flatten_contexts, contexts_kwargs
 
     @classmethod
     def modify_context_params(
         cls,
         block_adapter: BlockAdapter,
-        **cache_context_kwargs,
+        **context_kwargs,
     ) -> Tuple[List[str], List[Dict[str, Any]]]:
 
         flatten_contexts = BlockAdapter.flatten(
             block_adapter.unique_blocks_name
         )
         contexts_kwargs = [
-            cache_context_kwargs.copy()
+            copy.deepcopy(context_kwargs)  # must deep copy
             for _ in range(
                 len(flatten_contexts),
             )
@@ -257,9 +289,41 @@ class CachedAdapter:
         for i in range(
             min(len(contexts_kwargs), len(flatten_modifiers)),
         ):
-            contexts_kwargs[i].update(
-                flatten_modifiers[i]._context_kwargs,
-            )
+            if "cache_config" in flatten_modifiers[i]._context_kwargs:
+                modifier_cache_config = flatten_modifiers[
+                    i
+                ]._context_kwargs.get("cache_config", None)
+                modifier_calibrator_config = flatten_modifiers[
+                    i
+                ]._context_kwargs.get("calibrator_config", None)
+                if modifier_cache_config is not None:
+                    assert isinstance(
+                        modifier_cache_config, BasicCacheConfig
+                    ), (
+                        f"cache_config must be BasicCacheConfig, but got "
+                        f"{type(modifier_cache_config)}."
+                    )
+                    contexts_kwargs[i]["cache_config"].update(
+                        **modifier_cache_config.as_dict()
+                    )
+                if modifier_calibrator_config is not None:
+                    assert isinstance(
+                        modifier_calibrator_config, CalibratorConfig
+                    ), (
+                        f"calibrator_config must be CalibratorConfig, but got "
+                        f"{type(modifier_calibrator_config)}."
+                    )
+                    if (
+                        contexts_kwargs[i].get("calibrator_config", None)
+                        is None
+                    ):
+                        contexts_kwargs[i][
+                            "calibrator_config"
+                        ] = modifier_calibrator_config
+                    else:
+                        contexts_kwargs[i]["calibrator_config"].update(
+                            **modifier_calibrator_config.as_dict()
+                        )
             cls._config_messages(**contexts_kwargs[i])
 
         return flatten_contexts, contexts_kwargs
@@ -273,7 +337,7 @@ class CachedAdapter:
             "calibrator_config", None
         )
         if cache_config is not None:
-            message = f"Collected Cache Config: {cache_config.strify()}"
+            message = f"Collected Context Config: {cache_config.strify()}"
             if calibrator_config is not None:
                 message += f", Calibrator Config: {calibrator_config.strify(details=True)}"
             else:
@@ -284,6 +348,7 @@ class CachedAdapter:
     def mock_blocks(
         cls,
         block_adapter: BlockAdapter,
+        contexts_kwargs: List[Dict],
     ) -> List[torch.nn.Module]:
 
         BlockAdapter.assert_normalized(block_adapter)
@@ -293,24 +358,28 @@ class CachedAdapter:
 
         # Apply cache on transformer: mock cached transformer blocks
         for (
-            cached_blocks,
+            unified_blocks,
             transformer,
             blocks_name,
             unique_blocks_name,
             dummy_blocks_names,
         ) in zip(
-            cls.collect_cached_blocks(block_adapter),
+            cls.collect_unified_blocks(
+                block_adapter,
+                contexts_kwargs,
+            ),
             block_adapter.transformer,
             block_adapter.blocks_name,
             block_adapter.unique_blocks_name,
             block_adapter.dummy_blocks_names,
         ):
             cls.mock_transformer(
-                cached_blocks,
+                unified_blocks,
                 transformer,
                 blocks_name,
                 unique_blocks_name,
                 dummy_blocks_names,
+                block_adapter,
             )
 
         return block_adapter.transformer
@@ -318,11 +387,12 @@ class CachedAdapter:
     @classmethod
     def mock_transformer(
         cls,
-        cached_blocks: Dict[str, torch.nn.ModuleList],
+        unified_blocks: Dict[str, torch.nn.ModuleList],
         transformer: torch.nn.Module,
         blocks_name: List[str],
         unique_blocks_name: List[str],
         dummy_blocks_names: List[str],
+        block_adapter: BlockAdapter,
     ) -> torch.nn.Module:
         dummy_blocks = torch.nn.ModuleList()
 
@@ -330,7 +400,28 @@ class CachedAdapter:
 
         assert isinstance(dummy_blocks_names, list)
 
-        @functools.wraps(original_forward)
+        from accelerate import hooks
+
+        _hf_hook: Optional[hooks.ModelHook] = None
+
+        if getattr(transformer, "_hf_hook", None) is not None:
+            _hf_hook = transformer._hf_hook  # hooks from accelerate.hooks
+            if hasattr(transformer, "_old_forward"):
+                logger.warning(
+                    "_hf_hook is not None, so, we have to re-direct transformer's "
+                    f"original_forward({id(original_forward)}) to transformer's "
+                    f"_old_forward({id(transformer._old_forward)})"
+                )
+                original_forward = transformer._old_forward
+
+        # TODO: remove group offload hooks the re-apply after cache applied.
+        # hooks = _diffusers_hook.hooks.copy(); _diffusers_hook.hooks.clear()
+        # re-apply hooks to transformer after cache applied.
+        # from diffusers.hooks.hooks import HookFunctionReference, HookRegistry
+        # from diffusers.hooks.group_offloading import apply_group_offloading
+        context_manager: ContextManager = block_adapter.pipe._context_manager
+        assert isinstance(context_manager, ContextManager._supported_managers)
+
         def new_forward(self, *args, **kwargs):
             with ExitStack() as stack:
                 for name, context_name in zip(
@@ -339,7 +430,7 @@ class CachedAdapter:
                 ):
                     stack.enter_context(
                         unittest.mock.patch.object(
-                            self, name, cached_blocks[context_name]
+                            self, name, unified_blocks[context_name]
                         )
                     )
                 for dummy_name in dummy_blocks_names:
@@ -348,55 +439,85 @@ class CachedAdapter:
                             self, dummy_name, dummy_blocks
                         )
                     )
-                return original_forward(*args, **kwargs)
+                outputs = original_forward(*args, **kwargs)
+
+                if (
+                    context_manager.persistent_context
+                    and context_manager.is_pre_refreshed()
+                ):
+                    cls.apply_stats_hooks(block_adapter)
+
+            return outputs
+
+        def new_forward_with_hf_hook(self, *args, **kwargs):
+            # Compatible with model cpu offload
+            if _hf_hook is not None and hasattr(_hf_hook, "pre_forward"):
+                args, kwargs = _hf_hook.pre_forward(self, *args, **kwargs)
+
+            outputs = new_forward(self, *args, **kwargs)
+
+            if _hf_hook is not None and hasattr(_hf_hook, "post_forward"):
+                outputs = _hf_hook.post_forward(self, outputs)
+
+            return outputs
+
+        # NOTE: Still can't fully compatible with group offloading
+        transformer.forward = functools.update_wrapper(
+            functools.partial(new_forward_with_hf_hook, transformer),
+            new_forward_with_hf_hook,
+        )
 
-        transformer.forward = new_forward.__get__(transformer)
         transformer._original_forward = original_forward
         transformer._is_cached = True
 
         return transformer
 
     @classmethod
-    def collect_cached_blocks(
+    def collect_unified_blocks(
         cls,
         block_adapter: BlockAdapter,
+        contexts_kwargs: List[Dict],
     ) -> List[Dict[str, torch.nn.ModuleList]]:
 
         BlockAdapter.assert_normalized(block_adapter)
 
         total_cached_blocks: List[Dict[str, torch.nn.ModuleList]] = []
-        assert hasattr(block_adapter.pipe, "_cache_manager")
+        assert hasattr(block_adapter.pipe, "_context_manager")
         assert isinstance(
-            block_adapter.pipe._cache_manager,
-            CachedContextManager,
+            block_adapter.pipe._context_manager,
+            ContextManager._supported_managers,
         )
 
         for i in range(len(block_adapter.transformer)):
 
-            cached_blocks_bind_context = {}
+            unified_blocks_bind_context = {}
             for j in range(len(block_adapter.blocks[i])):
-                cached_blocks_bind_context[
+                cache_config: BasicCacheConfig = contexts_kwargs[
+                    i * len(block_adapter.blocks[i]) + j
+                ]["cache_config"]
+                unified_blocks_bind_context[
                     block_adapter.unique_blocks_name[i][j]
                 ] = torch.nn.ModuleList(
                     [
-                        CachedBlocks(
+                        UnifiedBlocks(
                             # 0. Transformer blocks configuration
                             block_adapter.blocks[i][j],
                             transformer=block_adapter.transformer[i],
                             forward_pattern=block_adapter.forward_pattern[i][j],
                             check_forward_pattern=block_adapter.check_forward_pattern,
                             check_num_outputs=block_adapter.check_num_outputs,
-                            # 1. Cache context configuration
+                            # 1. Cache/Prune context configuration
                             cache_prefix=block_adapter.blocks_name[i][j],
                             cache_context=block_adapter.unique_blocks_name[i][
                                 j
                             ],
-                            cache_manager=block_adapter.pipe._cache_manager,
+                            context_manager=block_adapter.pipe._context_manager,
+                            cache_type=cache_config.cache_type,
                         )
                     ]
                 )
 
-            total_cached_blocks.append(cached_blocks_bind_context)
+            total_cached_blocks.append(unified_blocks_bind_context)
 
         return total_cached_blocks
 
@@ -406,7 +527,7 @@ class CachedAdapter:
         block_adapter: BlockAdapter,
         contexts_kwargs: List[Dict],
     ):
-        block_adapter.pipe._cache_context_kwargs = contexts_kwargs[0]
+        block_adapter.pipe._context_kwargs = contexts_kwargs[0]
 
         params_shift = 0
         for i in range(len(block_adapter.transformer)):
@@ -417,40 +538,43 @@ class CachedAdapter:
             block_adapter.transformer[i]._has_separate_cfg = (
                 block_adapter.has_separate_cfg
             )
-            block_adapter.transformer[i]._cache_context_kwargs = (
-                contexts_kwargs[params_shift]
-            )
+            block_adapter.transformer[i]._context_kwargs = contexts_kwargs[
+                params_shift
+            ]
 
             blocks = block_adapter.blocks[i]
             for j in range(len(blocks)):
                 blocks[j]._forward_pattern = block_adapter.forward_pattern[i][j]
-                blocks[j]._cache_context_kwargs = contexts_kwargs[
-                    params_shift + j
-                ]
+                blocks[j]._context_kwargs = contexts_kwargs[params_shift + j]
 
             params_shift += len(blocks)
 
     @classmethod
+    @torch.compiler.disable
     def apply_stats_hooks(
         cls,
         block_adapter: BlockAdapter,
     ):
-        cache_manager = block_adapter.pipe._cache_manager
+        from cache_dit.caching.cache_blocks import (
+            apply_stats,
+        )
+
+        context_manager = block_adapter.pipe._context_manager
 
         for i in range(len(block_adapter.transformer)):
-            patch_cached_stats(
+            apply_stats(
                 block_adapter.transformer[i],
                 cache_context=block_adapter.unique_blocks_name[i][-1],
-                cache_manager=cache_manager,
+                context_manager=context_manager,
             )
             for blocks, unique_name in zip(
                 block_adapter.blocks[i],
                 block_adapter.unique_blocks_name[i],
             ):
-                patch_cached_stats(
+                apply_stats(
                     blocks,
                     cache_context=unique_name,
-                    cache_manager=cache_manager,
+                    context_manager=context_manager,
                 )
 
     @classmethod
@@ -478,11 +602,13 @@ class CachedAdapter:
                 original_call = pipe.__class__._original_call
                 pipe.__class__.__call__ = original_call
                 del pipe.__class__._original_call
-            if hasattr(pipe, "_cache_manager"):
-                cache_manager = pipe._cache_manager
-                if isinstance(cache_manager, CachedContextManager):
-                    cache_manager.clear_contexts()
-                del pipe._cache_manager
+            if hasattr(pipe, "_context_manager"):
+                context_manager = pipe._context_manager
+                if isinstance(
+                    context_manager, ContextManager._supported_managers
+                ):
+                    context_manager.clear_contexts()
+                del pipe._context_manager
             if hasattr(pipe, "_is_cached"):
                 del pipe.__class__._is_cached
 
@@ -497,22 +623,22 @@ class CachedAdapter:
         def _release_blocks_params(blocks):
             if hasattr(blocks, "_forward_pattern"):
                 del blocks._forward_pattern
-            if hasattr(blocks, "_cache_context_kwargs"):
-                del blocks._cache_context_kwargs
+            if hasattr(blocks, "_context_kwargs"):
+                del blocks._context_kwargs
 
         def _release_transformer_params(transformer):
             if hasattr(transformer, "_forward_pattern"):
                 del transformer._forward_pattern
             if hasattr(transformer, "_has_separate_cfg"):
                 del transformer._has_separate_cfg
-            if hasattr(transformer, "_cache_context_kwargs"):
-                del transformer._cache_context_kwargs
+            if hasattr(transformer, "_context_kwargs"):
+                del transformer._context_kwargs
             for blocks in BlockAdapter.find_blocks(transformer):
                 _release_blocks_params(blocks)
 
         def _release_pipeline_params(pipe):
-            if hasattr(pipe, "_cache_context_kwargs"):
-                del pipe._cache_context_kwargs
+            if hasattr(pipe, "_context_kwargs"):
+                del pipe._context_kwargs
 
         cls.release_hooks(
             pipe_or_adapter,
@@ -522,11 +648,24 @@ class CachedAdapter:
         )
 
         # release stats hooks
+        from cache_dit.caching.cache_blocks import (
+            remove_stats,
+        )
+
+        cls.release_hooks(
+            pipe_or_adapter, remove_stats, remove_stats, remove_stats
+        )
+
+        # maybe release parallelism stats
+        from cache_dit.parallelism.parallel_interface import (
+            remove_parallelism_stats,
+        )
+
         cls.release_hooks(
             pipe_or_adapter,
-            remove_cached_stats,
-            remove_cached_stats,
-            remove_cached_stats,
+            remove_parallelism_stats,
+            remove_parallelism_stats,
+            remove_parallelism_stats,
         )
 
     @classmethod