cache-dit 0.2.29__py3-none-any.whl → 0.2.31__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 = '0.2.29'
-__version_tuple__ = version_tuple = (0, 2, 29)
+__version__ = version = '0.2.31'
+__version_tuple__ = version_tuple = (0, 2, 31)
 
 __commit_id__ = commit_id = None

cache_dit/cache_factory/block_adapters/__init__.py

@@ -123,6 +123,7 @@ def hunyuanvideo_adapter(pipe, **kwargs) -> BlockAdapter:
     assert isinstance(pipe.transformer, HunyuanVideoTransformer3DModel)
     return BlockAdapter(
         pipe=pipe,
+        transformer=pipe.transformer,
         blocks=[
             pipe.transformer.transformer_blocks,
             pipe.transformer.single_transformer_blocks,
@@ -131,6 +132,8 @@ def hunyuanvideo_adapter(pipe, **kwargs) -> BlockAdapter:
             ForwardPattern.Pattern_0,
             ForwardPattern.Pattern_0,
         ],
+        # The type hint in diffusers is wrong
+        check_num_outputs=False,
         **kwargs,
     )
 
@@ -327,37 +330,6 @@ def bria_adapter(pipe, **kwargs) -> BlockAdapter:
     )
 
 
-@BlockAdapterRegistry.register("HunyuanDiT")
-def hunyuandit_adapter(pipe, **kwargs) -> BlockAdapter:
-    from diffusers import HunyuanDiT2DModel, HunyuanDiT2DControlNetModel
-
-    assert isinstance(
-        pipe.transformer,
-        (HunyuanDiT2DModel, HunyuanDiT2DControlNetModel),
-    )
-    return BlockAdapter(
-        pipe=pipe,
-        transformer=pipe.transformer,
-        blocks=pipe.transformer.blocks,
-        forward_pattern=ForwardPattern.Pattern_3,
-        **kwargs,
-    )
-
-
-@BlockAdapterRegistry.register("HunyuanDiTPAG")
-def hunyuanditpag_adapter(pipe, **kwargs) -> BlockAdapter:
-    from diffusers import HunyuanDiT2DModel
-
-    assert isinstance(pipe.transformer, HunyuanDiT2DModel)
-    return BlockAdapter(
-        pipe=pipe,
-        transformer=pipe.transformer,
-        blocks=pipe.transformer.blocks,
-        forward_pattern=ForwardPattern.Pattern_3,
-        **kwargs,
-    )
-
-
 @BlockAdapterRegistry.register("Lumina")
 def lumina_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import LuminaNextDiT2DModel
@@ -414,10 +386,12 @@ def pixart_adapter(pipe, **kwargs) -> BlockAdapter:
     )
 
 
-@BlockAdapterRegistry.register("Sana")
+@BlockAdapterRegistry.register("Sana", supported=False)
 def sana_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import SanaTransformer2DModel
 
+    # TODO: fix -> got multiple values for argument 'encoder_hidden_states'
+
     assert isinstance(pipe.transformer, SanaTransformer2DModel)
     return BlockAdapter(
         pipe=pipe,
@@ -428,20 +402,6 @@ def sana_adapter(pipe, **kwargs) -> BlockAdapter:
     )
 
 
-@BlockAdapterRegistry.register("ShapE")
-def shape_adapter(pipe, **kwargs) -> BlockAdapter:
-    from diffusers import PriorTransformer
-
-    assert isinstance(pipe.prior, PriorTransformer)
-    return BlockAdapter(
-        pipe=pipe,
-        transformer=pipe.prior,
-        blocks=pipe.prior.transformer_blocks,
-        forward_pattern=ForwardPattern.Pattern_3,
-        **kwargs,
-    )
-
-
 @BlockAdapterRegistry.register("StableAudio")
 def stabledudio_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import StableAudioDiTModel
@@ -459,21 +419,37 @@ def stabledudio_adapter(pipe, **kwargs) -> BlockAdapter:
 @BlockAdapterRegistry.register("VisualCloze")
 def visualcloze_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import FluxTransformer2DModel
+    from cache_dit.utils import is_diffusers_at_least_0_3_5
 
     assert isinstance(pipe.transformer, FluxTransformer2DModel)
