cache-dit 1.0.9__py3-none-any.whl → 1.0.10__py3-none-any.whl

cache_dit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '1.0.9'
-__version_tuple__ = version_tuple = (1, 0, 9)
+__version__ = version = '1.0.10'
+__version_tuple__ = version_tuple = (1, 0, 10)
 
 __commit_id__ = commit_id = None

cache_dit/cache_factory/__init__.py

@@ -8,6 +8,7 @@ from cache_dit.cache_factory.patch_functors import PatchFunctor
 
 from cache_dit.cache_factory.block_adapters import BlockAdapter
 from cache_dit.cache_factory.block_adapters import BlockAdapterRegistry
+from cache_dit.cache_factory.block_adapters import FakeDiffusionPipeline
 
 from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
 from cache_dit.cache_factory.cache_contexts import DBCacheConfig

cache_dit/cache_factory/block_adapters/__init__.py

@@ -1,5 +1,8 @@
 from cache_dit.cache_factory.forward_pattern import ForwardPattern
 from cache_dit.cache_factory.block_adapters.block_adapters import BlockAdapter
+from cache_dit.cache_factory.block_adapters.block_adapters import (
+    FakeDiffusionPipeline,
+)
 from cache_dit.cache_factory.block_adapters.block_adapters import ParamsModifier
 from cache_dit.cache_factory.block_adapters.block_registers import (
     BlockAdapterRegistry,
@@ -27,6 +30,7 @@ def flux_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_1,
                 ForwardPattern.Pattern_1,
             ],
+            check_forward_pattern=True,
             **kwargs,
         )
     else:
@@ -41,6 +45,7 @@ def flux_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_1,
                 ForwardPattern.Pattern_3,
             ],
+            check_forward_pattern=True,
             **kwargs,
         )
 
@@ -55,6 +60,7 @@ def mochi_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_0,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -69,6 +75,7 @@ def cogvideox_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_0,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -104,6 +111,7 @@ def wan_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_2,
                 ForwardPattern.Pattern_2,
             ],
+            check_forward_pattern=True,
             has_separate_cfg=True,
             **kwargs,
         )
@@ -114,6 +122,7 @@ def wan_adapter(pipe, **kwargs) -> BlockAdapter:
             transformer=pipe.transformer,
             blocks=pipe.transformer.blocks,
             forward_pattern=ForwardPattern.Pattern_2,
+            check_forward_pattern=True,
             has_separate_cfg=True,
             **kwargs,
         )
@@ -135,6 +144,7 @@ def hunyuanvideo_adapter(pipe, **kwargs) -> BlockAdapter:
             ForwardPattern.Pattern_0,
             ForwardPattern.Pattern_0,
         ],
+        check_forward_pattern=True,
         # The type hint in diffusers is wrong
         check_num_outputs=False,
         **kwargs,
@@ -159,6 +169,7 @@ def qwenimage_adapter(pipe, **kwargs) -> BlockAdapter:
             blocks=pipe.transformer.transformer_blocks,
             forward_pattern=ForwardPattern.Pattern_1,
             patch_functor=QwenImageControlNetPatchFunctor(),
+            check_forward_pattern=True,
             has_separate_cfg=True,
         )
     else:
@@ -167,6 +178,7 @@ def qwenimage_adapter(pipe, **kwargs) -> BlockAdapter:
             transformer=pipe.transformer,
             blocks=pipe.transformer.transformer_blocks,
             forward_pattern=ForwardPattern.Pattern_1,
+            check_forward_pattern=True,
             has_separate_cfg=True,
             **kwargs,
         )
@@ -182,6 +194,7 @@ def ltxvideo_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_2,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -196,6 +209,7 @@ def allegro_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_2,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -210,6 +224,7 @@ def cogview3plus_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_0,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -224,6 +239,7 @@ def cogview4_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_0,
+        check_forward_pattern=True,
         has_separate_cfg=True,
         **kwargs,
     )
@@ -239,6 +255,7 @@ def cosmos_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_2,
+        check_forward_pattern=True,
         has_separate_cfg=True,
         **kwargs,
     )
@@ -254,6 +271,7 @@ def easyanimate_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_0,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -271,6 +289,7 @@ def skyreelsv2_adapter(pipe, **kwargs) -> BlockAdapter:
         # encoder_hidden_states will never change in the blocks
         # forward loop.
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         has_separate_cfg=True,
         **kwargs,
     )
@@ -286,6 +305,7 @@ def sd3_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_1,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -300,6 +320,7 @@ def consisid_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_0,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -316,6 +337,7 @@ def dit_adapter(pipe, **kwargs) -> BlockAdapter:
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_3,
         patch_functor=DiTPatchFunctor(),
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -330,6 +352,7 @@ def amused_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_layers,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -350,6 +373,7 @@ def bria_adapter(pipe, **kwargs) -> BlockAdapter:
             ForwardPattern.Pattern_0,
             ForwardPattern.Pattern_0,
         ],
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -367,6 +391,7 @@ def lumina2_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.layers,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -381,6 +406,7 @@ def omnigen_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.layers,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -395,6 +421,7 @@ def pixart_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -409,6 +436,7 @@ def sana_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -423,6 +451,7 @@ def stabledudio_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -445,6 +474,7 @@ def visualcloze_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_1,
                 ForwardPattern.Pattern_1,
             ],
+            check_forward_pattern=True,
             **kwargs,
         )
     else:
@@ -459,6 +489,7 @@ def visualcloze_adapter(pipe, **kwargs) -> BlockAdapter:
                 ForwardPattern.Pattern_1,
                 ForwardPattern.Pattern_3,
             ],
+            check_forward_pattern=True,
             **kwargs,
         )
 
@@ -473,6 +504,7 @@ def auraflow_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.single_transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -495,6 +527,7 @@ def chroma_adapter(pipe, **kwargs) -> BlockAdapter:
             ForwardPattern.Pattern_3,
         ],
         patch_functor=ChromaPatchFunctor(),
+        check_forward_pattern=True,
         has_separate_cfg=True,
         **kwargs,
     )
@@ -510,6 +543,7 @@ def shape_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.prior,
         blocks=pipe.prior.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_3,
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -559,6 +593,7 @@ def hunyuandit_adapter(pipe, **kwargs) -> BlockAdapter:
         blocks=pipe.transformer.blocks,
         forward_pattern=ForwardPattern.Pattern_3,
         patch_functor=HunyuanDiTPatchFunctor(),
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -575,6 +610,7 @@ def hunyuanditpag_adapter(pipe, **kwargs) -> BlockAdapter:
         blocks=pipe.transformer.blocks,
         forward_pattern=ForwardPattern.Pattern_3,
         patch_functor=HunyuanDiTPatchFunctor(),
+        check_forward_pattern=True,
         **kwargs,
     )
 
@@ -613,6 +649,7 @@ def prx_adapter(pipe, **kwargs) -> BlockAdapter:
             transformer=pipe.transformer,
             blocks=pipe.transformer.blocks,
             forward_pattern=ForwardPattern.Pattern_3,
+            check_forward_pattern=True,
             check_num_outputs=False,
             **kwargs,
         )

cache_dit/cache_factory/block_adapters/block_adapters.py

@@ -6,7 +6,7 @@ from collections.abc import Iterable
 
 from typing import Any, Tuple, List, Optional, Union
 
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, ModelMixin
 from cache_dit.cache_factory.patch_functors import PatchFunctor
 from cache_dit.cache_factory.forward_pattern import ForwardPattern
 from cache_dit.cache_factory.params_modifier import ParamsModifier
@@ -16,12 +16,22 @@ from cache_dit.logger import init_logger
 logger = init_logger(__name__)
 
 
+class FakeDiffusionPipeline:
+    # A placeholder for pipelines when pipe is None.
+    def __init__(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+    ):
+        self.transformer = transformer  # Reference only
+
+
 @dataclasses.dataclass
 class BlockAdapter:
 
     # Transformer configurations.
     pipe: Union[
         DiffusionPipeline,
+        FakeDiffusionPipeline,
         Any,
     ] = None
 
@@ -73,7 +83,7 @@ class BlockAdapter:
         ]
     ] = None
 
