Skip to content

vllm_omni.diffusion.cache.cachedit.backend

Generic Cache-DiT backend lifecycle and integration.

CUSTOM_DIT_ENABLERS module-attribute 

CUSTOM_DIT_ENABLERS: dict[str, CacheDiTEnabler] = {}

CacheDiTEnabler module-attribute 

CacheDiTEnabler: TypeAlias = Callable[
    [Any, DiffusionCacheConfig],
    RefreshCacheContextFunc | CacheDiTEnableResult,
]

RefreshCacheContextFunc module-attribute 

RefreshCacheContextFunc: TypeAlias = Callable[
    [Any, int, bool], None
]

logger module-attribute 

logger = init_logger(__name__)

CacheDiTBackend

Bases: CacheBackend

Manage Cache-DiT through the common diffusion cache lifecycle.

disable 

disable(pipeline: SupportsComponentDiscovery) -> None

Remove Cache-DiT hooks so later requests use native forwards.

enable 

enable(pipeline: SupportsComponentDiscovery) -> None

refresh 

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

CacheDiTEnableResult dataclass

Custom-enabler result with the exact targets needed for teardown.

refresh instance-attribute 

refresh: RefreshCacheContextFunc

targets instance-attribute 

targets: tuple[Any, ...]

cache_summary 

cache_summary(
    pipeline: SupportsComponentDiscovery,
    details: bool = True,
) -> None

Log Cache-DiT statistics for every transformer on the pipeline.

enable_cache_for_dit 

enable_cache_for_dit(
    pipeline: SupportsComponentDiscovery,
    cache_config: DiffusionCacheConfig,
    block_adapter: BlockAdapter | None = None,
    adapter_cls: type[CachedAdapter] | None = None,
    get_pipeline_transformer: Callable[
        [SupportsComponentDiscovery], Module
    ] = _default_get_pipeline_transformer,
) -> RefreshCacheContextFunc

Enable Cache-DiT for a standard single-transformer DiT pipeline.