67 lines
1.8 KiB
Python
67 lines
1.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
StepCache: velocity-based step skipping for DreamZero-style DiT denoising.
|
|
|
|
Unlike block-level Cache-DiT, this backend skips entire DiT forwards during the
|
|
pipeline denoise loop when successive velocity predictions are highly aligned.
|
|
|
|
Usage:
|
|
from vllm_omni import Omni
|
|
|
|
omni = Omni(
|
|
model="...",
|
|
cache_backend="step_cache",
|
|
)
|
|
|
|
# Alternative: environment variable
|
|
# export DIFFUSION_CACHE_BACKEND=step_cache
|
|
"""
|
|
|
|
from vllm_omni.diffusion.cache.stepcache.backend import (
|
|
CUSTOM_STEP_CACHE_DIT_ENABLERS,
|
|
CUSTOM_STEPCACHE_ENABLERS,
|
|
STEP_CACHE_CONFIG_ATTR,
|
|
STEP_CACHE_DIT_CONFIG_ATTR,
|
|
STEP_CACHE_DIT_STATE_ATTR,
|
|
STEP_CACHE_STATE_ATTR,
|
|
StepCacheBackend,
|
|
StepCacheDiTBackend,
|
|
StepCacheDiTConfig,
|
|
StepCacheDiTState,
|
|
enable_dreamzero_step_cache_dit,
|
|
enable_dreamzero_stepcache,
|
|
get_step_cache_dit_config,
|
|
get_step_cache_dit_state,
|
|
get_stepcache_config,
|
|
get_stepcache_state,
|
|
is_step_cache_dit_active,
|
|
is_stepcache_active,
|
|
)
|
|
from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig
|
|
from vllm_omni.diffusion.cache.stepcache.state import StepCacheState
|
|
|
|
__all__ = [
|
|
"CUSTOM_STEP_CACHE_DIT_ENABLERS",
|
|
"CUSTOM_STEPCACHE_ENABLERS",
|
|
"STEP_CACHE_CONFIG_ATTR",
|
|
"STEP_CACHE_DIT_CONFIG_ATTR",
|
|
"STEP_CACHE_DIT_STATE_ATTR",
|
|
"STEP_CACHE_STATE_ATTR",
|
|
"StepCacheBackend",
|
|
"StepCacheConfig",
|
|
"StepCacheDiTBackend",
|
|
"StepCacheDiTConfig",
|
|
"StepCacheDiTState",
|
|
"StepCacheState",
|
|
"enable_dreamzero_step_cache_dit",
|
|
"enable_dreamzero_stepcache",
|
|
"get_step_cache_dit_config",
|
|
"get_step_cache_dit_state",
|
|
"get_stepcache_config",
|
|
"get_stepcache_state",
|
|
"is_step_cache_dit_active",
|
|
"is_stepcache_active",
|
|
]
|