cache-dit 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

cache_dit/cache_factory/cache_interface.py

@@ -4,8 +4,9 @@ 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_adapters import CachedAdapterV2
+from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
 from cache_dit.cache_factory.cache_contexts import CalibratorConfig
+from cache_dit.cache_factory.params_modifier import ParamsModifier
 
 from cache_dit.logger import init_logger
 
@@ -18,25 +19,20 @@ def enable_cache(
         DiffusionPipeline,
         BlockAdapter,
     ],
-    # Cache context kwargs
-    Fn_compute_blocks: int = 8,
-    Bn_compute_blocks: int = 0,
-    max_warmup_steps: int = 8,
-    max_cached_steps: int = -1,
-    max_continuous_cached_steps: int = -1,
-    residual_diff_threshold: float = 0.08,
-    # Cache CFG or not
-    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 = 1,
-    # New param only for v2 API
+    # Basic DBCache config: BasicCacheConfig
+    cache_config: BasicCacheConfig = BasicCacheConfig(),
+    # Calibrator config: TaylorSeerCalibratorConfig, etc.
     calibrator_config: Optional[CalibratorConfig] = None,
-    **other_cache_context_kwargs,
+    # Modify cache context params for specific blocks.
+    params_modifiers: Optional[
+        Union[
+            ParamsModifier,
+            List[ParamsModifier],
+            List[List[ParamsModifier]],
+        ]
+    ] = None,
+    # Other cache context kwargs: Deprecated cache kwargs
+    **kwargs,
 ) -> Union[
     DiffusionPipeline,
     BlockAdapter,
@@ -53,55 +49,51 @@ def enable_cache(
             The standard Diffusion Pipeline or custom BlockAdapter (from cache-dit or user-defined).
             For example: cache_dit.enable_cache(FluxPipeline(...)). Please check https://github.com/vipshop/cache-dit/blob/main/docs/BlockAdapter.md
             for the usgae of BlockAdapter.
-        Fn_compute_blocks (`int`, *required*, defaults to 8):
-            Specifies that `DBCache` uses the **first n** Transformer blocks to fit the information
-            at time step t, enabling the calculation of a more stable L1 diff and delivering more
-            accurate information to subsequent blocks. Please check https://github.com/vipshop/cache-dit/blob/main/docs/DBCache.md
-            for more details of DBCache.
-        Bn_compute_blocks: (`int`, *required*, defaults to 0):
-            Further fuses approximate information in the **last n** Transformer blocks to enhance
-            prediction accuracy. These blocks act as an auto-scaler for approximate hidden states
-            that use residual cache.
-        max_warmup_steps (`int`, *required*, defaults to 8):
-            DBCache does not apply the caching strategy when the number of running steps is less than
-            or equal to this value, ensuring the model sufficiently learns basic features during warmup.
-        max_cached_steps (`int`, *required*, defaults to -1):
-            DBCache disables the caching strategy when the previous cached steps exceed this value to
-            prevent precision degradation.
-        max_continuous_cached_steps (`int`, *required*, defaults to -1):
-            DBCache disables the caching strategy when the previous continous cached steps exceed this value to
-            prevent precision degradation.
-        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_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_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 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
-            [TaylorSeers: From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers](https://arxiv.org/pdf/2503.06923) algorithm
-            to further improve the precision of DBCache in cases where the cached steps are large,
-            namely, **Hybrid TaylorSeer + DBCache**. At timesteps with significant intervals,
-            the feature similarity in diffusion models decreases substantially, significantly
-            harming the generation quality.
-        enable_encoder_taylorseer (`bool`, *required*,  defaults to False):
-            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 1):
-            The order of taylorseer, higher values of n_derivatives will lead to longer computation time,
-            the recommended value is 1 or 2.
+        cache_config (`BasicCacheConfig`, *required*, defaults to BasicCacheConfig()):
+            Basic DBCache config for cache context, defaults to BasicCacheConfig(). The configurable params listed belows:
+                Fn_compute_blocks: (`int`, *required*, defaults to 8):
+                    Specifies that `DBCache` uses the **first n** Transformer blocks to fit the information
+                    at time step t, enabling the calculation of a more stable L1 diff and delivering more
+                    accurate information to subsequent blocks. Please check https://github.com/vipshop/cache-dit/blob/main/docs/DBCache.md
+                    for more details of DBCache.
+                Bn_compute_blocks: (`int`, *required*, defaults to 0):
+                    Further fuses approximate information in the **last n** Transformer blocks to enhance
+                    prediction accuracy. These blocks act as an auto-scaler for approximate hidden states
+                    that use residual cache.
+                residual_diff_threshold (`float`, *required*, defaults to 0.08):
+                    the value of residual diff threshold, a higher value leads to faster performance at the
+                    cost of lower precision.
+                max_warmup_steps (`int`, *required*, defaults to 8):
+                    DBCache does not apply the caching strategy when the number of running steps is less than
+                    or equal to this value, ensuring the model sufficiently learns basic features during warmup.
+                max_cached_steps (`int`, *required*, defaults to -1):
+                    DBCache disables the caching strategy when the previous cached steps exceed this value to
+                    prevent precision degradation.
+                max_continuous_cached_steps (`int`, *required*, defaults to -1):
+                    DBCache disables the caching strategy when the previous continous cached steps exceed this value to
+                    prevent precision degradation.
+                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_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 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.
         calibrator_config (`CalibratorConfig`, *optional*, defaults to None):
-            # config for calibrator, if calibrator_config is not None, means that user want to use CachedAdapterV2
-            # with specific calibrator, such as taylorseer, foca, and so on.
-        other_cache_context_kwargs: (`dict`, *optional*, defaults to {})
+            Config for calibrator, if calibrator_config is not None, means that user want to use DBCache
+            with specific calibrator, such as taylorseer, foca, and so on.
+        params_modifiers ('ParamsModifier', *optional*, defaults to None):
+            Modify cache context params for specific blocks. The configurable params listed belows:
+                cache_config (`BasicCacheConfig`, *required*, defaults to BasicCacheConfig()):
+                    The same as 'cache_config' param in cache_dit.enable_cache() interface.
+                calibrator_config (`CalibratorConfig`, *optional*, defaults to None):
+                    The same as 'calibrator_config' param in cache_dit.enable_cache() interface.
+                **kwargs: (`dict`, *optional*, defaults to {}):
+                    The same as 'kwargs' param in cache_dit.enable_cache() interface.
+        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.
 
@@ -116,51 +108,82 @@ def enable_cache(
     >>> 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)) is not None:
+    cache_context_kwargs = {}
+    if (cache_type := cache_context_kwargs.pop("cache_type", None)) is not None:
         if cache_type == CacheType.NONE:
             return pipe_or_adapter
 
-    cache_context_kwargs["cache_type"] = CacheType.DBCache
-    cache_context_kwargs["Fn_compute_blocks"] = Fn_compute_blocks
-    cache_context_kwargs["Bn_compute_blocks"] = Bn_compute_blocks
-    cache_context_kwargs["max_warmup_steps"] = max_warmup_steps
-    cache_context_kwargs["max_cached_steps"] = max_cached_steps
-    cache_context_kwargs["max_continuous_cached_steps"] = (
-        max_continuous_cached_steps
-    )
-    cache_context_kwargs["residual_diff_threshold"] = residual_diff_threshold
-    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
-    )
+    # WARNING: Deprecated cache config params. These parameters are now retained
+    # for backward compatibility but will be removed in the future.
+    deprecated_cache_kwargs = {
+        "Fn_compute_blocks": kwargs.get("Fn_compute_blocks", None),
+        "Bn_compute_blocks": kwargs.get("Bn_compute_blocks", None),
+        "max_warmup_steps": kwargs.get("max_warmup_steps", None),
+        "max_cached_steps": kwargs.get("max_cached_steps", None),
+        "max_continuous_cached_steps": kwargs.get(
+            "max_continuous_cached_steps", None
+        ),
+        "residual_diff_threshold": kwargs.get("residual_diff_threshold", None),
+        "enable_separate_cfg": kwargs.get("enable_separate_cfg", None),
+        "cfg_compute_first": kwargs.get("cfg_compute_first", None),
+        "cfg_diff_compute_separate": kwargs.get(
+            "cfg_diff_compute_separate", None
+        ),
+    }
+
+    deprecated_cache_kwargs = {
+        k: v for k, v in deprecated_cache_kwargs.items() if v is not None
+    }
+
+    if deprecated_cache_kwargs:
+        logger.warning(
+            "Manually settup DBCache context without BasicCacheConfig is "
+            "deprecated and will be removed in the future, please use "
+            "`cache_config` parameter instead!"
+        )
+        if cache_config is not None:
+            cache_config.update(**deprecated_cache_kwargs)
+        else:
+            cache_config = BasicCacheConfig(**deprecated_cache_kwargs)
+
+    if cache_config is not None:
+        cache_context_kwargs["cache_config"] = cache_config
+
+    # WARNING: Deprecated taylorseer params. These parameters are now retained
+    # for backward compatibility but will be removed in the future.
+    if (
+        kwargs.get("enable_taylorseer", None) is not None
+        or kwargs.get("enable_encoder_taylorseer", None) is not None
+    ):
+        logger.warning(
+            "Manually settup TaylorSeer calibrator without TaylorSeerCalibratorConfig is "
+            "deprecated and will be removed in the future, please use "
+            "`calibrator_config` parameter instead!"
+        )
+        from cache_dit.cache_factory.cache_contexts.calibrators import (
+            TaylorSeerCalibratorConfig,
+        )
 