-    check_forward_pattern: bool = True
+    check_forward_pattern: Optional[bool] = None
     check_num_outputs: bool = False
 
     # Pipeline Level Flags
@@ -110,12 +120,43 @@ class BlockAdapter:
     def __post_init__(self):
         if self.skip_post_init:
             return
+
+        self.maybe_fake_pipe()
         if any((self.pipe is not None, self.transformer is not None)):
             self.maybe_fill_attrs()
             self.maybe_patchify()
             self.maybe_skip_checks()
 
+    def maybe_fake_pipe(self):
+        if self.pipe is None:
+            self.pipe = FakeDiffusionPipeline()
+            logger.warning("pipe is None, use FakeDiffusionPipeline instead.")
+
     def maybe_skip_checks(self):
+        if self.check_forward_pattern is None:
+            if self.transformer is not None:
+                if self.nested_depth(self.transformer) == 0:
+                    transformer = self.transformer
+                elif self.nested_depth(self.transformer) == 1:
+                    transformer = self.transformer[0]
+                else:
+                    raise ValueError(
+                        "transformer nested depth can't more than 1, "
+                        f"current is: {self.nested_depth(self.transformer)}"
+                    )
+                if transformer.__module__.startswith("diffusers"):
+                    self.check_forward_pattern = True
+                    logger.info(
+                        f"Found transformer from diffusers: {transformer.__module__} "
+                        "enable check_forward_pattern by default."
+                    )
+                else:
+                    self.check_forward_pattern = False
+                    logger.info(
+                        f"Found transformer NOT from diffusers: {transformer.__module__} "
+                        "disable check_forward_pattern by default."
+                    )
+
         if getattr(self.transformer, "_hf_hook", None) is not None:
             logger.warning("_hf_hook is not None, force skip pattern check!")
             self.check_forward_pattern = False
@@ -208,7 +249,10 @@ class BlockAdapter:
             if self.transformer is not None:
                 self.patch_functor.apply(self.transformer, *args, **kwargs)
             else:
-                assert hasattr(self.pipe, "transformer")
+                assert hasattr(self.pipe, "transformer"), (
+                    "pipe.transformer can not be None when patch_functor "
+                    "is provided and transformer is None."
+                )
                 self.patch_functor.apply(self.pipe.transformer, *args, **kwargs)
 
     @staticmethod
@@ -224,6 +268,10 @@ class BlockAdapter:
             adapter.forward_pattern is not None
         ), "adapter.forward_pattern can not be None."
         pipe = adapter.pipe
+        if isinstance(pipe, FakeDiffusionPipeline):
+            raise ValueError(
+                "Can not auto block adapter for FakeDiffusionPipeline."
+            )
 
         assert hasattr(pipe, "transformer"), "pipe.transformer can not be None."
 

cache_dit/cache_factory/block_adapters/block_registers.py

@@ -1,7 +1,11 @@
-from typing import Any, Tuple, List, Dict, Callable
+import torch
+from typing import Any, Tuple, List, Dict, Callable, Union
 
 from diffusers import DiffusionPipeline
-from cache_dit.cache_factory.block_adapters.block_adapters import BlockAdapter
+from cache_dit.cache_factory.block_adapters.block_adapters import (
+    BlockAdapter,
+    FakeDiffusionPipeline,
+)
 
 from cache_dit.logger import init_logger
 
@@ -35,24 +39,42 @@ class BlockAdapterRegistry:
     @classmethod
     def get_adapter(
         cls,
-        pipe: DiffusionPipeline | str | Any,
+        pipe_or_module: DiffusionPipeline | torch.nn.Module | str | Any,
         **kwargs,
     ) -> BlockAdapter | None:
-        if not isinstance(pipe, str):
-            pipe_cls_name: str = pipe.__class__.__name__
+        if not isinstance(pipe_or_module, str):
+            cls_name: str = pipe_or_module.__class__.__name__
         else:
-            pipe_cls_name = pipe
+            cls_name = pipe_or_module
 
         for name in cls._adapters:
