diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 5efc001da..e5ac27aef 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -1,16 +1,38 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os 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 +def _dummy_snapshot_download(model_id): + return model_id + + +def omni_snapshot_download(model_id) -> str: + # TODO: this is just a workaround for quickly use modelscope, we should support + # modelscope in weight loading feature instead of using `snapshot_download` + if os.environ.get("VLLM_USE_MODELSCOPE", False): + from modelscope.hub.snapshot_download import snapshot_download + + return snapshot_download(model_id) + else: + return _dummy_snapshot_download(model_id) + + class Omni: """Unified entrypoint for both LLM and Diffusion models for better usability.""" def __init__(self, *args, **kwargs): model = args[0] if args else kwargs.get("model", "") + assert model != "", "Null model id detected, please specify a model id." + model = omni_snapshot_download(model) + if args: + args[0] = model + elif kwargs.get("model", "") != "": + kwargs["model"] = model if is_diffusion_model(model): self.instance: OmniLLM | OmniDiffusion = OmniDiffusion(*args, **kwargs) else: diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index a0028d572..6cd344de8 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -15,6 +15,9 @@ from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding from vllm_omni.request import OmniRequest for module_name, module in sys.modules.items(): + # only do patch on module of vllm, pass others + if "vllm" not in module_name: + continue if hasattr(module, "EngineCoreOutput") and module.EngineCoreOutput == _OriginalEngineCoreOutput: module.EngineCoreOutput = OmniEngineCoreOutput if hasattr(module, "EngineCoreOutputs") and module.EngineCoreOutputs == _OriginalEngineCoreOutputs: