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

cache_dit/cache_factory/block_adapters/__init__.py

@@ -153,7 +153,7 @@ def qwenimage_adapter(pipe, **kwargs) -> BlockAdapter:
     )
 
 
-@BlockAdapterRegistry.register("LTXVideo")
+@BlockAdapterRegistry.register("LTX")
 def ltxvideo_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import LTXVideoTransformer3DModel
 
@@ -248,7 +248,10 @@ def skyreelsv2_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.blocks,
-        forward_pattern=ForwardPattern.Pattern_2,
+        # NOTE: Use Pattern_3 instead of Pattern_2 because the
+        # encoder_hidden_states will never change in the blocks
+        # forward loop.
+        forward_pattern=ForwardPattern.Pattern_3,
         has_separate_cfg=True,
         **kwargs,
     )
@@ -285,6 +288,7 @@ def consisid_adapter(pipe, **kwargs) -> BlockAdapter:
 @BlockAdapterRegistry.register("DiT")
 def dit_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import DiTTransformer2DModel
+    from cache_dit.cache_factory.patch_functors import DiTPatchFunctor
 
     assert isinstance(pipe.transformer, DiTTransformer2DModel)
     return BlockAdapter(
@@ -292,6 +296,7 @@ def dit_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_3,
+        patch_functor=DiTPatchFunctor(),
         **kwargs,
     )
 
@@ -331,24 +336,13 @@ def bria_adapter(pipe, **kwargs) -> BlockAdapter:
 
 
 @BlockAdapterRegistry.register("Lumina")
-def lumina_adapter(pipe, **kwargs) -> BlockAdapter:
-    from diffusers import LuminaNextDiT2DModel
-
-    assert isinstance(pipe.transformer, LuminaNextDiT2DModel)
-    return BlockAdapter(
-        pipe=pipe,
-        transformer=pipe.transformer,
-        blocks=pipe.transformer.layers,
-        forward_pattern=ForwardPattern.Pattern_3,
-        **kwargs,
-    )
-
-
-@BlockAdapterRegistry.register("Lumina2")
 def lumina2_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import Lumina2Transformer2DModel
+    from diffusers import LuminaNextDiT2DModel
 
-    assert isinstance(pipe.transformer, Lumina2Transformer2DModel)
+    assert isinstance(
+        pipe.transformer, (Lumina2Transformer2DModel, LuminaNextDiT2DModel)
+    )
     return BlockAdapter(
         pipe=pipe,
         transformer=pipe.transformer,
@@ -386,12 +380,10 @@ def pixart_adapter(pipe, **kwargs) -> BlockAdapter:
     )
 
 
-@BlockAdapterRegistry.register("Sana", supported=False)
+@BlockAdapterRegistry.register("Sana")
 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,
@@ -469,6 +461,7 @@ def auraflow_adapter(pipe, **kwargs) -> BlockAdapter:
 @BlockAdapterRegistry.register("Chroma")
 def chroma_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import ChromaTransformer2DModel
+    from cache_dit.cache_factory.patch_functors import ChromaPatchFunctor
 
     assert isinstance(pipe.transformer, ChromaTransformer2DModel)
     return BlockAdapter(
@@ -482,6 +475,7 @@ def chroma_adapter(pipe, **kwargs) -> BlockAdapter:
             ForwardPattern.Pattern_1,
             ForwardPattern.Pattern_3,
         ],
+        patch_functor=ChromaPatchFunctor(),
         has_separate_cfg=True,
         **kwargs,
     )

cache_dit/cache_factory/block_adapters/block_adapters.py

@@ -16,8 +16,52 @@ logger = init_logger(__name__)
 
 
 class ParamsModifier:
-    def __init__(self, **kwargs):
-        self._context_kwargs = kwargs.copy()
+    def __init__(
+        self,
+        # Cache context kwargs
+        Fn_compute_blocks: Optional[int] = None,
+        Bn_compute_blocks: Optional[int] = None,
+        max_warmup_steps: Optional[int] = None,
+        max_cached_steps: Optional[int] = None,
+        max_continuous_cached_steps: Optional[int] = None,
+        residual_diff_threshold: Optional[float] = None,
+        # Cache CFG or not
+        enable_separate_cfg: Optional[bool] = None,
+        cfg_compute_first: Optional[bool] = None,
+        cfg_diff_compute_separate: Optional[bool] = None,
+        # Hybird TaylorSeer
+        enable_taylorseer: Optional[bool] = None,
+        enable_encoder_taylorseer: Optional[bool] = None,
+        taylorseer_cache_type: Optional[str] = None,
+        taylorseer_order: Optional[int] = None,
+        **other_cache_context_kwargs,
+    ):
+        self._context_kwargs = other_cache_context_kwargs.copy()
+        self._maybe_update_param("Fn_compute_blocks", Fn_compute_blocks)
+        self._maybe_update_param("Bn_compute_blocks", Bn_compute_blocks)
+        self._maybe_update_param("max_warmup_steps", max_warmup_steps)
+        self._maybe_update_param("max_cached_steps", max_cached_steps)
+        self._maybe_update_param(
+            "max_continuous_cached_steps", max_continuous_cached_steps
+        )
+        self._maybe_update_param(
+            "residual_diff_threshold", residual_diff_threshold
+        )
+        self._maybe_update_param("enable_separate_cfg", enable_separate_cfg)
+        self._maybe_update_param("cfg_compute_first", cfg_compute_first)
+        self._maybe_update_param(
+            "cfg_diff_compute_separate", cfg_diff_compute_separate
+        )
+        self._maybe_update_param("enable_taylorseer", enable_taylorseer)
+        self._maybe_update_param(
+            "enable_encoder_taylorseer", enable_encoder_taylorseer
+        )
+        self._maybe_update_param("taylorseer_cache_type", taylorseer_cache_type)
+        self._maybe_update_param("taylorseer_order", taylorseer_order)
+
+    def _maybe_update_param(self, key: str, value: Any):
+        if value is not None:
+            self._context_kwargs[key] = value
 
 
 @dataclasses.dataclass

cache_dit/cache_factory/block_adapters/block_registers.py

@@ -10,13 +10,14 @@ logger = init_logger(__name__)
 
 class BlockAdapterRegistry:
     _adapters: Dict[str, Callable[..., BlockAdapter]] = {}
-    _predefined_adapters_has_spearate_cfg: List[str] = [
+    _predefined_adapters_has_separate_cfg: List[str] = [
         "QwenImage",
         "Wan",
         "CogView4",
         "Cosmos",
         "SkyReelsV2",
         "Chroma",
+        "Lumina2",
     ]
 
     @classmethod
@@ -68,7 +69,7 @@ class BlockAdapterRegistry:
             return True
 
         pipe_cls_name = pipe_or_adapter.__class__.__name__
-        for name in cls._predefined_adapters_has_spearate_cfg:
+        for name in cls._predefined_adapters_has_separate_cfg:
             if pipe_cls_name.startswith(name):
                 return True
 

cache_dit/cache_factory/cache_adapters.py

@@ -114,27 +114,27 @@ class CachedAdapter:
         **cache_context_kwargs,
     ):
         # Check cache_context_kwargs
-        if cache_context_kwargs["enable_spearate_cfg"] is None:
+        if cache_context_kwargs["enable_separate_cfg"] is None:
             # Check cfg for some specific case if users don't set it as True
             if BlockAdapterRegistry.has_separate_cfg(block_adapter):
-                cache_context_kwargs["enable_spearate_cfg"] = True
+                cache_context_kwargs["enable_separate_cfg"] = True
                 logger.info(
-                    f"Use custom 'enable_spearate_cfg' from BlockAdapter: True. "
+                    f"Use custom 'enable_separate_cfg' from BlockAdapter: True. "
                     f"Pipeline: {block_adapter.pipe.__class__.__name__}."
                 )
             else:
-                cache_context_kwargs["enable_spearate_cfg"] = (
+                cache_context_kwargs["enable_separate_cfg"] = (
                     BlockAdapterRegistry.has_separate_cfg(block_adapter.pipe)
                 )
                 logger.info(
-                    f"Use default 'enable_spearate_cfg' from block adapter "
-                    f"register: {cache_context_kwargs['enable_spearate_cfg']}, "
+                    f"Use default 'enable_separate_cfg' from block adapter "
+                    f"register: {cache_context_kwargs['enable_separate_cfg']}, "
                     f"Pipeline: {block_adapter.pipe.__class__.__name__}."
                 )
         else:
             logger.info(
-                f"Use custom 'enable_spearate_cfg' from cache context "
-                f"kwargs: {cache_context_kwargs['enable_spearate_cfg']}. "
+                f"Use custom 'enable_separate_cfg' from cache context "
+                f"kwargs: {cache_context_kwargs['enable_separate_cfg']}. "
                 f"Pipeline: {block_adapter.pipe.__class__.__name__}."
             )
 

cache_dit/cache_factory/cache_contexts/cache_context.py

@@ -53,20 +53,20 @@ class CachedContext:  # Internal CachedContext Impl class
     enable_taylorseer: bool = False
     enable_encoder_taylorseer: bool = False
     taylorseer_cache_type: str = "hidden_states"  # residual or hidden_states
-    taylorseer_order: int = 2  # The order for TaylorSeer
+    taylorseer_order: int = 1  # The order for TaylorSeer
     taylorseer_kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict)
     taylorseer: Optional[TaylorSeer] = None
     encoder_tarlorseer: Optional[TaylorSeer] = None
 
-    # Support enable_spearate_cfg, such as Wan 2.1,
+    # Support enable_separate_cfg, such as Wan 2.1,
     # Qwen-Image. For model that fused CFG and non-CFG into single
-    # forward step, should set enable_spearate_cfg as False.
+    # forward step, should set enable_separate_cfg as False.
     # For example: CogVideoX, HunyuanVideo, Mochi.
-    enable_spearate_cfg: bool = False
+    enable_separate_cfg: bool = False
     # Compute cfg forward first or not, default False, namely,
     # 0, 2, 4, ..., -> non-CFG step; 1, 3, 5, ... -> CFG step.
     cfg_compute_first: bool = False
-    # Compute spearate diff values for CFG and non-CFG step,
+    # Compute separate diff values for CFG and non-CFG step,
     # default True. If False, we will use the computed diff from
     # current non-CFG transformer step for current CFG step.
     cfg_diff_compute_separate: bool = True
@@ -89,7 +89,7 @@ class CachedContext:  # Internal CachedContext Impl class
         if logger.isEnabledFor(logging.DEBUG):
             logger.info(f"Created _CacheContext: {self.name}")
         # Some checks for settings
