cache-dit 0.2.29__py3-none-any.whl → 0.2.30__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.30'
+__version_tuple__ = version_tuple = (0, 2, 30)
 
 __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,8 +487,27 @@ def chroma_adapter(pipe, **kwargs) -> BlockAdapter:
     )
 
 
-@BlockAdapterRegistry.register("HiDream")
+@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", supported=True)
 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
 
     assert isinstance(pipe.transformer, HiDreamImageTransformer2DModel)
@@ -520,13 +515,47 @@ def hidream_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=[
-            pipe.transformer.double_stream_blocks,
+            # pipe.transformer.double_stream_blocks,
             pipe.transformer.single_stream_blocks,
         ],
         forward_pattern=[
-            ForwardPattern.Pattern_4,
+            # ForwardPattern.Pattern_4,
             ForwardPattern.Pattern_3,
         ],
+        # The type hint in diffusers is wrong
         check_num_outputs=False,
         **kwargs,
     )
+
+
+@BlockAdapterRegistry.register("HunyuanDiT", supported=False)
+def hunyuandit_adapter(pipe, **kwargs) -> BlockAdapter:
+    # TODO: Patch Transformer forward
+    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", supported=False)
+def hunyuanditpag_adapter(pipe, **kwargs) -> BlockAdapter:
+    # TODO: Patch Transformer forward
+    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,
+    )

cache_dit/cache_factory/block_adapters/block_adapters.py

@@ -75,7 +75,7 @@ class BlockAdapter:
         List[List[ParamsModifier]],
     ] = None
 
-    check_num_outputs: bool = True
+    check_num_outputs: bool = False
 
     # Pipeline Level Flags
     # Patch Functor: Flux, etc.
@@ -111,9 +111,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 +130,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):
@@ -558,7 +560,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,
@@ -405,3 +359,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/pattern_3_4_5.py

@@ -1,5 +1,6 @@
 import torch
 
+from typing import Dict, Any
 from cache_dit.cache_factory import ForwardPattern
 from cache_dit.cache_factory.cache_blocks.pattern_base import (
     CachedBlocks_Pattern_Base,
@@ -31,7 +32,7 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
         # Call first `n` blocks to process the hidden states for
         # more stable diff calculation.
         # encoder_hidden_states: None Pattern 3, else 4, 5
-        hidden_states, encoder_hidden_states = self.call_Fn_blocks(
+        hidden_states, new_encoder_hidden_states = self.call_Fn_blocks(
             hidden_states,
             *args,
             **kwargs,
@@ -60,11 +61,10 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
         if can_use_cache:
             self.cache_manager.add_cached_step()
             del Fn_hidden_states_residual
-            hidden_states, encoder_hidden_states = (
+            hidden_states, new_encoder_hidden_states = (
                 self.cache_manager.apply_cache(
                     hidden_states,
-                    # None Pattern 3, else 4, 5
-                    encoder_hidden_states,
+                    new_encoder_hidden_states,  # encoder_hidden_states not use cache
                     prefix=(
                         f"{self.cache_prefix}_Bn_residual"
                         if self.cache_manager.is_cache_residual()
@@ -80,12 +80,12 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
             torch._dynamo.graph_break()
             # Call last `n` blocks to further process the hidden states
             # for higher precision.
-            hidden_states, encoder_hidden_states = self.call_Bn_blocks(
-                hidden_states,
-                encoder_hidden_states,
-                *args,
-                **kwargs,
-            )
+            if self.cache_manager.Bn_compute_blocks() > 0:
+                hidden_states, new_encoder_hidden_states = self.call_Bn_blocks(
+                    hidden_states,
+                    *args,
+                    **kwargs,
+                )
         else:
             self.cache_manager.set_Fn_buffer(
                 Fn_hidden_states_residual,
@@ -99,19 +99,20 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 )
             del Fn_hidden_states_residual
             torch._dynamo.graph_break()
+            old_encoder_hidden_states = new_encoder_hidden_states
             (
                 hidden_states,
-                encoder_hidden_states,
+                new_encoder_hidden_states,
                 hidden_states_residual,
-                # None Pattern 3, else 4, 5
-                encoder_hidden_states_residual,
             ) = self.call_Mn_blocks(  # middle
                 hidden_states,
-                # None Pattern 3, else 4, 5
-                encoder_hidden_states,
                 *args,
                 **kwargs,
             )
+            if new_encoder_hidden_states is not None:
+                new_encoder_hidden_states_residual = (
+                    new_encoder_hidden_states - old_encoder_hidden_states
+                )
             torch._dynamo.graph_break()
             if self.cache_manager.is_cache_residual():
                 self.cache_manager.set_Bn_buffer(
@@ -119,34 +120,32 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                     prefix=f"{self.cache_prefix}_Bn_residual",
                 )
             else:
-                # TaylorSeer
                 self.cache_manager.set_Bn_buffer(
                     hidden_states,
                     prefix=f"{self.cache_prefix}_Bn_hidden_states",
                 )
+
             if self.cache_manager.is_encoder_cache_residual():
-                self.cache_manager.set_Bn_encoder_buffer(
-                    # None Pattern 3, else 4, 5
-                    encoder_hidden_states_residual,
-                    prefix=f"{self.cache_prefix}_Bn_residual",
-                )
+                if new_encoder_hidden_states is not None:
+                    self.cache_manager.set_Bn_encoder_buffer(
+                        new_encoder_hidden_states_residual,
+                        prefix=f"{self.cache_prefix}_Bn_residual",
+                    )
             else:
-                # TaylorSeer
-                self.cache_manager.set_Bn_encoder_buffer(
-                    # None Pattern 3, else 4, 5
-                    encoder_hidden_states,
-                    prefix=f"{self.cache_prefix}_Bn_hidden_states",
-                )
+                if new_encoder_hidden_states is not None:
+                    self.cache_manager.set_Bn_encoder_buffer(
+                        new_encoder_hidden_states_residual,
+                        prefix=f"{self.cache_prefix}_Bn_hidden_states",
+                    )
             torch._dynamo.graph_break()
             # Call last `n` blocks to further process the hidden states
             # for higher precision.
-            hidden_states, encoder_hidden_states = self.call_Bn_blocks(
-                hidden_states,
-                # None Pattern 3, else 4, 5
-                encoder_hidden_states,
-                *args,
-                **kwargs,
-            )
+            if self.cache_manager.Bn_compute_blocks() > 0:
+                hidden_states, new_encoder_hidden_states = self.call_Bn_blocks(
+                    hidden_states,
+                    *args,
+                    **kwargs,
+                )
 
         torch._dynamo.graph_break()
 
@@ -154,12 +153,21 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
             hidden_states
             if self.forward_pattern.Return_H_Only
             else (
-                (hidden_states, encoder_hidden_states)
+                (hidden_states, new_encoder_hidden_states)
                 if self.forward_pattern.Return_H_First
-                else (encoder_hidden_states, hidden_states)
+                else (new_encoder_hidden_states, hidden_states)
             )
         )
 
+    @torch.compiler.disable
+    def maybe_update_kwargs(
+        self, encoder_hidden_states, kwargs: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        # if "encoder_hidden_states" in kwargs:
+        #     kwargs["encoder_hidden_states"] = encoder_hidden_states
+        # return kwargs
+        return kwargs
+
     def call_Fn_blocks(
         self,
         hidden_states: torch.Tensor,
@@ -172,7 +180,7 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
             f"Fn_compute_blocks {self.cache_manager.Fn_compute_blocks()} must be less than "
             f"the number of transformer blocks {len(self.transformer_blocks)}"
         )
-        encoder_hidden_states = None  # Pattern 3
+        new_encoder_hidden_states = None
         for block in self._Fn_blocks():
             hidden_states = block(
                 hidden_states,
@@ -180,25 +188,27 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 **kwargs,
             )
             if not isinstance(hidden_states, torch.Tensor):  # Pattern 4, 5
-                hidden_states, encoder_hidden_states = hidden_states
+                hidden_states, new_encoder_hidden_states = hidden_states
                 if not self.forward_pattern.Return_H_First:
-                    hidden_states, encoder_hidden_states = (
-                        encoder_hidden_states,
+                    hidden_states, new_encoder_hidden_states = (
+                        new_encoder_hidden_states,
                         hidden_states,
                     )
+            kwargs = self.maybe_update_kwargs(
+                new_encoder_hidden_states,
+                kwargs,
+            )
 
-        return hidden_states, encoder_hidden_states
+        return hidden_states, new_encoder_hidden_states
 
     def call_Mn_blocks(
         self,
         hidden_states: torch.Tensor,
-        # None Pattern 3, else 4, 5
-        encoder_hidden_states: torch.Tensor | None,
         *args,
         **kwargs,
     ):
         original_hidden_states = hidden_states
-        original_encoder_hidden_states = encoder_hidden_states
+        new_encoder_hidden_states = None
         for block in self._Mn_blocks():
             hidden_states = block(
                 hidden_states,
@@ -206,44 +216,33 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 **kwargs,
             )
             if not isinstance(hidden_states, torch.Tensor):  # Pattern 4, 5
-                hidden_states, encoder_hidden_states = hidden_states
+                hidden_states, new_encoder_hidden_states = hidden_states
                 if not self.forward_pattern.Return_H_First:
-                    hidden_states, encoder_hidden_states = (
-                        encoder_hidden_states,
+                    hidden_states, new_encoder_hidden_states = (
+                        new_encoder_hidden_states,
                         hidden_states,
                     )
+            kwargs = self.maybe_update_kwargs(
+                new_encoder_hidden_states,
+                kwargs,
+            )
 
         # compute hidden_states residual
         hidden_states = hidden_states.contiguous()
         hidden_states_residual = hidden_states - original_hidden_states
-        if (
-            original_encoder_hidden_states is not None
-            and encoder_hidden_states is not None
-        ):  # Pattern 4, 5
-            encoder_hidden_states_residual = (
-                encoder_hidden_states - original_encoder_hidden_states
-            )
-        else:
-            encoder_hidden_states_residual = None  # Pattern 3
 
         return (
             hidden_states,
-            encoder_hidden_states,
+            new_encoder_hidden_states,
             hidden_states_residual,
-            encoder_hidden_states_residual,
         )
 
     def call_Bn_blocks(
         self,
         hidden_states: torch.Tensor,
-        # None Pattern 3, else 4, 5
-        encoder_hidden_states: torch.Tensor | None,
         *args,
         **kwargs,
     ):
-        if self.cache_manager.Bn_compute_blocks() == 0:
-            return hidden_states, encoder_hidden_states
-
         assert self.cache_manager.Bn_compute_blocks() <= len(
             self.transformer_blocks
         ), (
@@ -264,11 +263,15 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                     **kwargs,
                 )
                 if not isinstance(hidden_states, torch.Tensor):  # Pattern 4,5
-                    hidden_states, encoder_hidden_states = hidden_states
+                    hidden_states, new_encoder_hidden_states = hidden_states
                     if not self.forward_pattern.Return_H_First:
-                        hidden_states, encoder_hidden_states = (
-                            encoder_hidden_states,
+                        hidden_states, new_encoder_hidden_states = (
+                            new_encoder_hidden_states,
                             hidden_states,
                         )
+                kwargs = self.maybe_update_kwargs(
+                    new_encoder_hidden_states,
+                    kwargs,
+                )
 
-        return hidden_states, encoder_hidden_states
+        return hidden_states, new_encoder_hidden_states

cache_dit/cache_factory/cache_contexts/cache_manager.py

@@ -733,17 +733,15 @@ class CachedContextManager:
                     encoder_prefix
                 )
 
-            assert (
-                encoder_hidden_states_prev is not None
-            ), f"{prefix}_encoder_buffer must be set before"
+            if encoder_hidden_states_prev is not None:
 
-            if self.is_encoder_cache_residual():
-                encoder_hidden_states = (
-                    encoder_hidden_states_prev + encoder_hidden_states
-                )
-            else:
-                # If encoder cache is not residual, we use the encoder hidden states directly
-                encoder_hidden_states = encoder_hidden_states_prev
+                if self.is_encoder_cache_residual():
+                    encoder_hidden_states = (
+                        encoder_hidden_states_prev + encoder_hidden_states
+                    )
+                else:
+                    # If encoder cache is not residual, we use the encoder hidden states directly
+                    encoder_hidden_states = encoder_hidden_states_prev
 
             encoder_hidden_states = encoder_hidden_states.contiguous()
 

cache_dit/cache_factory/cache_interface.py

@@ -1,11 +1,9 @@
-import torch
-from typing import Any, Tuple, List
+from typing import Any, Tuple, List, Union
 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 BlockAdapterRegistry
 from cache_dit.cache_factory.cache_adapters import CachedAdapter
-from cache_dit.cache_factory.cache_contexts import CachedContextManager
 
 from cache_dit.logger import init_logger
 
@@ -14,7 +12,10 @@ logger = init_logger(__name__)
 
 def enable_cache(
     # DiffusionPipeline or BlockAdapter
-    pipe_or_adapter: DiffusionPipeline | BlockAdapter,
+    pipe_or_adapter: Union[
+        DiffusionPipeline,
+        BlockAdapter,
+    ],
     # Cache context kwargs
     Fn_compute_blocks: int = 8,
     Bn_compute_blocks: int = 0,
@@ -32,7 +33,10 @@ def enable_cache(
     taylorseer_cache_type: str = "residual",
     taylorseer_order: int = 2,
     **other_cache_context_kwargs,
-) -> BlockAdapter:
+) -> Union[
+    DiffusionPipeline,
+    BlockAdapter,
+]:
     r"""
     Unified Cache API for  almost Any Diffusion Transformers (with Transformer Blocks
     that match the specific Input and Output patterns).
@@ -102,11 +106,11 @@ def enable_cache(
     >>> cache_dit.enable_cache(pipe) # One-line code with default cache options.
     >>> output = pipe(...) # Just call the pipe as normal.
     >>> stats = cache_dit.summary(pipe) # Then, get the summary of cache acceleration stats.
+    >>> cache_dit.disable_cache(pipe) # Disable cache and run original pipe.
     """
-
     # Collect cache context kwargs
     cache_context_kwargs = other_cache_context_kwargs.copy()
-    if cache_type := cache_context_kwargs.get("cache_type", None):
+    if (cache_type := cache_context_kwargs.get("cache_type", None)) is not None:
         if cache_type == CacheType.NONE:
             return pipe_or_adapter
 
@@ -145,79 +149,17 @@ def enable_cache(
 
 
 def disable_cache(
-    # DiffusionPipeline or BlockAdapter
-    pipe_or_adapter: DiffusionPipeline | BlockAdapter,
+    pipe_or_adapter: Union[
+        DiffusionPipeline,
+        BlockAdapter,
+    ],
 ):
-    from cache_dit.cache_factory.cache_blocks.utils import (
-        remove_cached_stats,
+    CachedAdapter.maybe_release_hooks(pipe_or_adapter)
+    logger.warning(
+        f"Cache Acceleration is disabled for: "
+        f"{pipe_or_adapter.__class__.__name__}."
     )
 
-    def _disable_blocks(blocks: torch.nn.ModuleList):
-        if blocks is None:
-            return
-        if hasattr(blocks, "_forward_pattern"):
-            del blocks._forward_pattern
-        if hasattr(blocks, "_cache_context_kwargs"):
-            del blocks._cache_context_kwargs
-        remove_cached_stats(blocks)
-
-    def _disable_transformer(transformer: torch.nn.Module):
-        if transformer is None or not BlockAdapter.is_cached(transformer):
-            return
-        if original_forward := getattr(transformer, "_original_forward"):
-            transformer.forward = original_forward.__get__(transformer)
-            del transformer._original_forward
-        if hasattr(transformer, "_is_cached"):
-            del transformer._is_cached
-        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
-        remove_cached_stats(transformer)
-        for blocks in BlockAdapter.find_blocks(transformer):
-            _disable_blocks(blocks)
-
-    def _disable_pipe(pipe: DiffusionPipeline):
-        if pipe is None or not BlockAdapter.is_cached(pipe):
-            return
-        if original_call := getattr(pipe, "_original_call"):
-            pipe.__class__.__call__ = original_call
-            del pipe.__class__._original_call
-        if cache_manager := getattr(pipe, "_cache_manager"):
-            assert isinstance(cache_manager, CachedContextManager)
-            cache_manager.clear_contexts()
-            del pipe._cache_manager
-        if hasattr(pipe, "_is_cached"):
-            del pipe.__class__._is_cached
-        if hasattr(pipe, "_cache_context_kwargs"):
-            del pipe._cache_context_kwargs
-        remove_cached_stats(pipe)
-
-    if isinstance(pipe_or_adapter, DiffusionPipeline):
-        pipe = pipe_or_adapter
-        _disable_pipe(pipe)
-        if hasattr(pipe, "transformer"):
-            _disable_transformer(pipe.transformer)
-        if hasattr(pipe, "transformer_2"):  # Wan 2.2
-            _disable_transformer(pipe.transformer_2)
-        pipe_cls_name = pipe.__class__.__name__
-        logger.warning(f"Cache Acceleration is disabled for: {pipe_cls_name}")
-    elif isinstance(pipe_or_adapter, BlockAdapter):
-        # BlockAdapter
-        adapter = pipe_or_adapter
-        BlockAdapter.assert_normalized(adapter)
-        _disable_pipe(adapter.pipe)
-        for transformer in BlockAdapter.flatten(adapter.transformer):
-            _disable_transformer(transformer)
-        for blocks in BlockAdapter.flatten(adapter.blocks):
-            _disable_blocks(blocks)
-        pipe_cls_name = adapter.pipe.__class__.__name__
-        logger.warning(f"Cache Acceleration is disabled for: {pipe_cls_name}")
-    else:
-        pass  # do nothing
-
 
 def supported_pipelines(
     **kwargs,

cache_dit/cache_factory/cache_types.py

@@ -22,11 +22,11 @@ def cache_type(type_hint: "CacheType | str") -> "CacheType":
     if isinstance(type_hint, CacheType):
         return type_hint
 
-    elif type_hint.lower() in (
-        "dual_block_cache",
-        "db_cache",
-        "dbcache",
-        "db",
+    elif type_hint.upper() in (
+        "DUAL_BLOCK_CACHE",
+        "DB_CACHE",
+        "DBCACHE",
+        "DB",
     ):
         return CacheType.DBCache
     return CacheType.NONE

cache_dit/cache_factory/patch_functors/functor_chroma.py

@@ -56,7 +56,8 @@ class ChromaPatchFunctor(PatchFunctor):
             transformer.forward = __patch_transformer_forward__.__get__(
                 transformer
             )
-            transformer._is_patched = True
+
+        transformer._is_patched = is_patched  # True or False
 
         cls_name = transformer.__class__.__name__
         logger.info(

cache_dit/cache_factory/patch_functors/functor_flux.py

@@ -57,7 +57,8 @@ class FluxPatchFunctor(PatchFunctor):
             transformer.forward = __patch_transformer_forward__.__get__(
                 transformer
             )
-            transformer._is_patched = True
+
+        transformer._is_patched = is_patched  # True or False
 
         cls_name = transformer.__class__.__name__
         logger.info(

cache_dit/utils.py

@@ -52,6 +52,9 @@ def summary(
             if hasattr(adapter_or_others, "transformer_2"):
                 transformer_2 = adapter_or_others.transformer_2
 
+        if not BlockAdapter.is_cached(transformer):
+            return [CacheStats()]
+
         blocks_stats: List[CacheStats] = []
         for blocks in BlockAdapter.find_blocks(transformer):
             blocks_stats.append(
@@ -212,7 +215,8 @@ def _summary(
         if logging:
             print(f"\n🤗Cache Options: {cls_name}\n\n{cache_options}")
     else:
-        logger.warning(f"Can't find Cache Options for: {cls_name}")
+        if logging:
+            logger.warning(f"Can't find Cache Options for: {cls_name}")
 
     if hasattr(module, "_cached_steps"):
         cached_steps: list[int] = module._cached_steps

{cache_dit-0.2.29.dist-info → cache_dit-0.2.30.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.2.29
+Version: 0.2.30
 Summary: 🤗 A Unified and Training-free Cache Acceleration Toolbox for Diffusion Transformers
 Author: DefTruth, vipshop.com, etc.
 Maintainer: DefTruth, vipshop.com, etc
@@ -59,7 +59,7 @@ Dynamic: requires-python
   🔥<b><a href="#unified">Unified Cache APIs</a> | <a href="#dbcache">DBCache</a> | <a href="#taylorseer">Hybrid TaylorSeer</a> | <a href="#cfg">Hybrid Cache CFG</a></b>🔥
   </p>
   <p align="center">
-  🎉Now, <b>cache-dit</b> covers <b>100%</b> Diffusers' <b>DiT-based</b> Pipelines🎉<br>
+  🎉Now, <b>cache-dit</b> covers <b>mainstream</b> Diffusers' <b>DiT-based</b> Pipelines🎉<br>
   🔥<b><a href="#supported">Qwen-Image</a> | <a href="#supported">FLUX.1</a> | <a href="#supported">Wan 2.1/2.2</a> | <a href="#supported"> ... </a> | <a href="#supported">CogVideoX</a></b>🔥
   </p>
 </div>
@@ -87,7 +87,6 @@ Dynamic: requires-python
 <summary> Previous News </summary>  
 
 - [2025-09-01] 📚[**Hybird Forward Pattern**](#unified) is supported! Please check [FLUX.1-dev](./examples/run_flux_adapter.py) as an example.
-- [2025-08-29] 🔥</b>Covers <b>100%</b> Diffusers' <b>DiT-based</b> Pipelines: **[BlockAdapter](#unified) + [Pattern Matching](#unified).**  
 - [2025-08-10] 🔥[**FLUX.1-Kontext-dev**](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) is supported! Please refer [run_flux_kontext.py](./examples/pipeline/run_flux_kontext.py) as an example.
 - [2025-07-18] 🎉First caching mechanism in [🤗huggingface/flux-fast](https://github.com/huggingface/flux-fast) with **[cache-dit](https://github.com/vipshop/cache-dit)**, check the [PR](https://github.com/huggingface/flux-fast/pull/13). 
 
@@ -130,19 +129,8 @@ pip3 install git+https://github.com/vipshop/cache-dit.git
 
 <div id="supported"></div>
 
-```python
->>> import cache_dit
->>> cache_dit.supported_pipelines()
-(31, ['Flux*', 'Mochi*', 'CogVideoX*', 'Wan*', 'HunyuanVideo*', 'QwenImage*', 'LTXVideo*', 
-'Allegro*', 'CogView3Plus*', 'CogView4*', 'Cosmos*', 'EasyAnimate*', 'SkyReelsV2*', 'SD3*', 
-'ConsisID*', 'DiT*', 'Amused*', 'Bria*', 'HunyuanDiT*', 'HunyuanDiTPAG*', 'Lumina*', 'Lumina2*',
-'OmniGen*', 'PixArt*', 'Sana*', 'ShapE*', 'StableAudio*', 'VisualCloze*', 'AuraFlow*', 
-'Chroma*', 'HiDream*'])
-```
-
 Currently, **cache-dit** library supports almost **Any** Diffusion Transformers (with **Transformer Blocks** that match the specific Input and Output **patterns**). Please check [🎉Unified Cache APIs](#unified) for more details. Here are just some of the tested models listed:
 
-
 - [🚀Qwen-Image-Edit](https://github.com/vipshop/cache-dit/raw/main/examples)  
 - [🚀Qwen-Image](https://github.com/vipshop/cache-dit/raw/main/examples)  
 - [🚀FLUX.1-dev](https://github.com/vipshop/cache-dit/raw/main/examples)  
@@ -154,35 +142,7 @@ Currently, **cache-dit** library supports almost **Any** Diffusion Transformers
 - [🚀Wan2.1-T2V](https://github.com/vipshop/cache-dit/raw/main/examples)
 - [🚀Wan2.1-FLF2V](https://github.com/vipshop/cache-dit/raw/main/examples)
 - [🚀HunyuanVideo](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀HunyuanDiT](https://github.com/vipshop/cache-dit/raw/main/examples)
 
-<details>
-<summary> More Pipelines </summary>  
-
-- [🚀mochi-1-preview](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀LTXVideo](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀Allegro](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀CogView3Plus](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀CogView4](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀Cosmos](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀EasyAnimate](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀SkyReelsV2](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀SD3](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀ConsisID](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀DiT](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀Amused](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀HunyuanDiTPAG](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀Lumina](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀Lumina2](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀OmniGen](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀PixArt](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀Sana](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀StableAudio](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀VisualCloze](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀AuraFlow](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀Chroma](https://github.com/vipshop/cache-dit/raw/main/examples)
-- [🚀HiDream](https://github.com/vipshop/cache-dit/raw/main/examples)
-  
 </details>
 
 ## 🎉Unified Cache APIs

{cache_dit-0.2.29.dist-info → cache_dit-0.2.30.dist-info}/RECORD

@@ -1,30 +1,30 @@
 cache_dit/__init__.py,sha256=kX9V-FegZG4c8LMwI4PTmMqH794MEW0pzDArdhC0cJw,1241
-cache_dit/_version.py,sha256=4_NDrwSRsA8gshfOOEHYB4RwOrbBlY3Re7Srt7YQl4M,706
+cache_dit/_version.py,sha256=6uKAYeE03adIcUS0SDwp52AaQx0KO8z_-07D_lPHrz8,706
 cache_dit/logger.py,sha256=0zsu42hN-3-rgGC_C29ms1IvVpV4_b4_SwJCKSenxBE,4304
-cache_dit/utils.py,sha256=bMeZw377_mACEj3nV1tn5DTqypBsbUVvZWJYjNQQxPg,10399
+cache_dit/utils.py,sha256=WK7eqgH6gCYNHXNLmWyxBDU0XSHTPg7CfOcyXlGXBqE,10510
 cache_dit/cache_factory/.gitignore,sha256=5Cb-qT9wsTUoMJ7vACDF7ZcLpAXhi5v-xdcWSRit988,23
 cache_dit/cache_factory/__init__.py,sha256=Iw6-iJLFbdzCsIDZXXOw371L-HPmoeZO_P9a3sDjP5s,1103
-cache_dit/cache_factory/cache_adapters.py,sha256=knNzV4BbCQmyiwsybFGl3LpTDEFtenykLb9-y_bbDpA,13905
-cache_dit/cache_factory/cache_interface.py,sha256=wy9QNNNNH4ncdyGepuvyJSJLbaTRCqqlvxtO6Os20yA,11317
-cache_dit/cache_factory/cache_types.py,sha256=FIFa6ZBfvvSMMHyBBhvarvgg2Y2wbRgITcG_uGylGe0,991
+cache_dit/cache_factory/cache_adapters.py,sha256=TA_0mEHMdSQDrt4rYASeX4-BD8pJOznSJfMV3hkrGuk,17851
+cache_dit/cache_factory/cache_interface.py,sha256=y1nY6R3MucRmAnG2UJRI_tIKrRk27FktGWLbfckf3zE,8543
+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/utils.py,sha256=XkVM9AXcB9zYq8-S8QKAsGz80r3tA6U3lBNGDGeHOe4,1871
-cache_dit/cache_factory/block_adapters/__init__.py,sha256=mtYPmsAYz4MGsMmanf6xZLaxZEkgE8gwB5mYhrM4nw4,15862
-cache_dit/cache_factory/block_adapters/block_adapters.py,sha256=sRTAfhDxdj3VAm7ejyklGwM7HMimQSEgVroVkLx7CR8,20997
-cache_dit/cache_factory/block_adapters/block_registers.py,sha256=79HpWTX7PO2ynY8I-KnF6pa-ETV4Dlpbxn5wvp_iyvw,2387
+cache_dit/cache_factory/block_adapters/__init__.py,sha256=EA-4mEVy-JJ5vRDo6C3nJIOXu0ZDNc6FQ-ZLAKHDtB0,17251
+cache_dit/cache_factory/block_adapters/block_adapters.py,sha256=OrKhuNdcGBCgSsPchdf4h32Ad-bQVUXNigMhPJ4cCvk,21069
+cache_dit/cache_factory/block_adapters/block_registers.py,sha256=ZeN2wGPmuf2u3puSsBx8x-rl3wRo8-cWcuWNcrssVfA,2553
 cache_dit/cache_factory/cache_blocks/__init__.py,sha256=OWjnpJxA8EJVoRzuyb5miuiRphUFj831-bbtWsTDjnM,2750
 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=CtBr6nvtAW8SAeEwPwiwWtPgrmwyb5ukb-j3IwFULJU,9953
+cache_dit/cache_factory/cache_blocks/pattern_3_4_5.py,sha256=nf2f5wdxp6tfq9AhFyMyBeKiZfxh63WG1g8q-c2BBSg,10182
 cache_dit/cache_factory/cache_blocks/pattern_base.py,sha256=XSDy3hsaKbAZPGZY92YgGA0qLgjQyIX8irQkb2R5T2c,20331
 cache_dit/cache_factory/cache_blocks/utils.py,sha256=dGOC1tMMOvcbvEgx44eTESKn_jsv-0RZ3tRHPa3wmQ4,1315
 cache_dit/cache_factory/cache_contexts/__init__.py,sha256=rqnJ5__zqnpVHK5A1OqWILpNh5Ss-0ZDTGgtxZMKGGo,250
 cache_dit/cache_factory/cache_contexts/cache_context.py,sha256=N88WLdd4KE9DuMWmpX8URcF55E2zWNwcKMxgVYkxMJY,13691
-cache_dit/cache_factory/cache_contexts/cache_manager.py,sha256=Tgk2-VFEhUp-oe-TFHzXay_YgbU8v90_Nx2G17ZlTds,34937
+cache_dit/cache_factory/cache_contexts/cache_manager.py,sha256=_NUXcMYYEIVfDHpc4HJr9RUjU5RUEkZmAgFGE8bh5Wc,34883
 cache_dit/cache_factory/cache_contexts/taylorseer.py,sha256=etSUIZzDvqW3ScKCbccTPcFaSmxV1T-xAXdk-p3e3wk,3802
 cache_dit/cache_factory/patch_functors/__init__.py,sha256=yK05iONMGILsTZ83ynrUUJtiJKJ_FDjxmVIzRLy416s,252
 cache_dit/cache_factory/patch_functors/functor_base.py,sha256=Ahk0fTfrHgNdEl-9JSkACvfyyv9G-Ei5OSz7XBIlX5o,357
-cache_dit/cache_factory/patch_functors/functor_chroma.py,sha256=Z0kPAib0TkXGzJIP9FRK559UlBVuGQSZIVFir6tHzJM,10022
-cache_dit/cache_factory/patch_functors/functor_flux.py,sha256=w_QaUwv7l7ypvFxWHzjHjAafLr1PxQcgv5N7VFjr6N8,9531
+cache_dit/cache_factory/patch_functors/functor_chroma.py,sha256=N3UzszCM55g3GHeVdyXkid2_n72VJrfqBM2gdtD52gw,10042
+cache_dit/cache_factory/patch_functors/functor_flux.py,sha256=rJsbGEIxWPGnZyGI4ZwLLBdg8u6ZItsOeh0UoD_bVwk,9551
 cache_dit/compile/__init__.py,sha256=FcTVzCeyypl-mxlc59_ehHL3lBNiDAFsXuRoJ-5Cfi0,56
 cache_dit/compile/utils.py,sha256=nN2OIrSdwRR5zGxJinKDqb07pXpvTNTF3g_OgLkeeBU,3858
 cache_dit/custom_ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -39,9 +39,9 @@ cache_dit/quantize/__init__.py,sha256=kWYoMAyZgBXu9BJlZjTQ0dRffW9GqeeY9_iTkXrb70
 cache_dit/quantize/quantize_ao.py,sha256=mGspqYgQtenl3QnKPtsSYsSD7LbVX93f1M940bhXKLU,6066
 cache_dit/quantize/quantize_interface.py,sha256=2s_R7xPSKuJeFpEGeLwRxnq_CqJcBG3a3lzyW5wh-UM,1241
 cache_dit/quantize/quantize_svdq.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit-0.2.29.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
-cache_dit-0.2.29.dist-info/METADATA,sha256=wagdLaiAIX7fs5Rsw89DMsj5KQl1t3zGsTosKJe2AlQ,24540
-cache_dit-0.2.29.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cache_dit-0.2.29.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
-cache_dit-0.2.29.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
-cache_dit-0.2.29.dist-info/RECORD,,
+cache_dit-0.2.30.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
+cache_dit-0.2.30.dist-info/METADATA,sha256=8Ln_X5fw14U3greCM7cSukrei1SRiMDpksFalg5ZBAU,22130
+cache_dit-0.2.30.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-0.2.30.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
+cache_dit-0.2.30.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-0.2.30.dist-info/RECORD,,