28 lines
1.1 KiB
Python
28 lines
1.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
# unify OmniLLM and OmniDiffusion entrypoints
|
|
|
|
from vllm_omni.diffusion.omni_diffusion import OmniDiffusion
|
|
from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
|
|
from vllm_omni.entrypoints.omni_llm import OmniLLM
|
|
|
|
|
|
class Omni:
|
|
def __init__(self, *args, **kwargs):
|
|
model = args[0] if args else kwargs.get("model", "")
|
|
if is_diffusion_model(model):
|
|
self.instance: OmniLLM | OmniDiffusion = OmniDiffusion(*args, **kwargs)
|
|
else:
|
|
self.instance: OmniLLM | OmniDiffusion = OmniLLM(*args, **kwargs)
|
|
|
|
def __getattr__(self, name):
|
|
"""Delegate attribute access to the chosen backend instance."""
|
|
return getattr(self.instance, name)
|
|
|
|
def generate(self, *args, **kwargs):
|
|
"""Convenience wrapper to call `generate` on the backend if available."""
|
|
if hasattr(self.instance, "generate"):
|
|
return getattr(self.instance, "generate")(*args, **kwargs)
|
|
raise AttributeError(f"'{self.instance.__class__.__name__}' has no attribute 'generate'")
|