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

{cache_dit-0.3.1.dist-info → cache_dit-0.3.3.dist-info}/RECORD

@@ -1,36 +1,32 @@
-cache_dit/__init__.py,sha256=Nd4a609z8PLFMSO8J0sUe2xRaFDIYK8778ff8yBU7uQ,1457
-cache_dit/_version.py,sha256=gGLpQUQx-ty9SEy9PYw9OgJWWzJLBnCpfJOfzL7SjlI,704
+cache_dit/__init__.py,sha256=sHRg0swXZZiw6lvSQ53fcVtN9JRayx0az2lXAz5OOGI,1510
+cache_dit/_version.py,sha256=lemL_4Kl75FgrO6lVuFrrtw6-Dcf9wtXBalKkXuzkO4,704
 cache_dit/logger.py,sha256=0zsu42hN-3-rgGC_C29ms1IvVpV4_b4_SwJCKSenxBE,4304
-cache_dit/utils.py,sha256=bERXpCaCpOPThXB8Rkk52yAjjLrvxbt12ntpzpWdfUQ,11131
+cache_dit/utils.py,sha256=AyYRwi5XBxYBH4GaXxOxv9-X24Te_IYOYwh54t_1d3A,10674
 cache_dit/cache_factory/.gitignore,sha256=5Cb-qT9wsTUoMJ7vACDF7ZcLpAXhi5v-xdcWSRit988,23
-cache_dit/cache_factory/__init__.py,sha256=Jj_Op6ACV35XilFPax3HEEsf_hOomjmogmNyWWteq_4,1539
-cache_dit/cache_factory/cache_interface.py,sha256=xpC-CWZDBfMb5BfnXnVW25xJhV8cYMRns-LKcPDksPU,9846
+cache_dit/cache_factory/__init__.py,sha256=vy9I6Ofkj9jWeUoOvh-cY5a9QlDDKfj2FVPlVTf7BeA,1390
+cache_dit/cache_factory/cache_interface.py,sha256=A_8bBsLfGOE5wM3_rniQKPJ223_-fSpNIq65uv00sF0,10620
 cache_dit/cache_factory/cache_types.py,sha256=ooukxQRG55uTLmaZ0SKw6gIeY6SQHhMxkbv55uj2Sqk,991
 cache_dit/cache_factory/forward_pattern.py,sha256=FumlCuZ-TSmSYH0hGBHctSJ-oGLCftdZjLygqhsmdR4,2258
+cache_dit/cache_factory/params_modifier.py,sha256=zYJJsInTYCaYHBZ7mZJOP-PZnkSg3iN1WPewNOayXos,3628
 cache_dit/cache_factory/utils.py,sha256=XkVM9AXcB9zYq8-S8QKAsGz80r3tA6U3lBNGDGeHOe4,1871
 cache_dit/cache_factory/block_adapters/__init__.py,sha256=33geXMz56TxFWMp0c-H4__MY5SGRzKMKj3TXnUYOMlc,17512
-cache_dit/cache_factory/block_adapters/block_adapters.py,sha256=HlmStNIny0rZiRBYw-xdYYViVk9AEt0XlquoacEGr1U,24203
+cache_dit/cache_factory/block_adapters/block_adapters.py,sha256=2TVK_KqiYXC7AKZ2s07fzdOzUoeUBc9P1SzQtLVzhf4,22249
 cache_dit/cache_factory/block_adapters/block_registers.py,sha256=2L7QeM4ygnaKQpC9PoJod0QRYyxidUKU2AYpysDCUwE,2572
