Skip to content

API Documentation

Unified Cache API for almost Any Diffusion Transformers (with Transformer Blocks that match the specific Input and Output patterns). For a good balance between performance and precision, DBCache is configured by default with F8B0, 8 warmup steps, and unlimited cached steps. All the configurable params are listed beflows.

API: enable_cache

def enable_cache(...) -> Union[DiffusionPipeline, BlockAdapter, Transformer]

Function Description

The enable_cache function serves as a unified caching interface designed to optimize the performance of diffusion transformer models by implementing an intelligent caching mechanism known as DBCache. This API is engineered to be compatible with nearly all diffusion transformer architectures that feature transformer blocks adhering to standard input-output patterns, eliminating the need for architecture-specific modifications.

By strategically caching intermediate outputs of transformer blocks during the diffusion process, DBCache significantly reduces redundant computations without compromising generation quality. The caching mechanism works by tracking residual differences between consecutive steps, allowing the model to reuse previously computed features when these differences fall below a configurable threshold. This approach maintains a balance between computational efficiency and output precision.

The default configuration (F8B0, 8 warmup steps, unlimited cached steps) is carefully tuned to provide an optimal tradeoff for most common use cases. The "F8B0" configuration indicates that the first 8 transformer blocks are used to compute stable feature differences, while no final blocks are employed for additional fusion. The warmup phase ensures the model establishes sufficient feature representation before caching begins, preventing potential degradation of output quality.

This function seamlessly integrates with both standard diffusion pipelines and custom block adapters, making it versatile for various deployment scenarios—from research prototyping to production environments where inference speed is critical. By abstracting the complexity of caching logic behind a simple interface, it enables developers to enhance model performance with minimal code changes.

Quick Start

>>> import cache_dit
>>> from diffusers import DiffusionPipeline
>>> pipe = DiffusionPipeline.from_pretrained("Qwen/Qwen-Image") # Can be any diffusion pipeline
>>> cache_dit.enable_cache(pipe) # One-line code with default cache options.
>>> output = pipe(...) # Just call the pipe as normal.
>>> stats = cache_dit.summary(pipe) # Then, get the summary of cache acceleration stats.
>>> cache_dit.disable_cache(pipe) # Disable cache and run original pipe.