-    return BlockAdapter(
-        pipe=pipe,
-        transformer=pipe.transformer,
-        blocks=[
-            pipe.transformer.transformer_blocks,
-            pipe.transformer.single_transformer_blocks,
-        ],
-        forward_pattern=[
-            ForwardPattern.Pattern_1,
-            ForwardPattern.Pattern_3,
-        ],
-        **kwargs,
-    )
+    if is_diffusers_at_least_0_3_5():
+        return BlockAdapter(
+            pipe=pipe,
+            transformer=pipe.transformer,
+            blocks=[
+                pipe.transformer.transformer_blocks,
+                pipe.transformer.single_transformer_blocks,
+            ],
+            forward_pattern=[
+                ForwardPattern.Pattern_1,
+                ForwardPattern.Pattern_1,
+            ],
+            **kwargs,
+        )
+    else:
+        return BlockAdapter(
+            pipe=pipe,
+            transformer=pipe.transformer,
+            blocks=[
+                pipe.transformer.transformer_blocks,
+                pipe.transformer.single_transformer_blocks,
+            ],
+            forward_pattern=[
+                ForwardPattern.Pattern_1,
+                ForwardPattern.Pattern_3,
+            ],
+            **kwargs,
+        )
 
 
 @BlockAdapterRegistry.register("AuraFlow")
@@ -511,9 +487,29 @@ def chroma_adapter(pipe, **kwargs) -> BlockAdapter:
     )
 
 
+@BlockAdapterRegistry.register("ShapE")
+def shape_adapter(pipe, **kwargs) -> BlockAdapter:
+    from diffusers import PriorTransformer
+
+    assert isinstance(pipe.prior, PriorTransformer)
+    return BlockAdapter(
+        pipe=pipe,
+        transformer=pipe.prior,
+        blocks=pipe.prior.transformer_blocks,
+        forward_pattern=ForwardPattern.Pattern_3,
+        **kwargs,
+    )
+
+
 @BlockAdapterRegistry.register("HiDream")
 def hidream_adapter(pipe, **kwargs) -> BlockAdapter:
+    # NOTE: Need to patch Transformer forward to fully support
+    # double_stream_blocks and single_stream_blocks, namely, need
+    # to remove the logics inside the blocks forward loop:
+    # https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_hidream_image.py#L893
+    # https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_hidream_image.py#L927
     from diffusers import HiDreamImageTransformer2DModel
+    from cache_dit.cache_factory.patch_functors import HiDreamPatchFunctor
 
     assert isinstance(pipe.transformer, HiDreamImageTransformer2DModel)
     return BlockAdapter(
@@ -524,9 +520,47 @@ def hidream_adapter(pipe, **kwargs) -> BlockAdapter:
             pipe.transformer.single_stream_blocks,
         ],
         forward_pattern=[
-            ForwardPattern.Pattern_4,
+            ForwardPattern.Pattern_0,
             ForwardPattern.Pattern_3,
         ],
-        check_num_outputs=False,
+        patch_functor=HiDreamPatchFunctor(),
+        # NOTE: The type hint in diffusers is wrong
+        check_forward_pattern=True,
+        check_num_outputs=True,
+        **kwargs,
+    )
+
+
+@BlockAdapterRegistry.register("HunyuanDiT")
+def hunyuandit_adapter(pipe, **kwargs) -> BlockAdapter:
+    from diffusers import HunyuanDiT2DModel, HunyuanDiT2DControlNetModel
+    from cache_dit.cache_factory.patch_functors import HunyuanDiTPatchFunctor
+
+    assert isinstance(
+        pipe.transformer,
+        (HunyuanDiT2DModel, HunyuanDiT2DControlNetModel),
+    )
+    return BlockAdapter(
+        pipe=pipe,
+        transformer=pipe.transformer,
+        blocks=pipe.transformer.blocks,
+        forward_pattern=ForwardPattern.Pattern_3,
+        patch_functor=HunyuanDiTPatchFunctor(),
+        **kwargs,
+    )
+
+
+@BlockAdapterRegistry.register("HunyuanDiTPAG")
+def hunyuanditpag_adapter(pipe, **kwargs) -> BlockAdapter:
+    from diffusers import HunyuanDiT2DModel
+    from cache_dit.cache_factory.patch_functors import HunyuanDiTPatchFunctor
+
+    assert isinstance(pipe.transformer, HunyuanDiT2DModel)
+    return BlockAdapter(
+        pipe=pipe,
+        transformer=pipe.transformer,
+        blocks=pipe.transformer.blocks,
+        forward_pattern=ForwardPattern.Pattern_3,
+        patch_functor=HunyuanDiTPatchFunctor(),
         **kwargs,
     )

cache_dit/cache_factory/block_adapters/block_adapters.py

@@ -75,7 +75,8 @@ class BlockAdapter:
         List[List[ParamsModifier]],
     ] = None
 
-    check_num_outputs: bool = True
+    check_forward_pattern: bool = True
+    check_num_outputs: bool = False
 
     # Pipeline Level Flags
     # Patch Functor: Flux, etc.
@@ -111,9 +112,9 @@ class BlockAdapter:
     def __post_init__(self):
         if self.skip_post_init:
             return
-        assert any((self.pipe is not None, self.transformer is not None))
-        self.maybe_fill_attrs()
-        self.maybe_patchify()
+        if any((self.pipe is not None, self.transformer is not None)):
+            self.maybe_fill_attrs()
+            self.maybe_patchify()
 
     def maybe_fill_attrs(self):
         # NOTE: This func should be call before normalize.
@@ -130,7 +131,9 @@ class BlockAdapter:
                 assert isinstance(blocks, torch.nn.ModuleList)
                 blocks_name = None
                 for attr_name in attr_names:
-                    if attr := getattr(transformer, attr_name, None):
+                    if (
+                        attr := getattr(transformer, attr_name, None)
+                    ) is not None:
                         if isinstance(attr, torch.nn.ModuleList) and id(
                             attr
                         ) == id(blocks):
@@ -389,11 +392,20 @@ class BlockAdapter:
         forward_pattern: ForwardPattern,
         **kwargs,
     ) -> bool:
+
+        if not kwargs.get("check_forward_pattern", True):
+            return True
+
         assert (
             forward_pattern.Supported
             and forward_pattern in ForwardPattern.supported_patterns()
         ), f"Pattern {forward_pattern} is not support now!"
 
+        # NOTE: Special case for HiDreamBlock
+        if hasattr(block, "block"):
+            if isinstance(block.block, torch.nn.Module):
+                block = block.block
+
         forward_parameters = set(
             inspect.signature(block.forward).parameters.keys()
         )
@@ -423,6 +435,14 @@ class BlockAdapter:
         logging: bool = True,
         **kwargs,
     ) -> bool:
+
+        if not kwargs.get("check_forward_pattern", True):
+            if logging:
+                logger.warning(
+                    f"Skipped Forward Pattern Check: {forward_pattern}"
+                )
+            return True
+
         assert (
             forward_pattern.Supported
             and forward_pattern in ForwardPattern.supported_patterns()
@@ -529,6 +549,7 @@ class BlockAdapter:
                     blocks,
                     forward_pattern=forward_pattern,
                     check_num_outputs=adapter.check_num_outputs,
+                    check_forward_pattern=adapter.check_forward_pattern,
                 ), (
                     "No block forward pattern matched, "
                     f"supported lists: {ForwardPattern.supported_patterns()}"
@@ -558,7 +579,7 @@ class BlockAdapter:
             assert isinstance(adapter[0], torch.nn.Module)
             return getattr(adapter[0], "_is_cached", False)
         else:
-            raise TypeError(f"Can't check this type: {adapter}!")
+            raise TypeError(f"Can't check this type: {type(adapter)}!")
 
     @classmethod
     def nested_depth(cls, obj: Any):

cache_dit/cache_factory/block_adapters/block_registers.py

@@ -1,4 +1,4 @@
-from typing import Any, Tuple, List, Dict
+from typing import Any, Tuple, List, Dict, Callable
 
 from diffusers import DiffusionPipeline
 from cache_dit.cache_factory.block_adapters.block_adapters import BlockAdapter
@@ -9,20 +9,23 @@ logger = init_logger(__name__)
 
 
 class BlockAdapterRegistry:
-    _adapters: Dict[str, BlockAdapter] = {}
-    _predefined_adapters_has_spearate_cfg: List[str] = {
+    _adapters: Dict[str, Callable[..., BlockAdapter]] = {}
+    _predefined_adapters_has_spearate_cfg: List[str] = [
         "QwenImage",
         "Wan",
         "CogView4",
         "Cosmos",
         "SkyReelsV2",
         "Chroma",
-    }
+    ]
 
     @classmethod
-    def register(cls, name):
-        def decorator(func):
-            cls._adapters[name] = func
+    def register(cls, name: str, supported: bool = True):
+        def decorator(
+            func: Callable[..., BlockAdapter]
+        ) -> Callable[..., BlockAdapter]:
+            if supported:
+                cls._adapters[name] = func
             return func
 
         return decorator

cache_dit/cache_factory/cache_adapters.py

@@ -4,7 +4,7 @@ import unittest
 import functools
 
 from contextlib import ExitStack
-from typing import Dict, List, Tuple, Any
+from typing import Dict, List, Tuple, Any, Union, Callable
 
 from diffusers import DiffusionPipeline
 
@@ -14,7 +14,10 @@ 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_blocks.utils import (
+    patch_cached_stats,
+    remove_cached_stats,
+)
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
@@ -29,9 +32,15 @@ class CachedAdapter:
     @classmethod
     def apply(
         cls,
-        pipe_or_adapter: DiffusionPipeline | BlockAdapter,
+        pipe_or_adapter: Union[
+            DiffusionPipeline,
+            BlockAdapter,
+        ],
         **cache_context_kwargs,
-    ) -> BlockAdapter:
+    ) -> Union[
+        DiffusionPipeline,
+        BlockAdapter,
+    ]:
         assert (
             pipe_or_adapter is not None
         ), "pipe or block_adapter can not both None!"
@@ -49,7 +58,7 @@ class CachedAdapter:
                 return cls.cachify(
                     block_adapter,
                     **cache_context_kwargs,
-                )
+                ).pipe
             else:
                 raise ValueError(
                     f"{pipe_or_adapter.__class__.__name__} is not officially supported "
@@ -82,7 +91,7 @@ class CachedAdapter:
             # 0. Must normalize block_adapter before apply cache
             block_adapter = BlockAdapter.normalize(block_adapter)
             if BlockAdapter.is_cached(block_adapter):
-                return block_adapter.pipe
+                return block_adapter
 
             # 1. Apply cache on pipeline: wrap cache context, must
             # call create_context before mock_blocks.
@@ -98,36 +107,6 @@ class CachedAdapter:
 
         return block_adapter
 
-    @classmethod
-    def patch_params(
-        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 check_context_kwargs(
         cls,
@@ -153,7 +132,9 @@ class CachedAdapter:
                     f"Pipeline: {block_adapter.pipe.__class__.__name__}."
                 )
 
-        if cache_type := cache_context_kwargs.pop("cache_type", None):
+        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!"
@@ -210,14 +191,14 @@ class CachedAdapter:
                         )
                     )
                 outputs = original_call(self, *args, **kwargs)
-                cls.patch_stats(block_adapter)
+                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.patch_params(block_adapter, contexts_kwargs)
+        cls.apply_params_hooks(block_adapter, contexts_kwargs)
 
         return block_adapter.pipe
 
@@ -261,33 +242,6 @@ class CachedAdapter:
 
         return flatten_contexts, contexts_kwargs
 
-    @classmethod
-    def patch_stats(
-        cls,
-        block_adapter: BlockAdapter,
-    ):
-        from cache_dit.cache_factory.cache_blocks.utils import (
-            patch_cached_stats,
-        )
-
-        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 mock_blocks(
         cls,
@@ -391,6 +345,7 @@ class CachedAdapter:
                             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],
@@ -405,3 +360,159 @@ class CachedAdapter:
             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, CachedContextManager):
+                    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_blocks/__init__.py

@@ -25,6 +25,7 @@ class CachedBlocks:
         transformer_blocks: torch.nn.ModuleList,
         transformer: torch.nn.Module = None,
         forward_pattern: ForwardPattern = None,
+        check_forward_pattern: bool = True,
         check_num_outputs: bool = True,
         # 1. Cache context configuration
         # 'transformer_blocks', 'blocks', 'single_transformer_blocks',
@@ -45,6 +46,7 @@ class CachedBlocks:
                 transformer_blocks,
                 transformer=transformer,
                 forward_pattern=forward_pattern,
+                check_forward_pattern=check_forward_pattern,
                 check_num_outputs=check_num_outputs,
                 # 1. Cache context configuration
                 cache_prefix=cache_prefix,
@@ -58,6 +60,7 @@ class CachedBlocks:
                 transformer_blocks,
                 transformer=transformer,
                 forward_pattern=forward_pattern,
+                check_forward_pattern=check_forward_pattern,
                 check_num_outputs=check_num_outputs,
                 # 1. Cache context configuration
                 cache_prefix=cache_prefix,