[Bugfix] Use Seed-TTS reference audio with MiniCPM-o 4.5 (#6628)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

tests/model_executor/models/minicpmo_4_5/test_llm2tts.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for the MiniCPM-o 4.5 stage 0 -> stage 1 bridge.
 
 Covers ``vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.llm2tts``:
@@ -25,6 +25,7 @@ import torch
 from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
     MiniCPMO45OmniForConditionalGeneration,
 )
+from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO45_REFERENCE_AUDIO_KEY
 from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import llm2tts
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -313,6 +314,44 @@ class TestPromptAndMultiModal:
         )
         assert out[0]["multi_modal_data"] == mm
 
+    def test_reference_audio_is_added_to_stage_handoff(self) -> None:
+        hidden_states = torch.zeros((2, _HIDDEN_DIM))
+        waveform = torch.tensor([0.25, -0.5], dtype=torch.float32)
+
+        stage_outputs = llm2tts(
+            [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden_states)],
+            prompt={"multi_modal_data": {"audio": (waveform, 16000)}},
+        )
+
+        intermediate_buffer = stage_outputs[0]["model_intermediate_buffer"]
+        assert intermediate_buffer["codes"]["ref"] == [0.25, -0.5]
+        assert intermediate_buffer["meta"]["ref_audio_sr"] == 16000
+
+    def test_serving_reference_audio_is_added_to_stage_handoff(self) -> None:
+        hidden_states = torch.zeros((2, _HIDDEN_DIM))
+        waveform = torch.tensor([0.125, -0.25], dtype=torch.float32)
+
+        stage_outputs = llm2tts(
+            [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden_states)],
+            prompt={MINICPMO45_REFERENCE_AUDIO_KEY: (waveform, 24000)},
+        )
+
+        intermediate_buffer = stage_outputs[0]["model_intermediate_buffer"]
+        assert intermediate_buffer["codes"]["ref"] == [0.125, -0.25]
+        assert intermediate_buffer["meta"]["ref_audio_sr"] == 24000
+
+    def test_missing_reference_audio_does_not_add_stage_handoff_fields(self) -> None:
+        hidden_states = torch.zeros((2, _HIDDEN_DIM))
+
+        stage_outputs = llm2tts(
+            [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden_states)],
+            prompt={"multi_modal_data": {}},
+        )
+
+        intermediate_buffer = stage_outputs[0]["model_intermediate_buffer"]
+        assert "codes" not in intermediate_buffer
+        assert "ref_audio_sr" not in intermediate_buffer.get("meta", {})
+
     def test_internal_streaming_context_is_accepted(self) -> None:
         hidden = torch.zeros((2, _HIDDEN_DIM))
         out = llm2tts(

tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from collections import defaultdict
 from types import SimpleNamespace
@@ -168,10 +168,14 @@ def test_duplex_turn_end_waits_for_terminal_codec_flush() -> None:
 def test_first_chunk_forwards_reference_voice_and_duplex_identity() -> None:
     manager = _manager()
     request = _request("req")
-    request.additional_information = {
+    request.model_intermediate_buffer = {
         "codes": {"ref": [0.1, -0.1]},
         "meta": {"ref_audio_sr": 16000},
     }
+    request.additional_information = {
+        "codes": {"ref": [0.9]},
+        "meta": {"ref_audio_sr": 8000},
+    }
 
     payload = tts2code2wav_async_chunk(
         manager,
@@ -193,10 +197,33 @@ def test_first_chunk_forwards_reference_voice_and_duplex_identity() -> None:
     assert payload.meta.turn_end is True
 
 
+def test_first_chunk_falls_back_to_legacy_reference_fields() -> None:
+    manager = _manager()
+    request = _request("req")
+    request.model_intermediate_buffer = {
+        "meta": {"ref_audio_sr": 16000},
+    }
+    request.additional_information = {
+        "codes": {"ref": [0.9]},
+        "meta": {"ref_audio_sr": 8000},
+    }
+
+    payload = tts2code2wav_async_chunk(
+        manager,
+        _duplex_delta(*range(7), turn_end=True),
+        request,
+        True,
+    )
+
+    assert payload is not None
+    assert payload.codes.ref.tolist() == pytest.approx([0.9])
+    assert payload.meta.ref_audio_sr == 16000
+
+
 def test_full_payload_forwards_all_codes_and_request_metadata() -> None:
     manager = _manager()
     request = _request("req")
-    request.additional_information = {
+    request.model_intermediate_buffer = {
         "codes": {"ref": [0.1, -0.1]},
         "meta": {
             "ref_audio_sr": 16000,