-cache_dit/cache_factory/cache_adapters/__init__.py,sha256=qB4bu1m3LgotOeNKluIkbQIf72PXpZWQMaSn1MOFEmY,149
-cache_dit/cache_factory/cache_adapters/cache_adapter.py,sha256=6WArUrTmtkZg147_Qef5jfzMVRg2hfYwvSB9Cvpf_HA,18297
-cache_dit/cache_factory/cache_adapters/v2/__init__.py,sha256=9PAH5YwpG_m0feE5eFQ7d2450nQR_Ctq8cd9Xu1Ldtk,96
-cache_dit/cache_factory/cache_adapters/v2/cache_adapter_v2.py,sha256=ove_pDh2QC3vjXWIYtrb8anc-NOmPIrDZN7hu16fjwU,18309
-cache_dit/cache_factory/cache_blocks/__init__.py,sha256=08Ox7kD05lkRKCOsVTdEZeKAWBheqpxfrAT1Nz7eclI,2916
+cache_dit/cache_factory/cache_adapters/__init__.py,sha256=py71WGD3JztQ1uk6qdLVbzYcQ1rvqFidNNaQYo7tqTo,79
+cache_dit/cache_factory/cache_adapters/cache_adapter.py,sha256=PuNFO0t9510MhOOJy93cz0uiG8PeWKsjgUWshNj76LQ,20906
+cache_dit/cache_factory/cache_blocks/__init__.py,sha256=mivvm8YOfqT7YHs8y_MzGOGztPw8LxAqKGXuSRXxCv0,3032
+cache_dit/cache_factory/cache_blocks/offload_utils.py,sha256=wusgcqaCrwEjvv7Guy-6VXhNOgPPUrBV2sSVuRmGuvo,3513
 cache_dit/cache_factory/cache_blocks/pattern_0_1_2.py,sha256=ElMps6_7uI74tSF9GDR_dEI0bZEhdzcepM29xFWnYo8,428
 cache_dit/cache_factory/cache_blocks/pattern_3_4_5.py,sha256=Bv56qETXhsREvCrNvnZpSqDIIHsi6Ze3FJW4Yk2x3uI,8597
-cache_dit/cache_factory/cache_blocks/pattern_base.py,sha256=d4H9kEB0AgnVMT8aF0Y54SUMUQUxw5HQ8gRkoCuTQ_A,14577
-cache_dit/cache_factory/cache_blocks/utils.py,sha256=dGOC1tMMOvcbvEgx44eTESKn_jsv-0RZ3tRHPa3wmQ4,1315
-cache_dit/cache_factory/cache_contexts/__init__.py,sha256=MQRxis-5gMhdJ6ZXIVN2nZEGPZoRLy59gSLniTYrWGY,437
-cache_dit/cache_factory/cache_contexts/cache_context.py,sha256=FWdgInClWY8VZBsZIevtYk--rX-RL8c3QfNOJtqR8a4,11855
-cache_dit/cache_factory/cache_contexts/cache_manager.py,sha256=Ig5VKoQ46iG3lKmsaMulYxd2vCm__2rY8NBvERwexwM,32719
-cache_dit/cache_factory/cache_contexts/taylorseer.py,sha256=4nxgSEZvDy-w-7XuJYzsyzdtF1_uFrDwlF06XBDFVKQ,3922
-cache_dit/cache_factory/cache_contexts/v2/__init__.py,sha256=GVafOd9BUa-Tyv7FZbTSkd4bGJPpMonb1AZv78qLeHU,385
-cache_dit/cache_factory/cache_contexts/v2/cache_context_v2.py,sha256=JkMJSm-zme9ayonSFq6Y6esCb6RMuGLvhVINM-LFj2Y,11776
-cache_dit/cache_factory/cache_contexts/v2/cache_manager_v2.py,sha256=ZRTl0M7jIPTIBS9lXoSh_pY6-hNu3JJ94WShv2CPWkk,32788
-cache_dit/cache_factory/cache_contexts/v2/calibrators/__init__.py,sha256=BLCV0EtOcu30iytErL_IK6J9ZwmpE6P9ffNt4OL-IaU,2343
-cache_dit/cache_factory/cache_contexts/v2/calibrators/base.py,sha256=mn6ZBkChGpGwN5csrHTUGMoX6BBPvqHXSLbIExiW-EU,748
-cache_dit/cache_factory/cache_contexts/v2/calibrators/foca.py,sha256=jrEkoiLgDR2fiX_scIpaLIDT0pTMc9stg6L9HBkgsZw,894
-cache_dit/cache_factory/cache_contexts/v2/calibrators/taylorseer.py,sha256=q5xBmT4EmpF_b3KPAjMIangTBvovE_c8ZfFjIN_E9tg,3834
+cache_dit/cache_factory/cache_blocks/pattern_base.py,sha256=wdh0bbcpKO08AW2FTsj9X_tTbFCLkDmBjrstMxTf7MQ,14668
+cache_dit/cache_factory/cache_blocks/pattern_utils.py,sha256=dGOC1tMMOvcbvEgx44eTESKn_jsv-0RZ3tRHPa3wmQ4,1315
+cache_dit/cache_factory/cache_contexts/__init__.py,sha256=T6Vak3x7Rs0Oy15Tou49p-rPQRA2jiuYtJBsbv1lBBU,388
+cache_dit/cache_factory/cache_contexts/cache_context.py,sha256=3EhaMCz3VUQ_NF81VgYwWoSEGIvhScPxPYhjL1OcgxE,15240
+cache_dit/cache_factory/cache_contexts/cache_manager.py,sha256=hSKAeP1CxmO3RFUxjFjAK1xdvVvTmeayh5jEHMaQXNE,30225
+cache_dit/cache_factory/cache_contexts/calibrators/__init__.py,sha256=mzYXO8tbytGpJJ9rpPu20kMoj1Iu_7Ym9tjfzV8rA98,5574
+cache_dit/cache_factory/cache_contexts/calibrators/base.py,sha256=mn6ZBkChGpGwN5csrHTUGMoX6BBPvqHXSLbIExiW-EU,748
+cache_dit/cache_factory/cache_contexts/calibrators/foca.py,sha256=nhHGs_hxwW1M942BQDMJb9-9IuHdnOxp774Jrna1bJI,891
+cache_dit/cache_factory/cache_contexts/calibrators/taylorseer.py,sha256=aGxr9SpytYznTepDWGPAxWDnuVMSuNyn6uNXnLh2acQ,4001
 cache_dit/cache_factory/patch_functors/__init__.py,sha256=oI6F3N9ezahRHaFUOZ1GfrAw1qFdKrxFXXmlwwehHj4,530
 cache_dit/cache_factory/patch_functors/functor_base.py,sha256=Ahk0fTfrHgNdEl-9JSkACvfyyv9G-Ei5OSz7XBIlX5o,357
 cache_dit/cache_factory/patch_functors/functor_chroma.py,sha256=xD0Q96VArp1vYBLQ0pcjRIyFB1i_Y7muZ2q07Hz8Oqs,13430
