Canlin GuoandHongsheng Liu
11c4c7f0ff
Signed-off-by: gcanlin <canlinguosdu@gmail.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
44 lines
1.4 KiB
Python
44 lines
1.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Global diffusion config — mirrors vLLM's set_current_vllm_config pattern.
|
|
|
|
Set once during model initialisation so that every ``Attention`` layer can
|
|
read the ``OmniDiffusionConfig`` without coupling ``__init__`` to the
|
|
runtime ``ForwardContext``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
_current_diffusion_config: OmniDiffusionConfig | None = None
|
|
|
|
|
|
@contextmanager
|
|
def set_current_diffusion_config(config: OmniDiffusionConfig):
|
|
"""Context manager that sets the global diffusion config for model init."""
|
|
global _current_diffusion_config
|
|
old = _current_diffusion_config
|
|
_current_diffusion_config = config
|
|
try:
|
|
yield
|
|
finally:
|
|
_current_diffusion_config = old
|
|
|
|
|
|
def get_current_diffusion_config() -> OmniDiffusionConfig:
|
|
"""Return the current diffusion config or raise."""
|
|
assert _current_diffusion_config is not None, (
|
|
"Diffusion config is not set. Wrap model construction with set_current_diffusion_config()."
|
|
)
|
|
return _current_diffusion_config
|
|
|
|
|
|
def get_current_diffusion_config_or_none() -> OmniDiffusionConfig | None:
|
|
"""Return the current diffusion config, or ``None`` if not set."""
|
|
return _current_diffusion_config
|