cache-dit 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

cache_dit/__init__.py

@@ -18,6 +18,9 @@ from cache_dit.cache_factory import BlockAdapter
 from cache_dit.cache_factory import ParamsModifier
 from cache_dit.cache_factory import ForwardPattern
 from cache_dit.cache_factory import PatchFunctor
+from cache_dit.cache_factory import CalibratorConfig
+from cache_dit.cache_factory import TaylorSeerCalibratorConfig
+from cache_dit.cache_factory import FoCaCalibratorConfig
 from cache_dit.cache_factory import supported_pipelines
 from cache_dit.cache_factory import get_adapter
 from cache_dit.compile import set_compile_configs

cache_dit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.0'
-__version_tuple__ = version_tuple = (0, 3, 0)
+__version__ = version = '0.3.1'
+__version_tuple__ = version_tuple = (0, 3, 1)
 
 __commit_id__ = commit_id = None

cache_dit/cache_factory/__init__.py

@@ -12,9 +12,16 @@ from cache_dit.cache_factory.block_adapters import BlockAdapterRegistry
 
 from cache_dit.cache_factory.cache_contexts import CachedContext
 from cache_dit.cache_factory.cache_contexts import CachedContextManager
+from cache_dit.cache_factory.cache_contexts import CachedContextV2
+from cache_dit.cache_factory.cache_contexts import CachedContextManagerV2
+from cache_dit.cache_factory.cache_contexts import CalibratorConfig  # no V1
+from cache_dit.cache_factory.cache_contexts import TaylorSeerCalibratorConfig
+from cache_dit.cache_factory.cache_contexts import FoCaCalibratorConfig
+
 from cache_dit.cache_factory.cache_blocks import CachedBlocks
 
 from cache_dit.cache_factory.cache_adapters import CachedAdapter
+from cache_dit.cache_factory.cache_adapters import CachedAdapterV2
 
 from cache_dit.cache_factory.cache_interface import enable_cache
 from cache_dit.cache_factory.cache_interface import disable_cache

cache_dit/cache_factory/block_adapters/block_adapters.py

@@ -9,6 +9,7 @@ from typing import Any, Tuple, List, Optional, Union
 from diffusers import DiffusionPipeline
 from cache_dit.cache_factory.forward_pattern import ForwardPattern
 from cache_dit.cache_factory.patch_functors import PatchFunctor
+from cache_dit.cache_factory.cache_contexts import CalibratorConfig
 
 from cache_dit.logger import init_logger
 
@@ -34,6 +35,8 @@ class ParamsModifier:
         enable_encoder_taylorseer: Optional[bool] = None,
         taylorseer_cache_type: Optional[str] = None,
         taylorseer_order: Optional[int] = None,
+        # New param only for v2 API
+        calibrator_config: Optional[CalibratorConfig] = None,
         **other_cache_context_kwargs,
     ):
         self._context_kwargs = other_cache_context_kwargs.copy()
@@ -52,12 +55,19 @@ class ParamsModifier:
         self._maybe_update_param(
             "cfg_diff_compute_separate", cfg_diff_compute_separate
         )
-        self._maybe_update_param("enable_taylorseer", enable_taylorseer)
-        self._maybe_update_param(
-            "enable_encoder_taylorseer", enable_encoder_taylorseer
-        )
-        self._maybe_update_param("taylorseer_cache_type", taylorseer_cache_type)
-        self._maybe_update_param("taylorseer_order", taylorseer_order)
+        # V1 only supports the Taylorseer calibrator. We have decided to
+        # keep this code for API compatibility reasons.
+        if calibrator_config is None:
+            self._maybe_update_param("enable_taylorseer", enable_taylorseer)
+            self._maybe_update_param(
+                "enable_encoder_taylorseer", enable_encoder_taylorseer
+            )
+            self._maybe_update_param(
+                "taylorseer_cache_type", taylorseer_cache_type
+            )
+            self._maybe_update_param("taylorseer_order", taylorseer_order)
+        else:
+            self._maybe_update_param("calibrator_config", calibrator_config)
 
     def _maybe_update_param(self, key: str, value: Any):
         if value is not None:

cache_dit/cache_factory/cache_adapters/__init__.py

@@ -0,0 +1,2 @@
+from cache_dit.cache_factory.cache_adapters.cache_adapter import CachedAdapter
+from cache_dit.cache_factory.cache_adapters.v2 import CachedAdapterV2

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

@@ -8,12 +8,12 @@ from typing import Dict, List, Tuple, Any, Union, Callable
 
 from diffusers import DiffusionPipeline
 
-from cache_dit.cache_factory import CacheType
-from cache_dit.cache_factory import BlockAdapter
-from cache_dit.cache_factory import ParamsModifier
-from cache_dit.cache_factory import BlockAdapterRegistry
-from cache_dit.cache_factory import CachedContextManager
-from cache_dit.cache_factory import CachedBlocks
+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_blocks import CachedBlocks
 from cache_dit.cache_factory.cache_blocks.utils import (
     patch_cached_stats,
     remove_cached_stats,

cache_dit/cache_factory/cache_adapters/v2/__init__.py

@@ -0,0 +1,3 @@
+from cache_dit.cache_factory.cache_adapters.v2.cache_adapter_v2 import (
+    CachedAdapterV2,
+)

cache_dit/cache_factory/cache_adapters/v2/cache_adapter_v2.py

@@ -0,0 +1,524 @@
+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 CachedContextManagerV2
+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 cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+# Unified Cached Adapter
+class CachedAdapterV2:
+
+    def __call__(self, *args, **kwargs):
+        return self.apply(*args, **kwargs)
+
+    @classmethod
+    def apply(
+        cls,
+        pipe_or_adapter: Union[
+            DiffusionPipeline,
+            BlockAdapter,
+        ],
+        **cache_context_kwargs,
+    ) -> Union[
+        DiffusionPipeline,
+        BlockAdapter,
+    ]:
+        assert (
+            pipe_or_adapter is not None
+        ), "pipe or block_adapter can not both None!"
+
+        if isinstance(pipe_or_adapter, DiffusionPipeline):
+            if BlockAdapterRegistry.is_supported(pipe_or_adapter):
+                logger.info(
+                    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
+                )
+                return cls.cachify(
+                    block_adapter,
+                    **cache_context_kwargs,
+                ).pipe
+            else:
+                raise ValueError(
+                    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!"
+            )
+            return cls.cachify(
+                pipe_or_adapter,
+                **cache_context_kwargs,
+            )
+
+    @classmethod
+    def cachify(
+        cls,
+        block_adapter: BlockAdapter,
+        **cache_context_kwargs,
+    ) -> BlockAdapter:
+
+        if block_adapter.auto:
+            block_adapter = BlockAdapter.auto_block_adapter(
+                block_adapter,
+            )
+
+        if BlockAdapter.check_block_adapter(block_adapter):
+
+            # 0. Must normalize block_adapter before apply cache
+            block_adapter = BlockAdapter.normalize(block_adapter)
+            if BlockAdapter.is_cached(block_adapter):
+                return block_adapter
+
+            # 1. Apply cache on pipeline: wrap cache context, must
+            # call create_context before mock_blocks.
+            cls.create_context(
+                block_adapter,
+                **cache_context_kwargs,
+            )
+
+            # 2. Apply cache on transformer: mock cached blocks
+            cls.mock_blocks(
+                block_adapter,
+            )
+
+        return block_adapter
+
+    @classmethod
+    def check_context_kwargs(
+        cls,
+        block_adapter: BlockAdapter,
+        **cache_context_kwargs,
+    ):
+        # Check cache_context_kwargs
+        if cache_context_kwargs["enable_separate_cfg"] is None:
+            # Check cfg for some specific case if users don't set it as True
+            if BlockAdapterRegistry.has_separate_cfg(block_adapter):
+                cache_context_kwargs["enable_separate_cfg"] = True
+                logger.info(
+                    f"Use custom 'enable_separate_cfg' from BlockAdapter: True. "
+                    f"Pipeline: {block_adapter.pipe.__class__.__name__}."
+                )
+            else:
+                cache_context_kwargs["enable_separate_cfg"] = (
+                    BlockAdapterRegistry.has_separate_cfg(block_adapter.pipe)
+                )
+                logger.info(
+                    f"Use default 'enable_separate_cfg' from block adapter "
+                    f"register: {cache_context_kwargs['enable_separate_cfg']}, "
+                    f"Pipeline: {block_adapter.pipe.__class__.__name__}."
+                )
+        else:
+            logger.info(
+                f"Use custom 'enable_separate_cfg' from cache context "
+                f"kwargs: {cache_context_kwargs['enable_separate_cfg']}. "
+                f"Pipeline: {block_adapter.pipe.__class__.__name__}."
+            )
+
+        if (
+            cache_type := cache_context_kwargs.pop("cache_type", None)
+        ) is not None:
+            assert (
+                cache_type == CacheType.DBCache
+            ), "Custom cache setting only support for DBCache now!"
+
+        return cache_context_kwargs
+
+    @classmethod
+    def create_context(
+        cls,
+        block_adapter: BlockAdapter,
+        **cache_context_kwargs,
+    ) -> DiffusionPipeline:
+
+        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
+        )
+        # 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 = CachedContextManagerV2(
+            name=f"{pipe_cls_name}_{hash(id(block_adapter.pipe))}",
+        )
+        block_adapter.pipe._cache_manager = cache_manager  # instance level
+
+        flatten_contexts, contexts_kwargs = cls.modify_context_params(
+            block_adapter, cache_manager, **cache_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,
+                            ),
+                        )
+                    )
+                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
+        block_adapter.pipe.__class__._is_cached = True
+
+        cls.apply_params_hooks(block_adapter, contexts_kwargs)
+
+        return block_adapter.pipe
+
+    @classmethod
+    def modify_context_params(
+        cls,
+        block_adapter: BlockAdapter,
+        cache_manager: CachedContextManagerV2,
+        **cache_context_kwargs,
+    ) -> Tuple[List[str], List[Dict[str, Any]]]:
+
+        flatten_contexts = BlockAdapter.flatten(
+            block_adapter.unique_blocks_name
+        )
+        contexts_kwargs = [
+            cache_context_kwargs.copy()
+            for _ in range(
+                len(flatten_contexts),
+            )
+        ]
+
+        for i in range(len(contexts_kwargs)):
+            contexts_kwargs[i]["name"] = flatten_contexts[i]
+
+        if block_adapter.params_modifiers is None:
+            return flatten_contexts, contexts_kwargs
+
+        flatten_modifiers: List[ParamsModifier] = BlockAdapter.flatten(
+            block_adapter.params_modifiers,
+        )
+
+        for i in range(
+            min(len(contexts_kwargs), len(flatten_modifiers)),
+        ):
+            contexts_kwargs[i].update(
+                flatten_modifiers[i]._context_kwargs,
+            )
+            contexts_kwargs[i], _ = cache_manager.collect_cache_kwargs(
+                default_attrs={}, **contexts_kwargs[i]
+            )
+
+        return flatten_contexts, contexts_kwargs
+
+    @classmethod
+    def mock_blocks(
+        cls,
+        block_adapter: BlockAdapter,
+    ) -> List[torch.nn.Module]:
+
+        BlockAdapter.assert_normalized(block_adapter)
+
+        if BlockAdapter.is_cached(block_adapter.transformer):
+            return block_adapter.transformer
+
+        # Apply cache on transformer: mock cached transformer blocks
+        for (
+            cached_blocks,
+            transformer,
+            blocks_name,
+            unique_blocks_name,
+            dummy_blocks_names,
+        ) in zip(
+            cls.collect_cached_blocks(block_adapter),
+            block_adapter.transformer,
+            block_adapter.blocks_name,
+            block_adapter.unique_blocks_name,
+            block_adapter.dummy_blocks_names,
+        ):
+            cls.mock_transformer(
+                cached_blocks,
+                transformer,
+                blocks_name,
+                unique_blocks_name,
+                dummy_blocks_names,
+            )
+
+        return block_adapter.transformer
+
+    @classmethod
+    def mock_transformer(
+        cls,
+        cached_blocks: Dict[str, torch.nn.ModuleList],
+        transformer: torch.nn.Module,
+        blocks_name: List[str],
+        unique_blocks_name: List[str],
+        dummy_blocks_names: List[str],
+    ) -> torch.nn.Module:
+        dummy_blocks = torch.nn.ModuleList()
+
+        original_forward = transformer.forward
+
+        assert isinstance(dummy_blocks_names, list)
+
+        @functools.wraps(original_forward)
+        def new_forward(self, *args, **kwargs):
+            with ExitStack() as stack:
+                for name, context_name in zip(
+                    blocks_name,
+                    unique_blocks_name,
+                ):
+                    stack.enter_context(
+                        unittest.mock.patch.object(
+                            self, name, cached_blocks[context_name]
+                        )
+                    )
+                for dummy_name in dummy_blocks_names:
+                    stack.enter_context(
+                        unittest.mock.patch.object(
+                            self, dummy_name, dummy_blocks
+                        )
+                    )
+                return original_forward(*args, **kwargs)
+
+        transformer.forward = new_forward.__get__(transformer)
+        transformer._original_forward = original_forward
+        transformer._is_cached = True
+
+        return transformer
+
+    @classmethod
+    def collect_cached_blocks(
+        cls,
+        block_adapter: BlockAdapter,
+    ) -> 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 isinstance(
+            block_adapter.pipe._cache_manager, CachedContextManagerV2
+        )
+
+        for i in range(len(block_adapter.transformer)):
+
+            cached_blocks_bind_context = {}
+            for j in range(len(block_adapter.blocks[i])):
+                cached_blocks_bind_context[
+                    block_adapter.unique_blocks_name[i][j]
+                ] = torch.nn.ModuleList(
+                    [
+                        CachedBlocks(
+                            # 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
+                            cache_prefix=block_adapter.blocks_name[i][j],
+                            cache_context=block_adapter.unique_blocks_name[i][
+                                j
+                            ],
+                            cache_manager=block_adapter.pipe._cache_manager,
+                        )
+                    ]
+                )
+
+            total_cached_blocks.append(cached_blocks_bind_context)
+
+        return total_cached_blocks
+
+    @classmethod
+    def apply_params_hooks(
+        cls,
+        block_adapter: BlockAdapter,
+        contexts_kwargs: List[Dict],
+    ):
+        block_adapter.pipe._cache_context_kwargs = contexts_kwargs[0]
+
+        params_shift = 0
+        for i in range(len(block_adapter.transformer)):
+
+            block_adapter.transformer[i]._forward_pattern = (
+                block_adapter.forward_pattern
+            )
+            block_adapter.transformer[i]._has_separate_cfg = (
+                block_adapter.has_separate_cfg
+            )
+            block_adapter.transformer[i]._cache_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
+                ]
+
+            params_shift += len(blocks)
+
+    @classmethod
+    def apply_stats_hooks(
+        cls,
+        block_adapter: BlockAdapter,
+    ):
+        cache_manager = block_adapter.pipe._cache_manager
+
+        for i in range(len(block_adapter.transformer)):
+            patch_cached_stats(
+                block_adapter.transformer[i],
+                cache_context=block_adapter.unique_blocks_name[i][-1],
+                cache_manager=cache_manager,
+            )
+            for blocks, unique_name in zip(
+                block_adapter.blocks[i],
+                block_adapter.unique_blocks_name[i],
+            ):
+                patch_cached_stats(
+                    blocks,
+                    cache_context=unique_name,
+                    cache_manager=cache_manager,
+                )
+
+    @classmethod
+    def maybe_release_hooks(
+        cls,
+        pipe_or_adapter: Union[
+            DiffusionPipeline,
+            BlockAdapter,
+        ],
+    ):
+        # release model hooks
+        def _release_blocks_hooks(blocks):
+            return
+
+        def _release_transformer_hooks(transformer):
+            if hasattr(transformer, "_original_forward"):
+                original_forward = transformer._original_forward
+                transformer.forward = original_forward.__get__(transformer)
+                del transformer._original_forward
+            if hasattr(transformer, "_is_cached"):
+                del transformer._is_cached
+
+        def _release_pipeline_hooks(pipe):
+            if hasattr(pipe, "_original_call"):
+                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, CachedContextManagerV2):
+                    cache_manager.clear_contexts()
+                del pipe._cache_manager
+            if hasattr(pipe, "_is_cached"):
+                del pipe.__class__._is_cached
+
+        cls.release_hooks(
+            pipe_or_adapter,
+            _release_blocks_hooks,
+            _release_transformer_hooks,
+            _release_pipeline_hooks,
+        )
+
+        # release params hooks
+        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
+
+        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
+            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
+
+        cls.release_hooks(
+            pipe_or_adapter,
+            _release_blocks_params,
+            _release_transformer_params,
+            _release_pipeline_params,
+        )
+
+        # release stats hooks
+        cls.release_hooks(
+            pipe_or_adapter,
+            remove_cached_stats,
+            remove_cached_stats,
+            remove_cached_stats,
+        )
+
+    @classmethod
+    def release_hooks(
+        cls,
+        pipe_or_adapter: Union[
+            DiffusionPipeline,
+            BlockAdapter,
+        ],
+        _release_blocks: Callable,
+        _release_transformer: Callable,
+        _release_pipeline: Callable,
+    ):
+        if isinstance(pipe_or_adapter, DiffusionPipeline):
+            pipe = pipe_or_adapter
+            _release_pipeline(pipe)
+            if hasattr(pipe, "transformer"):
+                _release_transformer(pipe.transformer)
+            if hasattr(pipe, "transformer_2"):  # Wan 2.2
+                _release_transformer(pipe.transformer_2)
+        elif isinstance(pipe_or_adapter, BlockAdapter):
+            adapter = pipe_or_adapter
+            BlockAdapter.assert_normalized(adapter)
+            _release_pipeline(adapter.pipe)
+            for transformer in BlockAdapter.flatten(adapter.transformer):
+                _release_transformer(transformer)
+            for blocks in BlockAdapter.flatten(adapter.blocks):
+                _release_blocks(blocks)

cache_dit/cache_factory/cache_contexts/__init__.py

@@ -3,3 +3,10 @@ from cache_dit.cache_factory.cache_contexts.cache_context import CachedContext
 from cache_dit.cache_factory.cache_contexts.cache_manager import (
     CachedContextManager,
 )
+from cache_dit.cache_factory.cache_contexts.v2 import (
+    CachedContextV2,
+    CachedContextManagerV2,
+    CalibratorConfig,
+    TaylorSeerCalibratorConfig,
+    FoCaCalibratorConfig,
+)

cache_dit/cache_factory/cache_contexts/v2/__init__.py

@@ -0,0 +1,13 @@
+from cache_dit.cache_factory.cache_contexts.v2.calibrators import (
+    Calibrator,
+    CalibratorBase,
+    CalibratorConfig,
+    TaylorSeerCalibratorConfig,
+    FoCaCalibratorConfig,
+)
+from cache_dit.cache_factory.cache_contexts.v2.cache_context_v2 import (
+    CachedContextV2,
+)
+from cache_dit.cache_factory.cache_contexts.v2.cache_manager_v2 import (
+    CachedContextManagerV2,
+)