-    # V1 only supports the Taylorseer calibrator. We have decided to
-    # keep this code for API compatibility reasons.
-    if calibrator_config is None:
-        cache_context_kwargs["enable_taylorseer"] = enable_taylorseer
-        cache_context_kwargs["enable_encoder_taylorseer"] = (
-            enable_encoder_taylorseer
+        calibrator_config = TaylorSeerCalibratorConfig(
+            enable_calibrator=kwargs.get("enable_taylorseer"),
+            enable_encoder_calibrator=kwargs.get("enable_encoder_taylorseer"),
+            calibrator_cache_type=kwargs.get(
+                "taylorseer_cache_type", "residual"
+            ),
+            taylorseer_order=kwargs.get("taylorseer_order", 1),
         )
-        cache_context_kwargs["taylorseer_cache_type"] = taylorseer_cache_type
-        cache_context_kwargs["taylorseer_order"] = taylorseer_order
-    else:
+
+    if calibrator_config is not None:
         cache_context_kwargs["calibrator_config"] = calibrator_config
 
+    if params_modifiers is not None:
+        cache_context_kwargs["params_modifiers"] = params_modifiers
+
     if isinstance(pipe_or_adapter, (DiffusionPipeline, BlockAdapter)):
-        if calibrator_config is None:
-            return CachedAdapter.apply(
-                pipe_or_adapter,
-                **cache_context_kwargs,
-            )
-        else:
-            logger.warning("You are using the un-stable V2 API!")
-            pipe_or_adapter._is_v2_api = True
-            return CachedAdapterV2.apply(
-                pipe_or_adapter,
-                **cache_context_kwargs,
-            )
+        return CachedAdapter.apply(
+            pipe_or_adapter,
+            **cache_context_kwargs,
+        )
     else:
         raise ValueError(
             f"type: {type(pipe_or_adapter)} is not valid, "
@@ -175,13 +198,7 @@ def disable_cache(
         BlockAdapter,
     ],
 ):
-    if getattr(pipe_or_adapter, "_is_v2_api", False):
-        logger.warning("You are using the un-stable V2 API!")
-        CachedAdapterV2.maybe_release_hooks(pipe_or_adapter)
-        del pipe_or_adapter._is_v2_api
-    else:
-        CachedAdapter.maybe_release_hooks(pipe_or_adapter)
-
+    CachedAdapter.maybe_release_hooks(pipe_or_adapter)
     logger.warning(
         f"Cache Acceleration is disabled for: "
         f"{pipe_or_adapter.__class__.__name__}."

cache_dit/cache_factory/params_modifier.py

@@ -0,0 +1,87 @@
+from typing import Optional
+
+from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
+from cache_dit.cache_factory.cache_contexts import CalibratorConfig
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class ParamsModifier:
+    def __init__(
+        self,
+        # Basic DBCache config: BasicCacheConfig
+        cache_config: BasicCacheConfig = None,
+        # Calibrator config: TaylorSeerCalibratorConfig, etc.
+        calibrator_config: Optional[CalibratorConfig] = None,
+        # Other cache context kwargs: Deprecated cache kwargs
+        **kwargs,
+    ):
+        self._context_kwargs = {}
+
+        # WARNING: Deprecated cache config params. These parameters are now retained
+        # for backward compatibility but will be removed in the future.
+        deprecated_cache_kwargs = {
+            "Fn_compute_blocks": kwargs.get("Fn_compute_blocks", None),
+            "Bn_compute_blocks": kwargs.get("Bn_compute_blocks", None),
+            "max_warmup_steps": kwargs.get("max_warmup_steps", None),
+            "max_cached_steps": kwargs.get("max_cached_steps", None),
+            "max_continuous_cached_steps": kwargs.get(
+                "max_continuous_cached_steps", None
+            ),
+            "residual_diff_threshold": kwargs.get(
+                "residual_diff_threshold", None
+            ),
+            "enable_separate_cfg": kwargs.get("enable_separate_cfg", None),
+            "cfg_compute_first": kwargs.get("cfg_compute_first", None),
+            "cfg_diff_compute_separate": kwargs.get(
+                "cfg_diff_compute_separate", None
+            ),
+        }
+
+        deprecated_cache_kwargs = {
+            k: v for k, v in deprecated_cache_kwargs.items() if v is not None
+        }
+
+        if deprecated_cache_kwargs:
+            logger.warning(
+                "Manually settup DBCache context without BasicCacheConfig is "
+                "deprecated and will be removed in the future, please use "
+                "`cache_config` parameter instead!"
+            )
+            if cache_config is not None:
+                cache_config.update(**deprecated_cache_kwargs)
+            else:
+                cache_config = BasicCacheConfig(**deprecated_cache_kwargs)
+
+        if cache_config is not None:
+            self._context_kwargs["cache_config"] = cache_config
+        # WARNING: Deprecated taylorseer params. These parameters are now retained
+        # for backward compatibility but will be removed in the future.
+        if (
+            kwargs.get("enable_taylorseer", None) is not None
+            or kwargs.get("enable_encoder_taylorseer", None) is not None
+        ):
+            logger.warning(
+                "Manually settup TaylorSeer calibrator without TaylorSeerCalibratorConfig is "
+                "deprecated and will be removed in the future, please use "
+                "`calibrator_config` parameter instead!"
+            )
+            from cache_dit.cache_factory.cache_contexts.calibrators import (
+                TaylorSeerCalibratorConfig,
+            )
+
+            calibrator_config = TaylorSeerCalibratorConfig(
+                enable_calibrator=kwargs.get("enable_taylorseer"),
+                enable_encoder_calibrator=kwargs.get(
+                    "enable_encoder_taylorseer"
+                ),
+                calibrator_cache_type=kwargs.get(
+                    "taylorseer_cache_type", "residual"
+                ),
+                taylorseer_order=kwargs.get("taylorseer_order", 1),
+            )
+
+        if calibrator_config is not None:
+            self._context_kwargs["calibrator_config"] = calibrator_config

cache_dit/metrics/__init__.py

@@ -4,10 +4,12 @@ from cache_dit.metrics.metrics import compute_mse
 from cache_dit.metrics.metrics import compute_video_psnr
 from cache_dit.metrics.metrics import compute_video_ssim
 from cache_dit.metrics.metrics import compute_video_mse
-from cache_dit.metrics.metrics import entrypoint
 from cache_dit.metrics.fid import FrechetInceptionDistance
+from cache_dit.metrics.fid import compute_fid
+from cache_dit.metrics.fid import compute_video_fid
 from cache_dit.metrics.config import set_metrics_verbose
 from cache_dit.metrics.config import get_metrics_verbose
+from cache_dit.metrics.metrics import entrypoint
 
 
 def main():

cache_dit/utils.py

@@ -9,7 +9,9 @@ from pprint import pprint
 from diffusers import DiffusionPipeline
 
 from typing import Dict, Any, List, Union
+from cache_dit.cache_factory import CacheType
 from cache_dit.cache_factory import BlockAdapter
+from cache_dit.cache_factory import BasicCacheConfig
 from cache_dit.cache_factory import CalibratorConfig
 from cache_dit.logger import init_logger
 
@@ -162,7 +164,6 @@ def strify(
         cache_options = stats.cache_options
         cached_steps = len(stats.cached_steps)
     elif isinstance(adapter_or_others, dict):
-        from cache_dit.cache_factory import CacheType
 
         # Assume cache_context_kwargs
         cache_options = adapter_or_others
@@ -180,31 +181,21 @@ def strify(
     if not cache_options:
         return "NONE"
 
-    def calibrator_str():
-        if not getattr(adapter_or_others, "_is_v2_api", False):
-            taylorseer_order = 0
-            if "taylorseer_order" in cache_options:
-                taylorseer_order = cache_options["taylorseer_order"]
-            return (
-                f"T{int(cache_options.get('enable_taylorseer', False))}"
-                f"O{taylorseer_order}"
-            )
+    def basic_cache_str():
+        cache_config: BasicCacheConfig = cache_options.get("cache_config", None)
+        if cache_config is not None:
+            return cache_config.strify()
+        return "NONE"
 
+    def calibrator_str():
         calibrator_config: CalibratorConfig = cache_options.get(
             "calibrator_config", None
         )
+        if calibrator_config is not None:
+            return calibrator_config.strify()
+        return "T0O0"
 
-        return calibrator_config.strify() if calibrator_config else "NONE"
-
-    cache_type_str = (
-        f"DBCACHE_F{cache_options.get('Fn_compute_blocks', 1)}"
-        f"B{cache_options.get('Bn_compute_blocks', 0)}_"
-        f"W{cache_options.get('max_warmup_steps', 0)}"
-        f"M{max(0, cache_options.get('max_cached_steps', -1))}"
-        f"MC{max(0, cache_options.get('max_continuous_cached_steps', -1))}_"
-        f"{calibrator_str()}_"
-        f"R{cache_options.get('residual_diff_threshold', 0.08)}"
-    )
+    cache_type_str = f"{basic_cache_str()}_{calibrator_str()}"
 
     if cached_steps:
         cache_type_str += f"_S{cached_steps}"