Skip to content

vllm_omni.diffusion.cache.stepcache.backend

StepCache backend implementation.

This module provides the stepcache backend that implements the CacheBackend interface for DreamZero-style velocity-based step skipping in the denoise loop.

CUSTOM_STEPCACHE_ENABLERS module-attribute 

CUSTOM_STEPCACHE_ENABLERS = {
    "DreamZeroPipeline": enable_dreamzero_stepcache
}

CUSTOM_STEP_CACHE_DIT_ENABLERS module-attribute 

CUSTOM_STEP_CACHE_DIT_ENABLERS = CUSTOM_STEPCACHE_ENABLERS

STEP_CACHE_DIT_CONFIG_ATTR module-attribute 

STEP_CACHE_DIT_CONFIG_ATTR = STEP_CACHE_CONFIG_ATTR

STEP_CACHE_DIT_STATE_ATTR module-attribute 

STEP_CACHE_DIT_STATE_ATTR = STEP_CACHE_STATE_ATTR

StepCacheDiTBackend module-attribute 

StepCacheDiTBackend = StepCacheBackend

StepCacheDiTConfig module-attribute 

StepCacheDiTConfig = StepCacheConfig

StepCacheDiTState module-attribute 

StepCacheDiTState = StepCacheState

enable_dreamzero_step_cache_dit module-attribute 

enable_dreamzero_step_cache_dit = enable_dreamzero_stepcache

logger module-attribute 

logger = init_logger(__name__)

StepCacheBackend

Bases: CacheBackend

Velocity cosine step-skipping cache backend for DreamZero.

Attaches :class:StepCacheConfig and :class:StepCacheState to supported pipelines. The denoise loop calls :meth:StepCacheState.should_run_step to decide whether to run predict_noise.

Example

from vllm_omni.diffusion.data import DiffusionCacheConfig backend = StepCacheBackend(DiffusionCacheConfig()) backend.enable(pipeline) backend.refresh(pipeline, num_inference_steps=16)

enable 

enable(pipeline: Any) -> None

Enable stepcache on a supported pipeline.

refresh 

refresh(
    pipeline: Any,
    num_inference_steps: int,
    verbose: bool = True,
) -> None

Refresh stepcache state for a new generation.

enable_dreamzero_stepcache 

enable_dreamzero_stepcache(
    pipeline: Any, config: DiffusionCacheConfig
) -> None

Enable stepcache for DreamZeroPipeline.

get_step_cache_dit_config 

get_step_cache_dit_config(
    pipeline: Any,
) -> StepCacheConfig | None

get_step_cache_dit_state 

get_step_cache_dit_state(
    pipeline: Any,
) -> StepCacheState | None

get_stepcache_config 

get_stepcache_config(
    pipeline: Any,
) -> StepCacheConfig | None

get_stepcache_state 

get_stepcache_state(pipeline: Any) -> StepCacheState | None

is_step_cache_dit_active 

is_step_cache_dit_active(pipeline: Any) -> bool

is_stepcache_active 

is_stepcache_active(pipeline: Any) -> bool

Return True when stepcache is enabled on pipeline.