-            if pipe_cls_name.startswith(name):
-                return cls._adapters[name](pipe, **kwargs)
+            if cls_name.startswith(name):
+                if not isinstance(pipe_or_module, DiffusionPipeline):
+                    assert isinstance(pipe_or_module, torch.nn.Module)
+                    # NOTE: Make pre-registered adapters support Transformer-only case.
+                    # WARN: This branch is not officially supported and only for testing
+                    # purpose. We construct a fake diffusion pipeline that contains the
+                    # given transformer module. Currently, only works for DiT models which
+                    # only have one transformer module. Case like multiple transformers
+                    # is not supported, e.g, Wan2.2. Please use BlockAdapter directly for
+                    # such cases.
+                    return cls._adapters[name](
+                        FakeDiffusionPipeline(pipe_or_module), **kwargs
+                    )
+                else:
+                    return cls._adapters[name](pipe_or_module, **kwargs)
 
         return None
 
     @classmethod
     def has_separate_cfg(
         cls,
-        pipe_or_adapter: DiffusionPipeline | BlockAdapter | Any,
+        pipe_or_adapter: Union[
+            DiffusionPipeline,
+            FakeDiffusionPipeline,
+            BlockAdapter,
+            Any,
+        ],
     ) -> bool:
 
         # Prefer custom setting from block adapter.
@@ -60,11 +82,16 @@ class BlockAdapterRegistry:
             return pipe_or_adapter.has_separate_cfg
 
         has_separate_cfg = False
+        if isinstance(pipe_or_adapter, FakeDiffusionPipeline):
+            return False
+
         if isinstance(pipe_or_adapter, DiffusionPipeline):
-            has_separate_cfg = cls.get_adapter(
+            adapter = cls.get_adapter(
                 pipe_or_adapter,
                 skip_post_init=True,  # check cfg setting only
-            ).has_separate_cfg
+            )
+            if adapter is not None:
+                has_separate_cfg = adapter.has_separate_cfg
 
         if has_separate_cfg:
             return True
@@ -77,11 +104,11 @@ class BlockAdapterRegistry:
         return False
 
     @classmethod
-    def is_supported(cls, pipe) -> bool:
-        pipe_cls_name: str = pipe.__class__.__name__
+    def is_supported(cls, pipe_or_module) -> bool:
+        cls_name: str = pipe_or_module.__class__.__name__
 
         for name in cls._adapters:
-            if pipe_cls_name.startswith(name):
+            if cls_name.startswith(name):
                 return True
         return False
 

cache_dit/cache_factory/cache_adapters/cache_adapter.py

@@ -5,10 +5,11 @@ import functools
 from contextlib import ExitStack
 from typing import Dict, List, Tuple, Any, Union, Callable, Optional
 
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, ModelMixin
 
 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 FakeDiffusionPipeline
 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 ContextManager
@@ -32,6 +33,9 @@ class CachedAdapter:
         pipe_or_adapter: Union[
             DiffusionPipeline,
             BlockAdapter,
+            # Transformer-only
+            torch.nn.Module,
+            ModelMixin,
         ],
         **context_kwargs,
     ) -> Union[
@@ -42,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 "
@@ -62,10 +68,12 @@ class CachedAdapter:
                 ):
                     block_adapter.params_modifiers = params_modifiers
 
-                return cls.cachify(
-                    block_adapter,
-                    **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 "
@@ -182,8 +190,6 @@ class CachedAdapter:
         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
@@ -193,38 +199,58 @@ class CachedAdapter:
             "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._context_manager = context_manager  # instance level
-
         flatten_contexts, contexts_kwargs = cls.modify_context_params(
             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(
-                        context_manager.enter_context(
-                            context_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)
@@ -353,6 +379,7 @@ class CachedAdapter:
                 blocks_name,
                 unique_blocks_name,
                 dummy_blocks_names,
+                block_adapter,
             )
 
         return block_adapter.transformer
@@ -365,6 +392,7 @@ class CachedAdapter:
         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()
 
@@ -391,6 +419,8 @@ class CachedAdapter:
         # 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:
@@ -410,6 +440,13 @@ class CachedAdapter:
                         )
                     )
                 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):
@@ -513,6 +550,7 @@ class CachedAdapter:
             params_shift += len(blocks)
 
     @classmethod
+    @torch.compiler.disable
     def apply_stats_hooks(
         cls,
         block_adapter: BlockAdapter,