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>
83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
Unified cache setup function for diffusion models.
|
|
|
|
This module provides the main entry point for setting up cache adapters
|
|
on transformer models. It handles cache type selection, adapter instantiation,
|
|
and application to the model.
|
|
"""
|
|
|
|
from typing import Any, Optional
|
|
|
|
import torch
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.cache.base import CacheAdapter
|
|
from vllm_omni.diffusion.cache.registry import get_cache_adapter
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def setup_cache(
|
|
transformer: torch.nn.Module,
|
|
cache_type: str = "none",
|
|
cache_config: Optional[dict[str, Any]] = None,
|
|
) -> Optional[CacheAdapter]:
|
|
"""
|
|
Setup cache adapter for transformer (one-time initialization).
|
|
|
|
This is the main entry point for applying caching to diffusion transformers.
|
|
It's called once during pipeline initialization and applies the appropriate
|
|
cache adapter based on configuration.
|
|
|
|
Args:
|
|
transformer: Transformer module to apply cache to
|
|
cache_type: Type of cache to apply. Options:
|
|
- "none": No caching (default)
|
|
- "tea_cache": TeaCache adaptive caching
|
|
- Additional types can be registered via register_cache_adapter()
|
|
cache_config: Cache-specific configuration dictionary. Contents depend
|
|
on cache type:
|
|
- For "tea_cache": {"rel_l1_thresh": 0.2, "model_type": "QwenImagePipeline"}
|
|
- For other types: see their respective adapter documentation
|
|
|
|
Returns:
|
|
CacheAdapter instance for state management, or None if cache_type="none"
|
|
|
|
Example:
|
|
>>> # Setup TeaCache during pipeline initialization
|
|
>>> adapter = setup_cache(
|
|
... transformer,
|
|
... cache_type="tea_cache",
|
|
... cache_config={"rel_l1_thresh": 0.2}
|
|
... )
|
|
>>>
|
|
>>> # Reset cache state before each generation
|
|
>>> if adapter is not None:
|
|
... adapter.reset(transformer)
|
|
"""
|
|
# No caching requested
|
|
if cache_type == "none":
|
|
logger.debug("No cache adapter requested (cache_type='none')")
|
|
return None
|
|
|
|
# Get adapter instance from registry
|
|
try:
|
|
adapter = get_cache_adapter(cache_type, cache_config or {})
|
|
except ValueError as e:
|
|
logger.error(f"Failed to get cache adapter: {e}")
|
|
raise
|
|
|
|
# Apply cache to transformer
|
|
logger.info(f"Applying {cache_type} cache to transformer")
|
|
try:
|
|
adapter.apply(transformer)
|
|
except Exception as e:
|
|
logger.error(f"Failed to apply cache adapter: {e}")
|
|
raise
|
|
|
|
logger.info(f"Cache adapter '{cache_type}' successfully applied")
|
|
return adapter
|