Parameter Description

  • pipe_or_adapter(DiffusionPipeline, BlockAdapter or Transformer, required):
    The standard Diffusion Pipeline or custom BlockAdapter (from cache-dit or user-defined). For example: cache_dit.enable_cache(pipe).

  • cache_config(DBCacheConfig, required, defaults to DBCacheConfig()):
    Basic DBCache config for cache context, defaults to DBCacheConfig(). The configurable parameters are listed below:

    • Fn_compute_blocks: (int, required, defaults to 8):
      Specifies that DBCache uses thefirst nTransformer blocks to fit the information at time step t, enabling the calculation of a more stable L1 difference and delivering more accurate information to subsequent blocks.
    • Bn_compute_blocks: (int, required, defaults to 0):
      Further fuses approximate information in thelast nTransformer 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 difference threshold, a higher value leads to faster performance at the cost of lower precision.
    • max_accumulated_residual_diff_threshold: (float, optional, defaults to None):
      The maximum accumulated relative l1 diff threshold for Cache. If set, when the accumulated relative l1 diff exceeds this threshold, the caching strategy will be disabled for current step. This is useful for some cases where the input condition changes significantly in a single step. Default None means this feature is disabled.
    • 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.
    • warmup_interval: (int, required, defaults to 1):
      Skip interval in warmup steps, e.g., when warmup_interval is 2, only 0, 2, 4, ... steps in warmup steps will be computed, others will use dynamic cache.
    • 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 continuous cached steps exceed this value to prevent precision degradation.
    • enable_separate_cfg: (bool, required, defaults to None):
      Whether to use separate cfg or not, such as in Wan 2.1, Qwen-Image. For models that fuse CFG and non-CFG into a single forward step, set enable_separate_cfg as False. Examples include: CogVideoX, HunyuanVideo, Mochi, etc.
    • cfg_compute_first: (bool, required, defaults to False):
      Whether to compute cfg forward first, default is False, meaning:
      0, 2, 4, ... -> non-CFG step; 1, 3, 5, ... -> CFG step.
    • cfg_diff_compute_separate: (bool, required, defaults to True):
      Whether to compute separate difference values for CFG and non-CFG steps, default is True. If False, we will use the computed difference from the current non-CFG transformer step for the current CFG step.
    • num_inference_steps (int, optional, defaults to None):
      num_inference_steps for DiffusionPipeline, used to adjust some internal settings for better caching performance. For example, we will refresh the cache once the executed steps exceed num_inference_steps if num_inference_steps is provided.
    • steps_computation_mask: (List[int], optional, defaults to None):
      This param introduce LeMiCa/EasyCache style compute mask for steps. It is a list of length num_inference_steps indicating whether to compute each step or not. 1 means must compute, 0 means use dynamic/static cache. If provided, will override other settings to decide whether to compute each step.
    • steps_computation_policy: (str, optional, defaults to "dynamic"):
      The computation policy for steps when using steps_computation_mask. It can be "dynamic" or "static". "dynamic" means using dynamic cache for steps marked as 0 in steps_computation_mask, while "static" means using static cache for those steps.
    • force_refresh_step_hint: (int, optional, defaults to None):
      The step index hint to force refresh the cache. If provided, the cache will be refreshed at the beginning of this step. This is useful for some cases where the input condition changes significantly at a certain step. Default None means no force refresh. For example, in a 50-step inference, setting force_refresh_step_hint=25 will refresh the cache before executing step 25 and view the remaining 25 steps as a new inference context.
    • force_refresh_step_policy: (str, optional, defaults to "once"):
      The policy to apply when force refreshing the cache at the step specified by
      force_refresh_step_hint. It can be "once" or "repeat". "once" means only refresh once
      at the step specified by force_refresh_step_hint, while "repeat" means refresh at the
      step specified by force_refresh_step_hint and then repeat refreshing every
      force_refresh_step_hint steps, e.g., if force_refresh_step_hint=25 and the inference
      has 100 steps, then the cache will be refreshed at:
      • once policy: step 25, treat the remaining steps as a new inference context,
        no more refresh after step 25;
      • repeat policy: step 25, 50, 75, treat the steps between each refresh as a new
        inference context.

  • calibrator_config (CalibratorConfig, optional, defaults to None):
    Config for calibrator. If calibrator_config is not None, it means the user wants to use DBCache with a specific calibrator, such as taylorseer, foca, and so on.

  • params_modifiers ('ParamsModifier', optional, defaults to None):
    Modify cache context parameters for specific blocks. The configurable parameters are listed below:

    • cache_config: (DBCacheConfig, required, defaults to DBCacheConfig()):
      The same as the 'cache_config' parameter in the cache_dit.enable_cache() interface.
    • calibrator_config: (CalibratorConfig, optional, defaults to None):
      The same as the 'calibrator_config' parameter in the cache_dit.enable_cache() interface.
    • kwargs: (dict, optional, defaults to {}):
      The same as the 'kwargs' parameter in the cache_dit.enable_cache() interface.

  • parallelism_config (ParallelismConfig, optional, defaults to None):
    Config for Parallelism. If parallelism_config is not None, it means the user wants to enable parallelism for cache-dit.

    • backend: (ParallelismBackend, required, defaults to "ParallelismBackend.AUTO"):
      Parallelism backend, currently CACHE_DIT, NATIVE_PYTORCH, and NATIVE_HYBRID are supported. For context parallelism, use the CACHE_DIT backend; for tensor parallelism, use NATIVE_PYTORCH; and for hybrid CP+TP, use NATIVE_HYBRID. Legacy NATIVE_DIFFUSER inputs are still accepted and automatically converted to CACHE_DIT for backward compatibility.
    • ulysses_size: (int, optional, defaults to None):
      The size of Ulysses cluster. If ulysses_size is not None, enable Ulysses style parallelism. This setting is only valid when backend is CACHE_DIT.
    • ring_size: (int, optional, defaults to None):
      The size of ring for ring parallelism. If ring_size is not None, enable ring attention. This setting is only valid when backend is CACHE_DIT.
    • tp_size: (int, optional, defaults to None):
      The size of tensor parallelism. If tp_size is not None, enable tensor parallelism. This setting is only valid when backend is NATIVE_PYTORCH.

  • attention_backend (str, optional, defaults to None):
    Custom attention backend in cache-dit for non-parallelism case. If attention_backend is not None, set the attention backend for the transformer module. Supported backends include: "native", "_sdpa_cudnn", "sage", "flash", "flash", "_native_npu", etc. Prefer attention_backend in parallelism_config when both are provided.

  • quantize_config (QuantizeConfig, optional, defaults to None):
    Config for quantization. If quantize_config is not None, it means the user wants to quantize the model for better performance. Supported quantization types include: float8_per_row, float8_per_tensor, float8_per_block, float8_weight_only, int8_per_row, int8_per_tensor, int8_weight_only, int4_weight_only, etc. Please check config.py for more details.

  • kwargs (dict, optional, defaults to {}):
    Other cache context keyword arguments. Please check cache_contexts/cache_context.py for more details.