[Misc] Temporarily support downloading models from ModelScope by snapshot download (#132)

Signed-off-by: MengqingCao <cmq0113@163.com>

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:

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: