65ca13102d
Signed-off-by: Prajwal A <prajwalanagani@gmail.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: Jiangyun Zhu <riverclouds.zhu@qq.com> Signed-off-by: Samit <285365963@qq.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Samit <285365963@qq.com>
42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
Cache module for diffusion model inference acceleration.
|
|
|
|
This module provides a unified cache adapter system for different caching strategies:
|
|
- TeaCache: Timestep Embedding Aware Cache for adaptive transformer caching
|
|
- (Future: DeepCache, DistriFusion, etc.)
|
|
|
|
The cache adapter system uses a registry pattern where cache types can be
|
|
dynamically registered and instantiated. Cache is configured via OmniDiffusionConfig
|
|
or DIFFUSION_CACHE_ADAPTER environment variable.
|
|
"""
|
|
|
|
from vllm_omni.diffusion.cache.apply import setup_cache
|
|
from vllm_omni.diffusion.cache.base import CacheAdapter
|
|
from vllm_omni.diffusion.cache.registry import (
|
|
CacheType,
|
|
get_cache_adapter,
|
|
register_cache_adapter,
|
|
)
|
|
from vllm_omni.diffusion.cache.teacache import (
|
|
CacheContext,
|
|
TeaCacheConfig,
|
|
apply_teacache_hook,
|
|
)
|
|
|
|
# Import teacache adapter to trigger registration
|
|
from vllm_omni.diffusion.cache.teacache.adapter import TeaCacheAdapter # noqa: F401
|
|
|
|
__all__ = [
|
|
"CacheAdapter",
|
|
"CacheType",
|
|
"TeaCacheConfig",
|
|
"CacheContext",
|
|
"setup_cache",
|
|
"get_cache_adapter",
|
|
"register_cache_adapter",
|
|
"apply_teacache_hook",
|
|
]
|