-        if self.enable_spearate_cfg:
+        if self.enable_separate_cfg:
             if self.cfg_diff_compute_separate:
                 assert self.cfg_compute_first is False, (
                     "cfg_compute_first must set as False if "
@@ -108,12 +108,12 @@ class CachedContext:  # Internal CachedContext Impl class
 
         if self.enable_taylorseer:
             self.taylorseer = TaylorSeer(**self.taylorseer_kwargs)
-            if self.enable_spearate_cfg:
+            if self.enable_separate_cfg:
                 self.cfg_taylorseer = TaylorSeer(**self.taylorseer_kwargs)
 
         if self.enable_encoder_taylorseer:
             self.encoder_tarlorseer = TaylorSeer(**self.taylorseer_kwargs)
-            if self.enable_spearate_cfg:
+            if self.enable_separate_cfg:
                 self.cfg_encoder_taylorseer = TaylorSeer(
                     **self.taylorseer_kwargs
                 )
@@ -145,7 +145,7 @@ class CachedContext:  # Internal CachedContext Impl class
         # incr     step: prev 0 -> 1; prev 1 -> 2
         # current  step: incr step - 1
         self.transformer_executed_steps += 1
-        if not self.enable_spearate_cfg:
+        if not self.enable_separate_cfg:
             self.executed_steps += 1
         else:
             # 0,1 -> 0 + 1, 2,3 -> 1 + 1, ...
@@ -183,7 +183,7 @@ class CachedContext:  # Internal CachedContext Impl class
 
         # mark_step_begin of TaylorSeer must be called after the cache is reset.
         if self.enable_taylorseer or self.enable_encoder_taylorseer:
-            if self.enable_spearate_cfg:
+            if self.enable_separate_cfg:
                 # Assume non-CFG steps: 0, 2, 4, 6, ...
                 if not self.is_separate_cfg_step():
                     taylorseer, encoder_taylorseer = self.get_taylorseers()
@@ -269,7 +269,7 @@ class CachedContext:  # Internal CachedContext Impl class
         return self.transformer_executed_steps - 1
 
     def is_separate_cfg_step(self):
-        if not self.enable_spearate_cfg:
+        if not self.enable_separate_cfg:
             return False
         if self.cfg_compute_first:
             # CFG steps: 0, 2, 4, 6, ...

cache_dit/cache_factory/cache_contexts/cache_manager.py

@@ -74,8 +74,8 @@ class CachedContextManager:
                 del self._cached_context_manager[cached_context]
 
     def clear_contexts(self):
-        for cached_context in self._cached_context_manager:
-            self.remove_context(cached_context)
+        for context_name in list(self._cached_context_manager.keys()):
+            self.remove_context(context_name)
 
     @contextlib.contextmanager
     def enter_context(self, cached_context: CachedContext | str):
@@ -364,10 +364,10 @@ class CachedContextManager:
         return cached_context.Bn_compute_blocks
 
     @torch.compiler.disable
-    def enable_spearate_cfg(self) -> bool:
+    def enable_separate_cfg(self) -> bool:
         cached_context = self.get_context()
         assert cached_context is not None, "cached_context must be set before"
-        return cached_context.enable_spearate_cfg
+        return cached_context.enable_separate_cfg
 
     @torch.compiler.disable
     def is_separate_cfg_step(self) -> bool:
@@ -410,7 +410,7 @@ class CachedContextManager:
 
         if all(
             (
-                self.enable_spearate_cfg(),
+                self.enable_separate_cfg(),
                 self.is_separate_cfg_step(),
                 not self.cfg_diff_compute_separate(),
                 self.get_current_step_residual_diff() is not None,

cache_dit/cache_factory/cache_contexts/taylorseer.py

@@ -1,4 +1,6 @@
 import math
+import torch
+from typing import List, Dict
 
 
 class TaylorSeer:
@@ -17,7 +19,7 @@ class TaylorSeer:
         self.reset_cache()
 
     def reset_cache(self):
-        self.state = {
+        self.state: Dict[str, List[torch.Tensor]] = {
             "dY_prev": [None] * self.ORDER,
             "dY_current": [None] * self.ORDER,
         }
@@ -36,15 +38,19 @@ class TaylorSeer:
             return True
         return False
 
-    def approximate_derivative(self, Y):
+    def approximate_derivative(self, Y: torch.Tensor) -> List[torch.Tensor]:
         # n-th order Taylor expansion:
         # Y(t) = Y(0) + dY(0)/dt * t + d^2Y(0)/dt^2 * t^2 / 2!
         #        + ... + d^nY(0)/dt^n * t^n / n!
         # TODO: Custom Triton/CUDA kernel for better performance,
         # especially for large n_derivatives.
-        dY_current = [None] * self.ORDER
+        dY_current: List[torch.Tensor] = [None] * self.ORDER
         dY_current[0] = Y
         window = self.current_step - self.last_non_approximated_step
+        if self.state["dY_prev"][0] is not None:
+            if dY_current[0].shape != self.state["dY_prev"][0].shape:
+                self.reset_cache()
+
         for i in range(self.n_derivatives):
             if self.state["dY_prev"][i] is not None and self.current_step > 1:
                 dY_current[i + 1] = (
@@ -54,7 +60,7 @@ class TaylorSeer:
                 break
         return dY_current
 
-    def approximate_value(self):
+    def approximate_value(self) -> torch.Tensor:
         # TODO: Custom Triton/CUDA kernel for better performance,
         # especially for large n_derivatives.
         elapsed = self.current_step - self.last_non_approximated_step
@@ -69,7 +75,7 @@ class TaylorSeer:
     def mark_step_begin(self):
         self.current_step += 1
 
-    def update(self, Y):
+    def update(self, Y: torch.Tensor):
         # Directly call this method will ingnore the warmup
         # policy and force full computation.
         # Assume warmup steps is 3, and n_derivatives is 3.
@@ -87,7 +93,7 @@ class TaylorSeer:
         self.state["dY_current"] = self.approximate_derivative(Y)
         self.last_non_approximated_step = self.current_step
 
-    def step(self, Y):
+    def step(self, Y: torch.Tensor):
         self.mark_step_begin()
         if self.should_compute_full():
             self.update(Y)

cache_dit/cache_factory/cache_interface.py

@@ -24,14 +24,14 @@ def enable_cache(
     max_continuous_cached_steps: int = -1,
     residual_diff_threshold: float = 0.08,
     # Cache CFG or not
-    enable_spearate_cfg: bool | None = None,
+    enable_separate_cfg: bool = None,
     cfg_compute_first: bool = False,
     cfg_diff_compute_separate: bool = True,
     # Hybird TaylorSeer
     enable_taylorseer: bool = False,
     enable_encoder_taylorseer: bool = False,
     taylorseer_cache_type: str = "residual",
-    taylorseer_order: int = 2,
+    taylorseer_order: int = 1,
     **other_cache_context_kwargs,
 ) -> Union[
     DiffusionPipeline,
@@ -70,15 +70,15 @@ def enable_cache(
         residual_diff_threshold (`float`, *required*, defaults to 0.08):
             he value of residual diff threshold, a higher value leads to faster performance at the
             cost of lower precision.
-        enable_spearate_cfg (`bool`, *required*,  defaults to None):
+        enable_separate_cfg (`bool`, *required*,  defaults to None):
             Whether to do separate cfg or not, such as Wan 2.1, Qwen-Image. For model that fused CFG
-            and non-CFG into single forward step, should set enable_spearate_cfg as False, for example:
+            and non-CFG into single forward step, should set enable_separate_cfg as False, for example:
             CogVideoX, HunyuanVideo, Mochi, etc.
         cfg_compute_first (`bool`, *required*,  defaults to False):
             Compute cfg forward first or not, default False, namely, 0, 2, 4, ..., -> non-CFG step;
             1, 3, 5, ... -> CFG step.
         cfg_diff_compute_separate (`bool`, *required*,  defaults to True):
-            Compute spearate diff values for CFG and non-CFG step, default True. If False, we will
+            Compute separate diff values for CFG and non-CFG step, default True. If False, we will
             use the computed diff from current non-CFG transformer step for current CFG step.
         enable_taylorseer (`bool`, *required*,  defaults to False):
             Enable the hybird TaylorSeer for hidden_states or not. We have supported the
@@ -91,10 +91,10 @@ def enable_cache(
             Enable the hybird TaylorSeer for encoder_hidden_states or not.
         taylorseer_cache_type (`str`, *required*,  defaults to `residual`):
             The TaylorSeer implemented in cache-dit supports both `hidden_states` and `residual` as cache type.
-        taylorseer_order (`int`, *required*, defaults to 2):
+        taylorseer_order (`int`, *required*, defaults to 1):
             The order of taylorseer, higher values of n_derivatives will lead to longer computation time,
-            but may improve precision significantly.
-        other_cache_kwargs: (`dict`, *optional*, defaults to {})
+            the recommended value is 1 or 2.
+        other_cache_context_kwargs: (`dict`, *optional*, defaults to {})
             Other cache context kwargs, please check https://github.com/vipshop/cache-dit/blob/main/src/cache_dit/cache_factory/cache_contexts/cache_context.py
             for more details.
 
@@ -123,7 +123,7 @@ def enable_cache(
         max_continuous_cached_steps
     )
     cache_context_kwargs["residual_diff_threshold"] = residual_diff_threshold
-    cache_context_kwargs["enable_spearate_cfg"] = enable_spearate_cfg
+    cache_context_kwargs["enable_separate_cfg"] = enable_separate_cfg
     cache_context_kwargs["cfg_compute_first"] = cfg_compute_first
     cache_context_kwargs["cfg_diff_compute_separate"] = (
         cfg_diff_compute_separate

cache_dit/cache_factory/patch_functors/__init__.py

@@ -1,4 +1,5 @@
 from cache_dit.cache_factory.patch_functors.functor_base import PatchFunctor
+from cache_dit.cache_factory.patch_functors.functor_dit import DiTPatchFunctor
 from cache_dit.cache_factory.patch_functors.functor_flux import FluxPatchFunctor
 from cache_dit.cache_factory.patch_functors.functor_chroma import (
     ChromaPatchFunctor,

cache_dit/cache_factory/patch_functors/functor_chroma.py

@@ -1,10 +1,9 @@
-import inspect
-
 import torch
 import numpy as np
 from typing import Tuple, Optional, Dict, Any, Union
 from diffusers import ChromaTransformer2DModel
 from diffusers.models.transformers.transformer_chroma import (
+    ChromaTransformerBlock,
     ChromaSingleTransformerBlock,
     Transformer2DModelOutput,
 )
@@ -27,24 +26,31 @@ class ChromaPatchFunctor(PatchFunctor):
     def apply(
         self,
         transformer: ChromaTransformer2DModel,
-        blocks: torch.nn.ModuleList = None,
         **kwargs,
     ) -> ChromaTransformer2DModel:
         if hasattr(transformer, "_is_patched"):
             return transformer
 
-        if blocks is None:
-            blocks = transformer.single_transformer_blocks
-
         is_patched = False
-        for block in blocks:
-            if isinstance(block, ChromaSingleTransformerBlock):
-                forward_parameters = inspect.signature(
-                    block.forward
-                ).parameters.keys()
-                if "encoder_hidden_states" not in forward_parameters:
-                    block.forward = __patch_single_forward__.__get__(block)
-                    is_patched = True
+        for index_block, block in enumerate(transformer.transformer_blocks):
+            assert isinstance(block, ChromaTransformerBlock)
+            img_offset = 3 * len(transformer.single_transformer_blocks)
+            txt_offset = img_offset + 6 * len(transformer.transformer_blocks)
+            img_modulation = img_offset + 6 * index_block
+            text_modulation = txt_offset + 6 * index_block
+            block._img_modulation = img_modulation
+            block._text_modulation = text_modulation
+            block.forward = __patch_double_forward__.__get__(block)
+
+        for index_block, block in enumerate(
+            transformer.single_transformer_blocks
+        ):
+            assert isinstance(block, ChromaSingleTransformerBlock)
+            start_idx = 3 * index_block
+            block._start_idx = start_idx
+            block.forward = __patch_single_forward__.__get__(block)
+
+        is_patched = True
 
         cls_name = transformer.__class__.__name__
 
@@ -69,25 +75,123 @@ class ChromaPatchFunctor(PatchFunctor):
         return transformer
 
 
+# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_chroma.py
+def __patch_double_forward__(
+    self: ChromaTransformerBlock,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    pooled_temb: torch.Tensor,
+    image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    # TODO: Fuse controlnet into block forward
+    img_modulation = self._img_modulation
+    text_modulation = self._text_modulation
+    temb = torch.cat(
+        (
+            pooled_temb[:, img_modulation : img_modulation + 6],
+            pooled_temb[:, text_modulation : text_modulation + 6],
+        ),
+        dim=1,
+    )
+
+    temb_img, temb_txt = temb[:, :6], temb[:, 6:]
+    norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+        hidden_states, emb=temb_img
+    )
+
+    (
+        norm_encoder_hidden_states,
+        c_gate_msa,
+        c_shift_mlp,
+        c_scale_mlp,
+        c_gate_mlp,
+    ) = self.norm1_context(encoder_hidden_states, emb=temb_txt)
+    joint_attention_kwargs = joint_attention_kwargs or {}
+    if attention_mask is not None:
+        attention_mask = (
+            attention_mask[:, None, None, :] * attention_mask[:, None, :, None]
+        )
+
+    # Attention.
+    attention_outputs = self.attn(
+        hidden_states=norm_hidden_states,
+        encoder_hidden_states=norm_encoder_hidden_states,
+        image_rotary_emb=image_rotary_emb,
+        attention_mask=attention_mask,
+        **joint_attention_kwargs,
+    )
+
+    if len(attention_outputs) == 2:
+        attn_output, context_attn_output = attention_outputs
+    elif len(attention_outputs) == 3:
+        attn_output, context_attn_output, ip_attn_output = attention_outputs
+
+    # Process attention outputs for the `hidden_states`.
+    attn_output = gate_msa.unsqueeze(1) * attn_output
+    hidden_states = hidden_states + attn_output
+
+    norm_hidden_states = self.norm2(hidden_states)
+    norm_hidden_states = (
+        norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+    )
+
+    ff_output = self.ff(norm_hidden_states)
+    ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+    hidden_states = hidden_states + ff_output
+    if len(attention_outputs) == 3:
+        hidden_states = hidden_states + ip_attn_output
+
+    # Process attention outputs for the `encoder_hidden_states`.
+
+    context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+    encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+    norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+    norm_encoder_hidden_states = (
+        norm_encoder_hidden_states * (1 + c_scale_mlp[:, None])
+        + c_shift_mlp[:, None]
+    )
+
+    context_ff_output = self.ff_context(norm_encoder_hidden_states)
+    encoder_hidden_states = (
+        encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
+    )
+    if encoder_hidden_states.dtype == torch.float16:
+        encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+    return encoder_hidden_states, hidden_states
+
+
 # adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_chroma.py
 def __patch_single_forward__(
     self: ChromaSingleTransformerBlock,  # Almost same as FluxSingleTransformerBlock
     hidden_states: torch.Tensor,
-    encoder_hidden_states: torch.Tensor,
-    temb: torch.Tensor,
+    pooled_temb: torch.Tensor,
     image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    attention_mask: Optional[torch.Tensor] = None,
     joint_attention_kwargs: Optional[Dict[str, Any]] = None,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    text_seq_len = encoder_hidden_states.shape[1]
-    hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+) -> torch.Tensor:
+    # TODO: Fuse controlnet into block forward
+    start_idx = self._start_idx
+    temb = pooled_temb[:, start_idx : start_idx + 3]
 
     residual = hidden_states
     norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
     mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
     joint_attention_kwargs = joint_attention_kwargs or {}
+
+    if attention_mask is not None:
+        attention_mask = (
+            attention_mask[:, None, None, :] * attention_mask[:, None, :, None]
+        )
+
     attn_output = self.attn(
         hidden_states=norm_hidden_states,
         image_rotary_emb=image_rotary_emb,
+        attention_mask=attention_mask,
         **joint_attention_kwargs,
     )
 
@@ -98,11 +202,7 @@ def __patch_single_forward__(
     if hidden_states.dtype == torch.float16:
         hidden_states = hidden_states.clip(-65504, 65504)
 
-    encoder_hidden_states, hidden_states = (
-        hidden_states[:, :text_seq_len],
-        hidden_states[:, text_seq_len:],
-    )
-    return encoder_hidden_states, hidden_states
+    return hidden_states
 
 
 # Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_chroma.py
@@ -174,24 +274,13 @@ def __patch_transformer_forward__(
         joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})
 
     for index_block, block in enumerate(self.transformer_blocks):
-        img_offset = 3 * len(self.single_transformer_blocks)
-        txt_offset = img_offset + 6 * len(self.transformer_blocks)
-        img_modulation = img_offset + 6 * index_block
-        text_modulation = txt_offset + 6 * index_block
-        temb = torch.cat(
-            (
-                pooled_temb[:, img_modulation : img_modulation + 6],
-                pooled_temb[:, text_modulation : text_modulation + 6],
-            ),
-            dim=1,
-        )
         if torch.is_grad_enabled() and self.gradient_checkpointing:
             encoder_hidden_states, hidden_states = (
                 self._gradient_checkpointing_func(
                     block,
                     hidden_states,
                     encoder_hidden_states,
-                    temb,
+                    pooled_temb,
                     image_rotary_emb,
                     attention_mask,
                 )
@@ -201,12 +290,13 @@ def __patch_transformer_forward__(
             encoder_hidden_states, hidden_states = block(
                 hidden_states=hidden_states,
                 encoder_hidden_states=encoder_hidden_states,
-                temb=temb,
+                pooled_temb=pooled_temb,
                 image_rotary_emb=image_rotary_emb,
                 attention_mask=attention_mask,
                 joint_attention_kwargs=joint_attention_kwargs,
             )
 
+        # TODO: Fuse controlnet into block forward
         # controlnet residual
         if controlnet_block_samples is not None:
             interval_control = len(self.transformer_blocks) / len(
@@ -227,43 +317,43 @@ def __patch_transformer_forward__(
                     + controlnet_block_samples[index_block // interval_control]
                 )
 
+    hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
     for index_block, block in enumerate(self.single_transformer_blocks):
-        start_idx = 3 * index_block
-        temb = pooled_temb[:, start_idx : start_idx + 3]
         if torch.is_grad_enabled() and self.gradient_checkpointing:
-            encoder_hidden_states, hidden_states = (
-                self._gradient_checkpointing_func(
-                    block,
-                    hidden_states,
-                    encoder_hidden_states,
-                    temb,
-                    image_rotary_emb,
-                )
+            hidden_states = self._gradient_checkpointing_func(
+                block,
+                hidden_states,
+                pooled_temb,
+                image_rotary_emb,
+                attention_mask,
+                joint_attention_kwargs,
             )
 
         else:
-            encoder_hidden_states, hidden_states = block(
+            hidden_states = block(
                 hidden_states=hidden_states,
-                encoder_hidden_states=encoder_hidden_states,
-                temb=temb,
+                pooled_temb=pooled_temb,
                 image_rotary_emb=image_rotary_emb,
                 attention_mask=attention_mask,
                 joint_attention_kwargs=joint_attention_kwargs,
             )
 
+        # TODO: Fuse controlnet into block forward
         # controlnet residual
         if controlnet_single_block_samples is not None:
             interval_control = len(self.single_transformer_blocks) / len(
                 controlnet_single_block_samples
             )
             interval_control = int(np.ceil(interval_control))
-            hidden_states = (
-                hidden_states
+            hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                hidden_states[:, encoder_hidden_states.shape[1] :, ...]
                 + controlnet_single_block_samples[
                     index_block // interval_control
                 ]
             )
 
+    hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
     temb = pooled_temb[:, -2:]
     hidden_states = self.norm_out(hidden_states, temb)
     output = self.proj_out(hidden_states)

cache_dit/cache_factory/patch_functors/functor_dit.py

@@ -0,0 +1,130 @@
+import torch
+import torch.nn.functional as F
+
+from typing import Optional, Dict, Any
+from diffusers.models.transformers.dit_transformer_2d import (
+    DiTTransformer2DModel,
+    Transformer2DModelOutput,
+)
+from cache_dit.cache_factory.patch_functors.functor_base import (
+    PatchFunctor,
+)
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class DiTPatchFunctor(PatchFunctor):
+
+    def apply(
+        self,
+        transformer: DiTTransformer2DModel,
+        **kwargs,
+    ) -> DiTTransformer2DModel:
+        if hasattr(transformer, "_is_patched"):
+            return transformer
+
+        is_patched = False
+
+        transformer._norm1_emb = transformer.transformer_blocks[0].norm1.emb
+
+        is_patched = True
+
+        cls_name = transformer.__class__.__name__
+
+        if is_patched:
+            logger.warning(f"Patched {cls_name} for cache-dit.")
+            assert not getattr(transformer, "_is_parallelized", False), (
+                "Please call `cache_dit.enable_cache` before Parallelize, "
+                "the __patch_transformer_forward__ will overwrite the "
+                "parallized forward and cause a downgrade of performance."
+            )
+            transformer.forward = __patch_transformer_forward__.__get__(
+                transformer
+            )
+
+        transformer._is_patched = is_patched  # True or False
+
+        logger.info(
+            f"Applied {self.__class__.__name__} for {cls_name}, "
+            f"Patch: {is_patched}."
+        )
+
+        return transformer
+
+
+def __patch_transformer_forward__(
+    self: DiTTransformer2DModel,
+    hidden_states: torch.Tensor,
+    timestep: Optional[torch.LongTensor] = None,
+    class_labels: Optional[torch.LongTensor] = None,
+    cross_attention_kwargs: Dict[str, Any] = None,
+    return_dict: bool = True,
+):
+    height, width = (
+        hidden_states.shape[-2] // self.patch_size,
+        hidden_states.shape[-1] // self.patch_size,
+    )
+    hidden_states = self.pos_embed(hidden_states)
+
+    # 2. Blocks
+    for block in self.transformer_blocks:
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            hidden_states = self._gradient_checkpointing_func(
+                block,
+                hidden_states,
+                None,
+                None,
+                None,
+                timestep,
+                cross_attention_kwargs,
+                class_labels,
+            )
+        else:
+            hidden_states = block(
+                hidden_states,
+                attention_mask=None,
+                encoder_hidden_states=None,
+                encoder_attention_mask=None,
+                timestep=timestep,
+                cross_attention_kwargs=cross_attention_kwargs,
+                class_labels=class_labels,
+            )
+
+    # 3. Output
+    # conditioning = self.transformer_blocks[0].norm1.emb(timestep, class_labels, hidden_dtype=hidden_states.dtype)
+    conditioning = self._norm1_emb(
+        timestep, class_labels, hidden_dtype=hidden_states.dtype
+    )
+    shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
+    hidden_states = (
+        self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
+    )
+    hidden_states = self.proj_out_2(hidden_states)
+
+    # unpatchify
+    height = width = int(hidden_states.shape[1] ** 0.5)
+    hidden_states = hidden_states.reshape(
+        shape=(
+            -1,
+            height,
+            width,
+            self.patch_size,
+            self.patch_size,
+            self.out_channels,
+        )
+    )
+    hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
+    output = hidden_states.reshape(
+        shape=(
+            -1,
+            self.out_channels,
+            height * self.patch_size,
+            width * self.patch_size,
+        )
+    )
+
+    if not return_dict:
+        return (output,)
+
+    return Transformer2DModelOutput(sample=output)

{cache_dit-0.2.33.dist-info → cache_dit-0.2.34.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.2.33
+Version: 0.2.34
 Summary: 🤗 A Unified and Training-free Cache Acceleration Toolbox for Diffusion Transformers
 Author: DefTruth, vipshop.com, etc.
 Maintainer: DefTruth, vipshop.com, etc
@@ -60,40 +60,97 @@ Dynamic: requires-python
   </p>
   <p align="center">
     🎉Now, <b>cache-dit</b> covers <b>most</b> mainstream Diffusers' <b>DiT</b> Pipelines🎉<br>
-    🔥<a href="#supported">Qwen-Image</a> | <a href="#supported">FLUX.1</a> | <a href="#supported">Qwen-Image-Lightning</a> | <a href="#supported"> Wan 2.1/2.2 </a>🔥<br>
-    🔥<a href="#supported">HunyuanImage-2.1</a> | <a href="#supported">HunyuanVideo</a> | <a href="#supported">HunyuanDiT</a> | <a href="#supported">HiDream</a> | <a href="#supported">Mochi</a>🔥<br>
-    🔥<a href="#supported">CogView3Plus</a> | <a href="#supported">CogView4</a> | <a href="#supported">Chroma</a> | <a href="#supported"> LTXVideo </a> | <a href="#supported">CogVideoX 1/1.5</a>🔥<br>
-    🔥<a href="#supported">Cosmos</a> | <a href="#supported">SkyReelsV2</a> | <a href="#supported">VisualCloze</a> | <a href="#supported"> OmniGen </a> | <a href="#supported">Lumina 1/2</a>🔥<br>
-    🔥<a href="#supported">Allegro</a> | <a href="#supported">EasyAnimate</a> | <a href="#supported">SD 3/3.5</a> | <a href="#supported"> ... </a> | <a href="#supported">PixArt</a>🔥
+    🔥<a href="#supported">Qwen-Image</a> | <a href="#supported">FLUX.1</a> | <a href="#supported">Qwen-Image-Lightning</a> | <a href="#supported"> Wan 2.1 </a> | <a href="#supported"> Wan 2.2 </a>🔥<br>
+    🔥<a href="#supported">HunyuanImage-2.1</a> | <a href="#supported">HunyuanVideo</a> | <a href="#supported">HunyuanDiT</a> | <a href="#supported">HiDream</a> | <a href="#supported">AuraFlow</a>🔥<br>
+    🔥<a href="#supported">CogView3Plus</a> | <a href="#supported">CogView4</a> | <a href="#supported">LTXVideo</a> | <a href="#supported">CogVideoX</a> | <a href="#supported">CogVideoX 1.5</a> | <a href="#supported">ConsisID</a>🔥<br>
+    🔥<a href="#supported">Cosmos</a> | <a href="#supported">SkyReelsV2</a> | <a href="#supported">VisualCloze</a> | <a href="#supported">OmniGen 1/2</a> | <a href="#supported">Lumina 1/2</a> | <a href="#supported">PixArt</a>🔥<br>
+    🔥<a href="#supported">Chroma</a> | <a href="#supported">Sana</a> | <a href="#supported">Allegro</a> | <a href="#supported">Mochi</a> | <a href="#supported">SD 3/3.5</a> | <a href="#supported">Amused</a> | <a href="#supported"> ... </a> | <a href="#supported">DiT-XL</a>🔥
   </p>
 </div>
 <div align='center'>
   <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/wan2.2.C0_Q0_NONE.gif width=124px>
   <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/wan2.2.C1_Q0_DBCACHE_F1B0_W2M8MC2_T1O2_R0.08.gif width=124px>
-  <img src=./assets/gifs/hunyuan_video.C0_L0_Q0_NONE.gif width=126px>
-  <img src=./assets/gifs/hunyuan_video.C0_L0_Q0_DBCACHE_F1B0_W8M0MC2_T0O2_R0.12_S27.gif width=126px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/hunyuan_video.C0_L0_Q0_NONE.gif width=126px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/hunyuan_video.C0_L0_Q0_DBCACHE_F1B0_W8M0MC2_T0O2_R0.12_S27.gif width=126px>
   <p><b>🔥Wan2.2 MoE</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:2.0x↑🎉 | <b>HunyuanVideo</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:2.1x↑🎉</p>
   <img src=https://github.com/vipshop/cache-dit/raw/main/assets/qwen-image.C0_Q0_NONE.png width=160px>
   <img src=https://github.com/vipshop/cache-dit/raw/main/assets/qwen-image.C1_Q0_DBCACHE_F8B0_W8M0MC0_T1O4_R0.12_S23.png width=160px>
-  <img src=./assets/flux.C0_Q0_NONE_T23.69s.png width=90px>
-  <img src=./assets/flux.C0_Q0_DBCACHE_F1B0_W4M0MC0_T1O2_R0.15_S16_T11.39s.png width=90px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/flux.C0_Q0_NONE_T23.69s.png width=90px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/flux.C0_Q0_DBCACHE_F1B0_W4M0MC0_T1O2_R0.15_S16_T11.39s.png width=90px>
   <p><b>🔥Qwen-Image</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.8x↑🎉 | <b>FLUX.1-dev</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:2.1x↑🎉</p>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/flux-kontext-cat.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/flux-kontext.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/flux-kontext.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.08_S10.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/flux-kontext.C0_L0_Q0_DBCACHE_F1B0_W8M0MC2_T0O2_R0.12_S12.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/flux-kontext.C0_L0_Q0_DBCACHE_F1B0_W2M0MC2_T0O2_R0.15_S15.png width=100px>
+  <p><b>🔥FLUX-Kontext-dev</b> | Baseline | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.3x↑🎉 | 1.7x↑🎉 | 2.0x↑ 🎉</p>
   <img src=https://github.com/vipshop/cache-dit/raw/main/assets/qwen-image-lightning.4steps.C0_L1_Q0_NONE.png width=160px>
   <img src=https://github.com/vipshop/cache-dit/raw/main/assets/qwen-image-lightning.4steps.C0_L1_Q0_DBCACHE_F16B16_W2M1MC1_T0O2_R0.9_S1.png width=160px>
-  <img src=./assets/sd_3_5.C0_L0_Q0_NONE.png width=90px>
-  <img src=./assets/sd_3_5.C0_L0_Q0_DBCACHE_F1B0_W8M0MC3_T0O2_R0.12_S30.png width=90px>
-  <p><b>🔥Qwen-Image-Lightning</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.14x↑🎉 | <b>SD 3.5</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:2.5x↑🎉</p>
-  <img src=./assets/hidream.C0_L0_Q0_NONE.png width=100px>
-  <img src=./assets/hidream.C0_L0_Q0_DBCACHE_F1B0_W8M0MC0_T0O2_R0.08_S24.png width=100px>
-  <img src=./assets/cogview4.C0_L0_Q0_NONE.png width=100px>
-  <img src=./assets/cogview4.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.08_S15.png width=100px>
-  <img src=./assets/cogview4.C0_L0_Q0_DBCACHE_F1B0_W4M0MC4_T0O2_R0.2_S22.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/hunyuan-image-2.1.C0_L0_Q1_fp8_w8a16_wo_NONE.png width=90px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/hunyuan-image-2.1.C0_L0_Q1_fp8_w8a16_wo_DBCACHE_F8B0_W8M0MC2_T1O2_R0.12_S25.png width=90px>
+  <p><b>🔥Qwen...Lightning</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.14x↑🎉 | <b>HunyuanImage</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.7x↑🎉</p>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/examples/data/bear.png width=125px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/qwen-image-edit.C0_L0_Q0_NONE.png width=125px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/qwen-image-edit.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.08_S18.png width=125px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/qwen-image-edit.C0_L0_Q0_DBCACHE_F1B0_W8M0MC2_T0O2_R0.12_S24.png width=125px>
+  <p><b>🔥Qwen-Image-Edit</b> | Input w/o Edit | Baseline | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.6x↑🎉 | 1.9x↑🎉 </p>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/hidream.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/hidream.C0_L0_Q0_DBCACHE_F1B0_W8M0MC0_T0O2_R0.08_S24.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/cogview4.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/cogview4.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.08_S15.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/cogview4.C0_L0_Q0_DBCACHE_F1B0_W4M0MC4_T0O2_R0.2_S22.png width=100px>
   <p><b>🔥HiDream-I1</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.9x↑🎉 | <b>CogView4</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.4x↑🎉 | 1.7x↑🎉</p>
-  <img src=./assets/gifs/mochi.C0_L0_Q0_NONE.gif width=160px>
-  <img src=./assets/gifs/mochi.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.08_S34.gif width=160px>
-  <img src=./assets/hunyuan-image-2.1.C0_L0_Q1_fp8_w8a16_wo_NONE.png width=91px>
-  <img src=./assets/hunyuan-image-2.1.C0_L0_Q1_fp8_w8a16_wo_DBCACHE_F8B0_W8M0MC2_T1O2_R0.12_S25.png width=91px>
-  <p><b>🔥Mochi-1</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.8x↑🎉 | <b>HunyuanImage-2.1</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.7x↑🎉
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/cogview3_plus.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/cogview3_plus.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.08_S15.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/cogview3_plus.C0_L0_Q0_DBCACHE_F1B0_W8M0MC2_T0O2_R0.08_S25.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/chroma1-hd.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/chroma1-hd.C0_L0_Q0_DBCACHE_F1B0_W8M0MC0_T0O2_R0.08_S20.png width=100px>
+  <p><b>🔥CogView3</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.5x↑🎉 | 2.0x↑🎉| <b>Chroma1-HD</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.9x↑🎉</p>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/mochi.C0_L0_Q0_NONE.gif width=125px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/mochi.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.08_S34.gif width=125px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/skyreels_v2.C0_L0_Q0_NONE.gif width=125px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/skyreels_v2.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.12_S17.gif width=125px>
+  <p><b>🔥Mochi-1-preview</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.8x↑🎉 | <b>SkyReelsV2</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.6x↑🎉</p>
+  <img src=./examples/data/visualcloze/00555_00.jpg width=100px>
+  <img src=./examples/data/visualcloze/12265_00.jpg width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/visualcloze-512.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/visualcloze-512.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.08_S15.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/visualcloze-512.C0_L0_Q0_DBCACHE_F1B0_W8M0MC0_T0O2_R0.08_S18.png width=100px>
+  <p><b>🔥VisualCloze-512</b> | Model | Cloth | Baseline | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.4x↑🎉 | 1.7x↑🎉 </p>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/ltx-video.C0_L0_Q0_NONE.gif width=144px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/ltx-video.C0_L0_Q0_DBCACHE_F1B0_W8M0MC0_T0O2_R0.15_S13.gif width=144px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/cogvideox1.5.C0_L0_Q0_NONE.gif width=105px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/cogvideox1.5.C0_L0_Q0_DBCACHE_F1B0_W8M0MC0_T0O2_R0.12_S22.gif width=105px>
+  <p><b>🔥LTX-Video-0.9.7</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.7x↑🎉 | <b>CogVideoX1.5</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:2.0x↑🎉</p>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/omingen-v1.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/omingen-v1.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.08_S24.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/omingen-v1.C0_L0_Q0_DBCACHE_F1B0_W8M0MC0_T1O2_R0.08_S38.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/lumina2.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/lumina2.C0_L0_Q0_DBCACHE_F1B0_W2M0MC2_T0O2_R0.12_S14.png width=100px>
+  <p><b>🔥OmniGen-v1</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.5x↑🎉 | 3.3x↑🎉 | <b>Lumina2</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.9x↑🎉</p>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/allegro.C0_L0_Q0_NONE.gif width=117px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/allegro.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.26_S27.gif width=117px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/auraflow.C0_L0_Q0_NONE.png width=133px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/auraflow.C0_L0_Q0_DBCACHE_F1B0_W8M0MC2_T0O2_R0.08_S28.png width=133px>
+  <p><b>🔥Allegro</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.36x↑🎉 | <b>AuraFlow-v0.3</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:2.27x↑🎉 </p>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/sana.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/sana.C0_L0_Q0_DBCACHE_F8B0_W8M0MC2_T0O2_R0.25_S6.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/sana.C0_L0_Q0_DBCACHE_F1B0_W8M0MC2_T0O2_R0.3_S8.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/pixart-sigma.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/pixart-sigma.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.08_S28.png width=100px>
+  <p><b>🔥Sana</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.3x↑🎉 | 1.6x↑🎉| <b>PixArt-Sigma</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:2.3x↑🎉</p>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/pixart-alpha.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/pixart-alpha.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.05_S27.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/pixart-alpha.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.08_S32.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/sd_3_5.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/sd_3_5.C0_L0_Q0_DBCACHE_F1B0_W8M0MC3_T0O2_R0.12_S30.png width=100px>
+  <p><b>🔥PixArt-Alpha</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.6x↑🎉 | 1.8x↑🎉| <b>SD 3.5</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:2.5x↑🎉</p>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/amused.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/amused.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.34_S1.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/amused.C0_L0_Q0_DBCACHE_F8B0_W8M0MC0_T0O2_R0.38_S2.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/dit-xl.C0_L0_Q0_NONE.png width=100px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/dit-xl.C0_L0_Q0_DBCACHE_F1B0_W8M0MC2_T0O2_R0.15_S11.png width=100px>
+  <p><b>🔥Asumed</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.1x↑🎉 | 1.2x↑🎉 | <b>DiT-XL-256</b> | <a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:1.8x↑🎉
   <br>♥️ Please consider to leave a <b>⭐️ Star</b> to support us ~ ♥️</p>
 </div>
 
@@ -126,9 +183,10 @@ Dynamic: requires-python
 - [🔥Supported Models](#supported)
 - [🎉Unified Cache APIs](#unified)
   - [📚Forward Pattern Matching](#unified)
-  - [🎉Cache with One-line Code](#unified)
+  - [♥️Cache with One-line Code](#unified)
   - [🔥Automatic Block Adapter](#unified)
   - [📚Hybird Forward Pattern](#unified)
+  - [📚Implement Patch Functor](#unified)
   - [🤖Cache Acceleration Stats](#unified)
 - [⚡️Dual Block Cache](#dbcache)
 - [🔥Hybrid TaylorSeer](#taylorseer)
@@ -161,9 +219,15 @@ Currently, **cache-dit** library supports almost **Any** Diffusion Transformers
 - [🚀Qwen-Image-Lightning](https://github.com/vipshop/cache-dit/raw/main/examples)  
 - [🚀Qwen-Image-Edit](https://github.com/vipshop/cache-dit/raw/main/examples)  
 - [🚀Qwen-Image](https://github.com/vipshop/cache-dit/raw/main/examples)  
+- [🚀SkyReelsV2](https://github.com/vipshop/cache-dit/raw/main/examples)  
+- [🚀LTXVideo](https://github.com/vipshop/cache-dit/raw/main/examples)  
+- [🚀OmniGen](https://github.com/vipshop/cache-dit/raw/main/examples)  
+- [🚀Lumina2](https://github.com/vipshop/cache-dit/raw/main/examples)  
 - [🚀FLUX.1-dev](https://github.com/vipshop/cache-dit/raw/main/examples)  
 - [🚀FLUX.1-Fill-dev](https://github.com/vipshop/cache-dit/raw/main/examples)  
 - [🚀FLUX.1-Kontext-dev](https://github.com/vipshop/cache-dit/raw/main/examples)  
+- [🚀Chroma1-HD](https://github.com/vipshop/cache-dit/raw/main/examples)  
+- [🚀VisualCloze](https://github.com/vipshop/cache-dit/raw/main/examples)  
 - [🚀CogVideoX](https://github.com/vipshop/cache-dit/raw/main/examples)
 - [🚀CogVideoX1.5](https://github.com/vipshop/cache-dit/raw/main/examples)
 - [🚀CogView3-Plus](https://github.com/vipshop/cache-dit/raw/main/examples)
@@ -175,9 +239,16 @@ Currently, **cache-dit** library supports almost **Any** Diffusion Transformers
 - [🚀HunyuanVideo](https://github.com/vipshop/cache-dit/raw/main/examples)
 - [🚀HunyuanDiT](https://github.com/vipshop/cache-dit/raw/main/examples)
 - [🚀HiDream-I1-Full](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀AuraFlow-v0.3](https://github.com/vipshop/cache-dit/raw/main/examples)
 - [🚀PixArt-Alpha](https://github.com/vipshop/cache-dit/raw/main/examples)
 - [🚀PixArt-Sigma](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀NVIDIA Sana](https://github.com/vipshop/cache-dit/raw/main/examples)
 - [🚀SD-3/3.5](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀ConsisID](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀Allegro](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀Amused](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀DiT-XL](https://github.com/vipshop/cache-dit/raw/main/examples)
+- ...
 
 </details>
 
@@ -265,6 +336,75 @@ cache_dit.enable_cache(
 )
 ```
 
+Even sometimes you have more complex cases, such as **Wan 2.2 MoE**, which has more than one Transformer (namely `transformer` and `transformer_2`) in its structure. Fortunately, **cache-dit** can also handle this situation very well. Please refer to [📚Wan 2.2 MoE](./examples/pipeline/run_wan_2.2.py) as an example.
+
+```python
+from cache_dit import ForwardPattern, BlockAdapter, ParamsModifier
+
+cache_dit.enable_cache(
+    BlockAdapter(
+        pipe=pipe,
+        transformer=[
+            pipe.transformer,
+            pipe.transformer_2,
+        ],
+        blocks=[
+            pipe.transformer.blocks,
+            pipe.transformer_2.blocks,
+        ],
+        forward_pattern=[
+            ForwardPattern.Pattern_2,
+            ForwardPattern.Pattern_2,
+        ],
+        # Setup different cache params for each 'blocks'. You can 
+        # pass any specific cache params to ParamModifier, the old 
+        # value will be overwrite by the new one.
+        params_modifiers=[
+            ParamsModifier(
+                max_warmup_steps=4,
+                max_cached_steps=8,
+            ),
+            ParamsModifier(
+                max_warmup_steps=2,
+                max_cached_steps=20,
+            ),
+        ],
+        has_separate_cfg=True,
+    ),
+)
+```
+### 📚Implement Patch Functor
+
+For any PATTERN not {0...5}, we introduced the simple abstract concept of **Patch Functor**. Users can implement a subclass of Patch Functor to convert an unknown Pattern into a known PATTERN, and for some models, users may also need to fuse the operations within the blocks for loop into block forward. 
+
+![](https://github.com/vipshop/cache-dit/raw/main/assets/patch-functor.png)
+
+Some Patch functors have already been provided in cache-dit: [📚HiDreamPatchFunctor](./src/cache_dit/cache_factory/patch_functors/functor_hidream.py), [📚ChromaPatchFunctor](./src/cache_dit/cache_factory/patch_functors/functor_chroma.py), etc. After implementing Patch Functor, users need to set the `patch_functor` property of **BlockAdapter**.
+
+```python
+@BlockAdapterRegistry.register("HiDream")
+def hidream_adapter(pipe, **kwargs) -> BlockAdapter:
+    from diffusers import HiDreamImageTransformer2DModel
+    from cache_dit.cache_factory.patch_functors import HiDreamPatchFunctor
+
+    assert isinstance(pipe.transformer, HiDreamImageTransformer2DModel)
+    return BlockAdapter(
+        pipe=pipe,
+        transformer=pipe.transformer,
+        blocks=[
+            pipe.transformer.double_stream_blocks,
+            pipe.transformer.single_stream_blocks,
+        ],
+        forward_pattern=[
+            ForwardPattern.Pattern_0,
+            ForwardPattern.Pattern_3,
+        ],
+        # NOTE: Setup your custom patch functor here.
+        patch_functor=HiDreamPatchFunctor(),
+        **kwargs,
+    )
+```
+
 ### 🤖Cache Acceleration Stats Summary
 
 After finishing each inference of `pipe(...)`, you can call the `cache_dit.summary()` API on pipe to get the details of the **Cache Acceleration Stats** for the current inference. 
@@ -348,7 +488,7 @@ cache_dit.enable_cache(
     # Taylorseer cache type cache be hidden_states or residual.
     taylorseer_cache_type="residual",
     # Higher values of order will lead to longer computation time
-    taylorseer_order=2, # default is 2.
+    taylorseer_order=1, # default is 1.
     max_warmup_steps=3, # prefer: >= order + 1
     residual_diff_threshold=0.12
 )s
@@ -372,7 +512,7 @@ cache_dit.enable_cache(
 
 <div id="cfg"></div>
 
-cache-dit supports caching for **CFG (classifier-free guidance)**. For models that fuse CFG and non-CFG into a single forward step, or models that do not include CFG (classifier-free guidance) in the forward step, please set `enable_spearate_cfg` param to **False (default)**. Otherwise, set it to True. For examples:
+cache-dit supports caching for **CFG (classifier-free guidance)**. For models that fuse CFG and non-CFG into a single forward step, or models that do not include CFG (classifier-free guidance) in the forward step, please set `enable_separate_cfg` param to **False (default, None)**. Otherwise, set it to True. For examples:
 
 ```python
 cache_dit.enable_cache(
@@ -380,14 +520,14 @@ cache_dit.enable_cache(
     ...,
     # CFG: classifier free guidance or not
     # For model that fused CFG and non-CFG into single forward step,
-    # should set enable_spearate_cfg as False. For example, set it as True 
+    # should set enable_separate_cfg as False. For example, set it as True 
     # for Wan 2.1/Qwen-Image and set it as False for FLUX.1, HunyuanVideo, 
     # CogVideoX, Mochi, LTXVideo, Allegro, CogView3Plus, EasyAnimate, SD3, etc.
-    enable_spearate_cfg=True, # Wan 2.1, Qwen-Image, CogView4, Cosmos, SkyReelsV2, etc.
+    enable_separate_cfg=True, # Wan 2.1, Qwen-Image, CogView4, Cosmos, SkyReelsV2, etc.
     # Compute cfg forward first or not, default False, namely, 
     # 0, 2, 4, ..., -> non-CFG step; 1, 3, 5, ... -> CFG step.
     cfg_compute_first=False,
-    # Compute spearate diff values for CFG and non-CFG step, 
+    # Compute separate diff values for CFG and non-CFG step, 
     # default True. If False, we will use the computed diff from 
     # current non-CFG transformer step for current CFG step.
     cfg_diff_compute_separate=True,

{cache_dit-0.2.33.dist-info → cache_dit-0.2.34.dist-info}/RECORD

@@ -1,29 +1,30 @@
 cache_dit/__init__.py,sha256=kX9V-FegZG4c8LMwI4PTmMqH794MEW0pzDArdhC0cJw,1241
-cache_dit/_version.py,sha256=gTEHTWtuqv38KTvjBsXd5hC019b6d7AyfC8gLMY7KAo,706
+cache_dit/_version.py,sha256=CtkelOzOJFXtgJ0APT8pLd5zWrG63eLavWaOD_cX7xo,706
 cache_dit/logger.py,sha256=0zsu42hN-3-rgGC_C29ms1IvVpV4_b4_SwJCKSenxBE,4304
 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=dmNX68nBD52HtQvHnNAuSn1zjDWrQdycD0qXy-w-mwc,18212
-cache_dit/cache_factory/cache_interface.py,sha256=LpyCy-tQ_GcTRAYLpMMf9hFVIktABHI6CObn5Ll8bMw,8548
+cache_dit/cache_factory/cache_adapters.py,sha256=OFJlxxyODhoZstN4EfPgC7tE8M1ZdQFcE25gDNrW7NA,18212
+cache_dit/cache_factory/cache_interface.py,sha256=tHQv7i8Hp6nfbjZWHwDx3nEvCfxLeBw26aMYjyu6nMw,8541
 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=OZM5vJwmQIkoIwVmMxKXiHqKvs31NyAva1Z91C_ko3w,17547
-cache_dit/cache_factory/block_adapters/block_adapters.py,sha256=IqHV10aK2qA8kEVDi7EEoUSBt0GzwCUM4GpLNf8Jgww,21656
-cache_dit/cache_factory/block_adapters/block_registers.py,sha256=ZeN2wGPmuf2u3puSsBx8x-rl3wRo8-cWcuWNcrssVfA,2553
+cache_dit/cache_factory/block_adapters/__init__.py,sha256=33geXMz56TxFWMp0c-H4__MY5SGRzKMKj3TXnUYOMlc,17512
+cache_dit/cache_factory/block_adapters/block_adapters.py,sha256=zZbbsZYWbUClfa6He69w_Wdf8ZLhKwMAb9gURYEUmgQ,23725
+cache_dit/cache_factory/block_adapters/block_registers.py,sha256=2L7QeM4ygnaKQpC9PoJod0QRYyxidUKU2AYpysDCUwE,2572
 cache_dit/cache_factory/cache_blocks/__init__.py,sha256=08Ox7kD05lkRKCOsVTdEZeKAWBheqpxfrAT1Nz7eclI,2916
 cache_dit/cache_factory/cache_blocks/pattern_0_1_2.py,sha256=ElMps6_7uI74tSF9GDR_dEI0bZEhdzcepM29xFWnYo8,428
 cache_dit/cache_factory/cache_blocks/pattern_3_4_5.py,sha256=Bv56qETXhsREvCrNvnZpSqDIIHsi6Ze3FJW4Yk2x3uI,8597
 cache_dit/cache_factory/cache_blocks/pattern_base.py,sha256=d4H9kEB0AgnVMT8aF0Y54SUMUQUxw5HQ8gRkoCuTQ_A,14577
 cache_dit/cache_factory/cache_blocks/utils.py,sha256=dGOC1tMMOvcbvEgx44eTESKn_jsv-0RZ3tRHPa3wmQ4,1315
 cache_dit/cache_factory/cache_contexts/__init__.py,sha256=rqnJ5__zqnpVHK5A1OqWILpNh5Ss-0ZDTGgtxZMKGGo,250
-cache_dit/cache_factory/cache_contexts/cache_context.py,sha256=zqixcxV_LjnyoYDZ6q3HAC-hqYyVV6g0MWKBI2hA1nQ,11855
-cache_dit/cache_factory/cache_contexts/cache_manager.py,sha256=Mcj1upIpXT_CwO4AdY4ZNJSWoOXn3Lx2mBZRi_QuLbU,32710
-cache_dit/cache_factory/cache_contexts/taylorseer.py,sha256=hgLmgIkQgwbFTjxqtLUCJ3mgDGEcJK09B7RK8sBdPiI,3593
-cache_dit/cache_factory/patch_functors/__init__.py,sha256=06zdddrjvSCgBzJ0a8niRHd3ucF2qsbzlbL00d4aCvk,451
+cache_dit/cache_factory/cache_contexts/cache_context.py,sha256=FWdgInClWY8VZBsZIevtYk--rX-RL8c3QfNOJtqR8a4,11855
+cache_dit/cache_factory/cache_contexts/cache_manager.py,sha256=Ig5VKoQ46iG3lKmsaMulYxd2vCm__2rY8NBvERwexwM,32719
+cache_dit/cache_factory/cache_contexts/taylorseer.py,sha256=4nxgSEZvDy-w-7XuJYzsyzdtF1_uFrDwlF06XBDFVKQ,3922
+cache_dit/cache_factory/patch_functors/__init__.py,sha256=oI6F3N9ezahRHaFUOZ1GfrAw1qFdKrxFXXmlwwehHj4,530
 cache_dit/cache_factory/patch_functors/functor_base.py,sha256=Ahk0fTfrHgNdEl-9JSkACvfyyv9G-Ei5OSz7XBIlX5o,357
-cache_dit/cache_factory/patch_functors/functor_chroma.py,sha256=2iLxlsc-1dDHRveqCXaC07E9CeMNOuBNkvpJ1atpK7E,10048
+cache_dit/cache_factory/patch_functors/functor_chroma.py,sha256=xD0Q96VArp1vYBLQ0pcjRIyFB1i_Y7muZ2q07Hz8Oqs,13430
+cache_dit/cache_factory/patch_functors/functor_dit.py,sha256=SDjhzCWa6PoFNN4_upoQEf6DHvW1yJ7zuXMS2VvyJco,3904
 cache_dit/cache_factory/patch_functors/functor_flux.py,sha256=UMkyuEYjO7UO_zmXi9Djd-nD-XMgCUgE-qkYA3plWSM,9559
 cache_dit/cache_factory/patch_functors/functor_hidream.py,sha256=pi_vvpDy1lsgQHxu3eK3v93rdJL7oNwkt3WakRP8pbw,15375
 cache_dit/cache_factory/patch_functors/functor_hunyuan_dit.py,sha256=iSo5dD5uKnjQQeysDUIkKt0wdnK5bzXTc_F_lfHG70w,6401
@@ -40,9 +41,9 @@ cache_dit/metrics/metrics.py,sha256=8jvM1sF-nDxUuwCRy44QEoo4dYVLCQVh1QyAMs4eaQY,
 cache_dit/quantize/__init__.py,sha256=kWYoMAyZgBXu9BJlZjTQ0dRffW9GqeeY9_iTkXrb70A,59
 cache_dit/quantize/quantize_ao.py,sha256=Fx1KW4l3gdEkdrcAYtPoDW7WKBJWrs3glOHiEwW_TgE,6160
 cache_dit/quantize/quantize_interface.py,sha256=2s_R7xPSKuJeFpEGeLwRxnq_CqJcBG3a3lzyW5wh-UM,1241
-cache_dit-0.2.33.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
-cache_dit-0.2.33.dist-info/METADATA,sha256=GQBvDzKLXL3tABguCRqLNc-Z39h0AcMK_J37demDTu8,25977
-cache_dit-0.2.33.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cache_dit-0.2.33.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
-cache_dit-0.2.33.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
-cache_dit-0.2.33.dist-info/RECORD,,
+cache_dit-0.2.34.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
+cache_dit-0.2.34.dist-info/METADATA,sha256=BvEY08xjrGPcqTEZSHvSDtJP4sGZv1T6jzhGj-jQbvo,38284
+cache_dit-0.2.34.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-0.2.34.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
+cache_dit-0.2.34.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-0.2.34.dist-info/RECORD,,