@@ -42,7 +38,7 @@ cache_dit/compile/__init__.py,sha256=FcTVzCeyypl-mxlc59_ehHL3lBNiDAFsXuRoJ-5Cfi0
 cache_dit/compile/utils.py,sha256=nN2OIrSdwRR5zGxJinKDqb07pXpvTNTF3g_OgLkeeBU,3858
 cache_dit/custom_ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cache_dit/custom_ops/triton_taylorseer.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/metrics/__init__.py,sha256=RaUhl5dieF40RqnizGzR30qoJJ9dyMUEADwgwMaMQrE,575
+cache_dit/metrics/__init__.py,sha256=UjPJ69DyyjZDfERTpKAjZKOxOTx58aWnkze7VfH3en8,673
 cache_dit/metrics/clip_score.py,sha256=ERNCFQFJKzJdbIX9OAg-1LiSPuXUVHLOFxbf2gcENpc,3938
 cache_dit/metrics/config.py,sha256=ieOgD9ayz722RjVzk24bSIqS2D6o7TZjGk8KeXV-OLQ,551
 cache_dit/metrics/fid.py,sha256=ZM_FM0XERtpnkMUfphmw2aOdljrh1uba-pnYItu0q6M,18219
@@ -53,9 +49,9 @@ cache_dit/metrics/metrics.py,sha256=7UV-H2NRbhfr6dvrXEzU97Zy-BSQ5zEfm9CKtaK4ldg,
 cache_dit/quantize/__init__.py,sha256=kWYoMAyZgBXu9BJlZjTQ0dRffW9GqeeY9_iTkXrb70A,59
 cache_dit/quantize/quantize_ao.py,sha256=Fx1KW4l3gdEkdrcAYtPoDW7WKBJWrs3glOHiEwW_TgE,6160
 cache_dit/quantize/quantize_interface.py,sha256=2s_R7xPSKuJeFpEGeLwRxnq_CqJcBG3a3lzyW5wh-UM,1241
-cache_dit-0.3.1.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
-cache_dit-0.3.1.dist-info/METADATA,sha256=I3gHe9m40_Ja0VurS7CDBYx_x_4rpra8zN245gBKv-A,46536
-cache_dit-0.3.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cache_dit-0.3.1.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
-cache_dit-0.3.1.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
-cache_dit-0.3.1.dist-info/RECORD,,
+cache_dit-0.3.3.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
+cache_dit-0.3.3.dist-info/METADATA,sha256=2kUqLHOXsbb25iz6uO8Y3pzOVMSaRHs-st6o3imjX_o,34752
+cache_dit-0.3.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-0.3.3.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
+cache_dit-0.3.3.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-0.3.3.dist-info/RECORD,,

cache_dit/cache_factory/cache_adapters/v2/__init__.py

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

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)