cache-dit 0.2.33__py3-none-any.whl → 0.2.36__py3-none-any.whl

cache_dit/__init__.py

@@ -4,6 +4,10 @@ except ImportError:
     __version__ = "unknown version"
     version_tuple = (0, 0, "unknown version")
 
+from cache_dit.utils import summary
+from cache_dit.utils import strify
+from cache_dit.utils import disable_print
+from cache_dit.logger import init_logger
 from cache_dit.cache_factory import load_options
 from cache_dit.cache_factory import enable_cache
 from cache_dit.cache_factory import disable_cache
@@ -18,9 +22,7 @@ from cache_dit.cache_factory import supported_pipelines
 from cache_dit.cache_factory import get_adapter
 from cache_dit.compile import set_compile_configs
 from cache_dit.quantize import quantize
-from cache_dit.utils import summary
-from cache_dit.utils import strify
-from cache_dit.logger import init_logger
+
 
 NONE = CacheType.NONE
 DBCache = CacheType.DBCache

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.36'
+__version_tuple__ = version_tuple = (0, 2, 36)
 
 __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,