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

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

tests/model_executor/models/test_aura_qwen3_vl.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.model_executor.models.aura_omni.qwen3_vl import (
+    AuraQwen3VLForConditionalGeneration,
+    AuraQwen3VLProcessingInfo,
+    Qwen3VLProcessor,
+)
+from vllm_omni.model_executor.models.registry import OmniModelRegistry
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeCtx:
+    def __init__(self, config):
+        self.config = config
+        self.processor_type = None
+
+    def get_hf_config(self, typ=None):
+        if typ is not None:
+            raise TypeError("wrong config class")
+        return self.config
+
+    def get_hf_processor(self, typ=None, **kwargs):
+        self.processor_type = typ
+        return SimpleNamespace(kwargs=kwargs)
+
+
+def test_aura_qwen3_vl_processing_accepts_remote_qwen3_vl_config():
+    info = object.__new__(AuraQwen3VLProcessingInfo)
+    info.ctx = _FakeCtx(
+        SimpleNamespace(
+            model_type="qwen3_vl",
+            vision_config=SimpleNamespace(spatial_merge_size=2),
+            text_config=SimpleNamespace(hidden_size=4096),
+        )
+    )
+
+    assert info.get_hf_config().model_type == "qwen3_vl"
+
+
+def test_aura_qwen3_vl_processing_rejects_unrelated_remote_config():
+    info = object.__new__(AuraQwen3VLProcessingInfo)
+    info.ctx = _FakeCtx(SimpleNamespace(model_type="not_qwen3_vl"))
+
+    with pytest.raises(TypeError, match="wrong config class"):
+        info.get_hf_config()
+
+
+def test_aura_qwen3_vl_model_arch_registered():
+    model_cls = OmniModelRegistry._try_load_model_cls("AuraQwen3VLForConditionalGeneration")
+
+    assert model_cls is AuraQwen3VLForConditionalGeneration
+
+
+def test_aura_qwen3_vl_processing_forces_upstream_processor_class():
+    info = object.__new__(AuraQwen3VLProcessingInfo)
+    ctx = _FakeCtx(SimpleNamespace(model_type="qwen3_vl"))
+    info.ctx = ctx
+
+    processor = info.get_hf_processor(extra="value")
+
+    assert ctx.processor_type is Qwen3VLProcessor
+    assert processor.kwargs["use_fast"] is True
+    assert processor.kwargs["extra"] == "value"

tests/model_executor/stage_input_processors/test_aura_omni.py

@@ -0,0 +1,197 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+    PRECOMPUTED_TEXT_IDS_KEY,
+)
+from vllm_omni.model_executor.stage_input_processors.aura_omni import (
+    SILENT_TEXT,
+    asr2aura,
+    aura2tts,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _source_output(text: str, request_id: str = "req-1", token_ids: list[int] | None = None):
+    output = SimpleNamespace(text=text, cumulative_token_ids=token_ids or [1, 2, 3], multimodal_output={})
+    return SimpleNamespace(request_id=request_id, outputs=[output])
+
+
+def _source_delta_final_output(cumulative_text: str, request_id: str = "req-1"):
+    output = SimpleNamespace(
+        text="",
+        cumulative_text=cumulative_text,
+        cumulative_token_ids=[1, 2, 3],
+        multimodal_output={},
+    )
+    return SimpleNamespace(request_id=request_id, outputs=[output])
+
+
+def test_asr2aura_carries_video_payload_and_transcript():
+    prompt = {
+        "multi_modal_data": {"video": ["frame-0", "frame-1"]},
+        "additional_information": {"aura_system_prompt": ["system"]},
+    }
+
+    [next_input] = asr2aura([_source_output("What is happening now?")], prompt=[prompt])
+
+    assert next_input["multi_modal_data"] == {"video": ["frame-0", "frame-1"]}
+    assert "<|video_pad|>" in next_input["prompt"]
+    assert "What is happening now?" in next_input["prompt"]
+    assert next_input["prompt"].startswith("<|im_start|>system\nsystem")
+
+
+def test_asr2aura_drops_audio_before_qwen3_vl_stage():
+    prompt = {
+        "multi_modal_data": {
+            "audio": ("wave", 16000),
+            "video": ["frame-0", "frame-1"],
+        },
+    }
+
+    [next_input] = asr2aura([_source_output("Check the video")], prompt=[prompt])
+
+    assert next_input["multi_modal_data"] == {"video": ["frame-0", "frame-1"]}
+    assert "<|video_pad|>" in next_input["prompt"]
+
+
+def test_asr2aura_reads_video_stashed_for_downstream_stage():
+    prompt = {
+        "multi_modal_data": {"audio": ("wave", 16000)},
+        "additional_information": {
+            "deferred_multi_modal_data": {"video": ["frame-0", "frame-1"]},
+        },
+    }
+
+    [next_input] = asr2aura([_source_output("Check the video")], prompt=[prompt])
+
+    assert next_input["multi_modal_data"] == {"video": ["frame-0", "frame-1"]}
+    assert "<|video_pad|>" in next_input["prompt"]
+
+
+def test_asr2aura_supports_video_only_observation():
+    prompt = {"multi_modal_data": {"video": ["frame-0", "frame-1"]}}
+
+    [next_input] = asr2aura([_source_output("")], prompt=[prompt])
+
+    assert "<|video_pad|>" in next_input["prompt"]
+    assert "<|im_start|>assistant" in next_input["prompt"]
+
+
+def test_aura2tts_builds_qwen3_tts_prompt_information():
+    prompt = {
+        "additional_information": {
+            "tts_language": ["Chinese"],
+            "tts_instruct": ["Calm voice."],
+            "tts_ref_audio": ["ref.wav"],
+            "tts_ref_text": ["Reference transcript sample."],
+        }
+    }
+
+    [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt])
+
+    assert len(tts_input["prompt_token_ids"]) >= 32
+    assert tts_input["additional_information"]["text"] == ["Hello."]
+    assert PRECOMPUTED_TEXT_IDS_KEY not in tts_input["additional_information"]
+    assert tts_input["additional_information"]["task_type"] == ["Base"]
+    assert tts_input["additional_information"]["language"] == ["Chinese"]
+    assert tts_input["additional_information"]["ref_audio"] == ["ref.wav"]
+    assert tts_input["additional_information"]["ref_text"] == ["Reference transcript sample."]
+    assert tts_input["additional_information"]["x_vector_only_mode"] == [False]
+    assert tts_input["additional_information"]["instruct"] == ["Calm voice."]
+
+
+def test_aura2tts_prefers_streaming_cumulative_text():
+    prompt = {
+        "additional_information": {
+            "tts_ref_audio": ["ref.wav"],
+            "tts_ref_text": ["Reference transcript sample."],
+        }
+    }
+
+    [tts_input] = aura2tts(
+        [_source_delta_final_output("The complete AURA reply.")],
+        prompt=[prompt],
+    )
+
+    assert tts_input["additional_information"]["text"] == ["The complete AURA reply."]
+
+
+def test_aura2tts_supports_base_ref_audio_override():
+    prompt = {
+        "additional_information": {
+            "tts_ref_audio": ["custom.wav"],
+            "tts_ref_text": ["custom transcript"],
+        }
+    }
+
+    [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt])
+
+    assert tts_input["additional_information"]["task_type"] == ["Base"]
+    assert tts_input["additional_information"]["ref_audio"] == ["custom.wav"]
+    assert tts_input["additional_information"]["ref_text"] == ["custom transcript"]
+    assert tts_input["additional_information"]["x_vector_only_mode"] == [False]
+
+
+def test_aura2tts_supports_x_vector_only_mode_for_base():
+    prompt = {
+        "additional_information": {
+            "tts_task_type": ["Base"],
+            "tts_x_vector_only_mode": [True],
+            "tts_ref_audio": ["ref.wav"],
+            "tts_ref_text": ["Reference transcript sample."],
+        }
+    }
+
+    [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt])
+
+    assert tts_input["additional_information"]["x_vector_only_mode"] == [True]
+
+
+def test_aura2tts_supports_custom_voice_mode():
+    prompt = {
+        "additional_information": {
+            "tts_task_type": ["CustomVoice"],
+            "tts_speaker": ["vivian"],
+        }
+    }
+
+    [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt])
+
+    assert tts_input["additional_information"]["task_type"] == ["CustomVoice"]
+    assert tts_input["additional_information"]["speaker"] == ["Vivian"]
+    assert "ref_audio" not in tts_input["additional_information"]
+
+
+def test_aura2tts_passes_token_ids_to_qwen3_tts_when_enabled():
+    prompt = {
+        "additional_information": {
+            "tts_ref_audio": ["ref.wav"],
+            "tts_ref_text": ["Reference transcript sample."],
+            "tts_pass_token_ids": [True],
+        }
+    }
+
+    [tts_input] = aura2tts(
+        [
+            _source_output(
+                "Hello.",
+                token_ids=[151644, 77091, 198, 108386, 1773, 151645, 198],
+            )
+        ],
+        prompt=[prompt],
+    )
+
+    assert tts_input["additional_information"][PRECOMPUTED_TEXT_IDS_KEY] == [
+        [151644, 77091, 198, 108386, 1773, 151645, 198, 151644, 77091, 198]
+    ]
+    assert "text" not in tts_input["additional_information"]
+
+
+def test_aura2tts_drops_silent_response():
+    assert aura2tts([_source_output(SILENT_TEXT)]) == []