Refactor model mode initialization and enhance error handling in AsyncOmni. Update worker type retrieval in omni_stage and add audio token handling in Qwen3OmniMoe classes for improved robustness.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -39,7 +39,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.scheduler_max_num_seqs = vllm_config.scheduler_config.max_num_seqs
         self.connector = self.create_connector(model_config)
         super().__init__(model_config)
-        self.model_mode = getattr(model_config, "worker_type", "ar")
+        self.model_mode = getattr(model_config, "worker_type", None) or "ar"
         # State specific to Chunk management
         self.custom_process_next_stage_input_func = None
         custom_process_next_stage_input_func = getattr(model_config, "custom_process_next_stage_input_func", None)

vllm_omni/entrypoints/async_omni.py

@@ -401,7 +401,10 @@ class AsyncOmni(OmniBase):
     ) -> AsyncGenerator[OmniRequestOutput, None]:
         all_stages_finished = {stage_id: False for stage_id in range(final_stage_id_for_e2e + 1)}
         submit_flag = True
+        _loop_iter = 0
+        _last_progress_ts = time.time()
         while not all(all_stages_finished.values()):
+            _loop_iter += 1
             for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]):
                 if all_stages_finished[stage_id]:
                     continue
@@ -410,6 +413,7 @@ class AsyncOmni(OmniBase):
                 except asyncio.QueueEmpty:
                     await asyncio.sleep(0.001)
                     continue
+                _last_progress_ts = time.time()
                 engine_outputs, finished, output_to_yield = self._process_single_result(
                     result,
                     stage,
@@ -418,11 +422,20 @@ class AsyncOmni(OmniBase):
                 )
                 if submit_flag and stage_id == 0:
                     submit_flag = False
-                    prompt_token_ids = engine_outputs.prompt_token_ids
+                    prompt_token_ids = getattr(engine_outputs, "prompt_token_ids", None)
+                    if prompt_token_ids is None:
+                        prompt_token_ids = []
                     engine_input = copy.deepcopy(prompt)
-                    next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
+                    try:
+                        next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
+                    except Exception:
+                        raise
                     engine_input["prompt_token_ids"] = [0] * next_prompt_len
                     engine_input["multi_modal_data"] = engine_input["mm_processor_kwargs"] = None
+                    for _mm_key in ("mm_kwargs", "mm_hashes", "mm_placeholders", "multi_modal_uuids"):
+                        engine_input.pop(_mm_key, None)
+                    if engine_input.get("type") == "multimodal":
+                        engine_input["type"] = "token"
                     for i in range(1, len(self.stage_list)):
                         task = {
                             "request_id": request_id,
@@ -709,6 +722,8 @@ class AsyncOmni(OmniBase):
         OMNI: Required by upstream OpenAIServingModels.__init__ which
         accesses engine_client.renderer.
         """
+        if self.input_processor is None:
+            return None
         return self.input_processor.renderer
 
     async def do_log_stats(self) -> None:

vllm_omni/entrypoints/omni_stage.py

@@ -207,7 +207,7 @@ def _sequential_init_lock(engine_args: dict[str, Any], stage_init_timeout: int =
 
 
 def _resolve_worker_cls(engine_args: dict[str, Any]) -> None:
-    worker_type = engine_args.pop("worker_type", None)
+    worker_type = engine_args.get("worker_type", None)
     if not worker_type:
         return
     worker_cls = engine_args.get("worker_cls")

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -340,6 +340,11 @@ class Qwen3OmniMoeThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProc
                 seq_len=seq_len,
                 mm_counts=mm_counts,
             )
+            if audio_tokens is None:
+                feature_extractor = self.get_feature_extractor()
+                max_audio_samples = feature_extractor.chunk_length * feature_extractor.sampling_rate
+                max_audio_tokens = int(_get_feat_extract_output_lengths(torch.tensor([max_audio_samples])).item())
+                audio_tokens = {"audio": max_audio_tokens}
             mm_max_tokens["audio"] = audio_tokens["audio"]
 
         return mm_max_tokens
@@ -769,8 +774,7 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix
             feature_lens=audio_feature_lengths,
             aftercnn_lens=audio_output_lengths,
         )
-        # OMNI: audio_tower.forward() returns hidden_states tensor directly
-        audio_features = audio_outputs
+        audio_features = audio_outputs if isinstance(audio_outputs, torch.Tensor) else audio_outputs.last_hidden_state
         return audio_features.split(audio_output_lengths.tolist())