From b06a5148f6b89b1036b9c44e6acce6469ed9f0b4 Mon Sep 17 00:00:00 2001 From: dsinghvi Date: Mon, 22 Dec 2025 13:20:26 +0530 Subject: [PATCH] bug resolution for omni models registration not present in vllm (#397) Signed-off-by: Divyansh Singhvi --- vllm_omni/engine/arg_utils.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index d3e2268ea..57324374f 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -10,6 +10,17 @@ from vllm_omni.config import OmniModelConfig logger = init_logger(__name__) +def register_omni_models_to_vllm(): + from vllm.model_executor.models import ModelRegistry + + from vllm_omni.model_executor.models.registry import _OMNI_MODELS + + supported_archs = ModelRegistry.get_supported_archs() + for arch, (mod_folder, mod_relname, cls_name) in _OMNI_MODELS.items(): + if arch not in supported_archs: + ModelRegistry.register_model(arch, f"vllm_omni.model_executor.models.{mod_folder}.{mod_relname}:{cls_name}") + + @dataclass class OmniEngineArgs(EngineArgs): """Engine arguments for omni models, extending base EngineArgs. @@ -38,11 +49,20 @@ class OmniEngineArgs(EngineArgs): # we need to draw the text config from the corresponding model stage. return getattr(config_dict["hf_config"], config_dict["hf_config_name"]).get_text_config() + def _ensure_omni_models_registered(self): + if hasattr(self, "_omni_models_registered"): + return True + register_omni_models_to_vllm() + self._omni_models_registered = True + return True + def create_model_config(self) -> OmniModelConfig: """Create an OmniModelConfig from these engine arguments. Returns: OmniModelConfig instance with all configuration fields set """ + # register omni models to avoid model not found error + self._ensure_omni_models_registered() # First, get the base ModelConfig from the parent class base_config = super().create_model_config() @@ -99,7 +119,16 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): # we need to draw the text config from the corresponding model stage. return getattr(config_dict["hf_config"], config_dict["hf_config_name"]).get_text_config() + def _ensure_omni_models_registered(self): + if hasattr(self, "_omni_models_registered"): + return True + register_omni_models_to_vllm() + self._omni_models_registered = True + return True + def create_model_config(self) -> OmniModelConfig: + # register omni models to avoid model not found error + self._ensure_omni_models_registered() # First, get the base ModelConfig from the parent class base_config = super().create_model_config()