535bd35b42
Signed-off-by: lishunyang12 <lishunyang12@users.noreply.github.com> Signed-off-by: Armaan Amatya <armaanamatya2014@gmail.com> Co-authored-by: lishunyang12 <lishunyang12@users.noreply.github.com> Co-authored-by: Armaan Amatya <38102001+armaanamatya@users.noreply.github.com>
41 lines
1.3 KiB
Python
41 lines
1.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""
|
|
Cache module for diffusion model inference acceleration.
|
|
|
|
This module provides a unified cache backend system for different caching strategies:
|
|
- TeaCache: Timestep Embedding Aware Cache for adaptive transformer caching
|
|
- MagCache: Magnitude-based Cache for adaptive transformer caching
|
|
- cache-dit: DBCache, SCM, and TaylorSeer caching strategies
|
|
|
|
Cache backends are instantiated directly via their constructors and configured via OmniDiffusionConfig.
|
|
"""
|
|
|
|
from vllm_omni.diffusion.cache.base import CacheBackend
|
|
from vllm_omni.diffusion.cache.exact_projection_cache import ExactProjectionCache
|
|
from vllm_omni.diffusion.cache.prompt_embed_cache import (
|
|
PromptEmbedCache,
|
|
install_prompt_embed_cache,
|
|
resolve_prompt_embed_cache_config,
|
|
uninstall_prompt_embed_cache,
|
|
)
|
|
from vllm_omni.diffusion.cache.teacache import (
|
|
CacheContext,
|
|
TeaCacheConfig,
|
|
apply_teacache_hook,
|
|
)
|
|
from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend
|
|
|
|
__all__ = [
|
|
"CacheBackend",
|
|
"ExactProjectionCache",
|
|
"CacheContext",
|
|
"TeaCacheBackend",
|
|
"TeaCacheConfig",
|
|
"apply_teacache_hook",
|
|
"PromptEmbedCache",
|
|
"install_prompt_embed_cache",
|
|
"resolve_prompt_embed_cache_config",
|
|
"uninstall_prompt_embed_cache",
|
|
]
|