[Misc] Fix type interface mismatch (#478)

Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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

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: