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

cache_dit/cache_factory/cache_adapters/v2/cache_adapter_v2.py

@@ -1,524 +0,0 @@
-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/taylorseer.py

@@ -1,102 +0,0 @@
-import math
-import torch
-from typing import List, Dict
-
-
-class TaylorSeer:
-    def __init__(
-        self,
-        n_derivatives=2,
-        max_warmup_steps=1,
-        skip_interval_steps=1,
-        compute_step_map=None,
-    ):
-        self.n_derivatives = n_derivatives
-        self.ORDER = n_derivatives + 1
-        self.max_warmup_steps = max_warmup_steps
-        self.skip_interval_steps = skip_interval_steps
-        self.compute_step_map = compute_step_map
-        self.reset_cache()
-
-    def reset_cache(self):
-        self.state: Dict[str, List[torch.Tensor]] = {
-            "dY_prev": [None] * self.ORDER,
-            "dY_current": [None] * self.ORDER,
-        }
-        self.current_step = -1
-        self.last_non_approximated_step = -1
-
-    def should_compute_full(self, step=None):
-        step = self.current_step if step is None else step
-        if self.compute_step_map is not None:
-            return self.compute_step_map[step]
-        if (
-            step < self.max_warmup_steps
-            or (step - self.max_warmup_steps + 1) % self.skip_interval_steps
-            == 0
-        ):
-            return True
-        return False
-
-    def approximate_derivative(self, Y: torch.Tensor) -> List[torch.Tensor]:
-        # n-th order Taylor expansion:
-        # Y(t) = Y(0) + dY(0)/dt * t + d^2Y(0)/dt^2 * t^2 / 2!
-        #        + ... + d^nY(0)/dt^n * t^n / n!
-        # TODO: Custom Triton/CUDA kernel for better performance,
-        # especially for large n_derivatives.
-        dY_current: List[torch.Tensor] = [None] * self.ORDER
-        dY_current[0] = Y
-        window = self.current_step - self.last_non_approximated_step
-        if self.state["dY_prev"][0] is not None:
-            if dY_current[0].shape != self.state["dY_prev"][0].shape:
-                self.reset_cache()
-
-        for i in range(self.n_derivatives):
-            if self.state["dY_prev"][i] is not None and self.current_step > 1:
-                dY_current[i + 1] = (
-                    dY_current[i] - self.state["dY_prev"][i]
-                ) / window
-            else:
-                break
-        return dY_current
-
-    def approximate_value(self) -> torch.Tensor:
-        # TODO: Custom Triton/CUDA kernel for better performance,
-        # especially for large n_derivatives.
-        elapsed = self.current_step - self.last_non_approximated_step
-        output = 0
-        for i, derivative in enumerate(self.state["dY_current"]):
-            if derivative is not None:
-                output += (1 / math.factorial(i)) * derivative * (elapsed**i)
-            else:
-                break
-        return output
-
-    def mark_step_begin(self):
-        self.current_step += 1
-
-    def update(self, Y: torch.Tensor):
-        # Directly call this method will ingnore the warmup
-        # policy and force full computation.
-        # Assume warmup steps is 3, and n_derivatives is 3.
-        # step 0: dY_prev    = [None, None,   None,    None   ]
-        #         dY_current = [Y0,   None,   None,    None   ]
-        # step 1: dY_prev    = [Y0,   None,   None,    None   ]
-        #         dY_current = [Y1,   dY1,    None,    None   ]
-        # step 2: dY_prev    = [Y1,   dY1,    None,    None   ]
-        #         dY_current = [Y2,   dY2/Y1, dY2/dY1, None   ]
-        # step 3: dY_prev    = [Y2,   dY2/Y1, dY2/dY1, None   ],
-        #         dY_current = [Y3,   dY3/Y2, dY3/dY2, dY3/dY1]
-        # step 4: dY_prev    = [Y3,   dY3/Y2, dY3/dY2, dY3/dY1]
-        #         dY_current = [Y4,   dY4/Y3, dY4/dY3, dY4/dY2]
-        self.state["dY_prev"] = self.state["dY_current"]
-        self.state["dY_current"] = self.approximate_derivative(Y)
-        self.last_non_approximated_step = self.current_step
-
-    def step(self, Y: torch.Tensor):
-        self.mark_step_begin()
-        if self.should_compute_full():
-            self.update(Y)
-            return Y
-        else:
-            return self.approximate_value()

cache_dit/cache_factory/cache_contexts/v2/__init__.py

@@ -1,13 +0,0 @@
-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,
-)