From abde501dfa301b1eb838bf97f8d8e281fee82826 Mon Sep 17 00:00:00 2001 From: dsinghvi Date: Fri, 26 Dec 2025 12:58:49 +0530 Subject: [PATCH] [Misc] Fix type interface mismatch (#478) Signed-off-by: Divyansh Singhvi Co-authored-by: Hongsheng Liu --- vllm_omni/worker/gpu_model_runner.py | 4 +++- vllm_omni/worker/npu/npu_model_runner.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 4ca0a6a0e..e426e3938 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -16,6 +16,8 @@ from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.worker.gpu_input_batch import CachedRequestState from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata +from vllm_omni.model_executor.models.output_templates import OmniOutput + if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput else: @@ -308,7 +310,7 @@ class OmniGPUModelRunner(GPUModelRunner): self.input_batch.refresh_metadata() @torch.inference_mode() - def extract_multimodal_outputs(self, hidden_states: torch.Tensor | list[torch.Tensor]) -> dict: + def extract_multimodal_outputs(self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput) -> dict: if hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs: text_hidden_states = hidden_states.text_hidden_states multimodal_outputs = hidden_states.multimodal_outputs diff --git a/vllm_omni/worker/npu/npu_model_runner.py b/vllm_omni/worker/npu/npu_model_runner.py index e1fb9a7a7..b6f12d84f 100644 --- a/vllm_omni/worker/npu/npu_model_runner.py +++ b/vllm_omni/worker/npu/npu_model_runner.py @@ -21,6 +21,8 @@ from vllm_ascend.ascend_forward_context import set_ascend_forward_context from vllm_ascend.utils import enable_sp, lmhead_tp_enable from vllm_ascend.worker.model_runner_v1 import NPUModelRunner +from vllm_omni.model_executor.models.output_templates import OmniOutput + if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput @@ -300,7 +302,7 @@ class OmniNPUModelRunner(NPUModelRunner): @torch.inference_mode() def extract_multimodal_outputs( - self, hidden_states: torch.Tensor | list[torch.Tensor] + self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput ) -> tuple[torch.Tensor, torch.Tensor | list[torch.Tensor] | dict]: """Extract multimodal outputs from hidden states.""" if hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs: