[Model] Aura support - Non async chunk path (#4257)

Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
Signed-off-by: R2-Y <ruiruyang2@gmail.com>

vllm_omni/model_executor/models/aura_omni/pipeline.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""AURA Omni pipeline topology.
+
+Semantic modules:
+  1. Qwen3-ASR: microphone audio -> transcript text
+  2. AURA/Qwen3-VL: transcript + video -> text or <|silent|>
+  3. Qwen3-TTS: response text -> audio
+
+The Qwen3-TTS module is represented as two native engine stages, Talker and
+Code2Wav, to reuse the existing vLLM-Omni implementation without wrapping or
+duplicating codec streaming behavior.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_AURA_PROC = "vllm_omni.model_executor.stage_input_processors.aura_omni"
+_QWEN3_TTS_PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_tts"
+
+
+AURA_OMNI_PIPELINE = PipelineConfig(
+    model_type="aura_omni",
+    model_arch="Qwen3ASRForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="asr",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="text",
+            model_arch="Qwen3ASRForConditionalGeneration",
+            sampling_constraints={"detokenize": True},
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="aura",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="text",
+            model_arch="AuraQwen3VLForConditionalGeneration",
+            custom_process_input_func=f"{_AURA_PROC}.asr2aura",
+            sampling_constraints={"detokenize": True},
+        ),
+        StagePipelineConfig(
+            stage_id=2,
+            model_stage="qwen3_tts",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(1,),
+            owns_tokenizer=True,
+            engine_output_type="latent",
+            model_arch="Qwen3TTSTalkerForConditionalGeneration",
+            custom_process_input_func=f"{_AURA_PROC}.aura2tts",
+            custom_process_next_stage_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_full_payload",
+            async_chunk_process_next_stage_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_async_chunk",
+            sampling_constraints={
+                "detokenize": False,
+                "stop_token_ids": [2150],
+            },
+        ),
+        StagePipelineConfig(
+            stage_id=3,
+            model_stage="code2wav",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(2,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            model_arch="Qwen3TTSCode2Wav",
+            sync_process_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_token_only",
+            sampling_constraints={"detokenize": True},
+            extras={"tts_args": {"max_instructions_length": 500}},
+        ),
+    ),
+)

vllm_omni/model_executor/models/aura_omni/qwen3_vl.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""AURA-compatible Qwen3-VL wrapper.
+
+AURA ships Qwen3-VL weights with a repository-local ``configuration_qwen3_vl``
+module in ``auto_map``. With ``trust_remote_code=True``, Transformers therefore
+returns a remote ``Qwen3VLConfig`` class. The model structure is still Qwen3-VL,
+but vLLM's stock multimodal processor checks for the upstream Transformers
+``Qwen3VLConfig`` type exactly. This wrapper keeps the stock model and processor
+logic, while accepting structurally-compatible remote Qwen3-VL configs.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from vllm.model_executor.models.qwen3_vl import (
+    Qwen3VLConfig,
+    Qwen3VLDummyInputsBuilder,
+    Qwen3VLForConditionalGeneration,
+    Qwen3VLMultiModalProcessor,
+    Qwen3VLProcessingInfo,
+    Qwen3VLProcessor,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+
+
+class AuraQwen3VLProcessingInfo(Qwen3VLProcessingInfo):
+    def get_hf_config(self) -> Any:
+        try:
+            return self.ctx.get_hf_config(Qwen3VLConfig)
+        except TypeError:
+            config = self.ctx.get_hf_config()
+            if (
+                getattr(config, "model_type", None) == "qwen3_vl"
+                and hasattr(config, "vision_config")
+                and hasattr(config, "text_config")
+            ):
+                return config
+            raise
+
+    def get_hf_processor(self, **kwargs: object) -> Any:
+        # AURA's config advertises `processing_qwen3_vl.py` in auto_map, but
+        # the checkpoint snapshot does not include that file. Force the
+        # upstream processor class while still allowing AURA's remote config.
+        return self.ctx.get_hf_processor(
+            Qwen3VLProcessor,
+            use_fast=kwargs.pop("use_fast", True),
+            **kwargs,
+        )
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen3VLMultiModalProcessor,
+    info=AuraQwen3VLProcessingInfo,
+    dummy_inputs=Qwen3VLDummyInputsBuilder,
+)
+class AuraQwen3VLForConditionalGeneration(Qwen3VLForConditionalGeneration):
+    pass

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -30,7 +30,7 @@ from vllm_omni.utils.speaker_cache import (
 )
 
 from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig
-from .prompt_embeds_builder import Qwen3TTSPromptEmbedsBuilder
+from .prompt_embeds_builder import PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder
 from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM
 from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config
 from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Encoder
@@ -40,6 +40,19 @@ logger = init_logger(__name__)
 _TRAILING_TEXT_COMPACT_MIN_FRAMES = 64
 
 
+def _has_tts_text_conditioning(info_dict: dict[str, Any], hidden_states: Any | None = None) -> bool:
+    text_list = info_dict.get("text")
+    if isinstance(text_list, list) and bool(text_list) and bool(text_list[0]):
+        return True
+    if PRECOMPUTED_TEXT_IDS_KEY in info_dict:
+        return True
+    if isinstance(hidden_states, dict):
+        tail = hidden_states.get("trailing_text")
+        if isinstance(tail, torch.Tensor):
+            return True
+    return False
+
+
 # ---------------------------------------------------------------------------
 # Components ported from the HuggingFace Qwen3-TTS reference implementation.
 # Only the classes actually needed by the vLLM AR Talker are kept here.
@@ -625,9 +638,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             except Exception:
                 is_prefill = span_len > 1
 
-        text_list = info_dict.get("text")
-        if not isinstance(text_list, list) or not text_list or not text_list[0]:
-            raise ValueError("Missing additional_information.text for Qwen3-TTS AR talker.")
+        if not _has_tts_text_conditioning(info_dict, hs):
+            raise ValueError("Missing Qwen3-TTS text conditioning: provide `text` or precomputed text token ids.")
 
         task_type = (info_dict.get("task_type") or ["CustomVoice"])[0]
         codec_streaming_val = meta.get("codec_streaming")
@@ -650,7 +662,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         if is_prefill:
             # Prefill (prompt embeddings)
             prompt_embeds_cpu = embed.get("prefill")
-
             # First prefill round: prompt_embeds_cpu is not yet populated.
             # Subsequent prefill rounds (multi-chunk): prompt_embeds_cpu is a Tensor stored by the first round.
             is_first_prefill = not isinstance(prompt_embeds_cpu, torch.Tensor) or prompt_embeds_cpu.ndim != 2
@@ -808,9 +819,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             hs = payload.get("hidden_states", {})
             meta = payload.get("meta", {})
 
-            text_list = info_dict.get("text")
-            if not isinstance(text_list, list) or not text_list or not text_list[0]:
-                raise ValueError("Missing additional_information.text for Qwen3-TTS AR talker.")
+            if not _has_tts_text_conditioning(info_dict, hs):
+                raise ValueError("Missing Qwen3-TTS text conditioning: provide `text` or precomputed text token ids.")
 
             task_type = (info_dict.get("task_type") or ["CustomVoice"])[0]
             codec_streaming_val = meta.get("codec_streaming")

vllm_omni/model_executor/models/registry.py

@@ -343,6 +343,11 @@ _OMNI_MODELS = {
         "minicpmo_4_5_omni_tts",
         "MiniCPMO45OmniTTSForConditionalGeneration",
     ),
+    "AuraQwen3VLForConditionalGeneration": (
+        "aura_omni",
+        "qwen3_vl",
+        "AuraQwen3VLForConditionalGeneration",
+    ),
 }