[Frontend] Add configurable output sample rate to Speech API (#6553)

Signed-off-by: 郑倩倩 <2458983186@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/entrypoints/openai_api/test_audio_format.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for audio output format handling in chat completions.
 
 Covers:
@@ -11,10 +12,15 @@ Covers:
 
 from __future__ import annotations
 
+from io import BytesIO
+
 import numpy as np
 import pytest
+import soundfile
+import torch
+import torchaudio
 
-from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
+from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin, StreamingAudioResampler
 from vllm_omni.entrypoints.openai.protocol.audio import (
     DEFAULT_AUDIO_FORMAT,
     SUPPORTED_AUDIO_FORMATS,
@@ -122,6 +128,81 @@ class TestCreateAudio:
         decoded = base64.b64decode(response.audio_data)
         assert decoded[:4] == b"RIFF"
 
+    def test_resamples_wav_to_requested_output_rate(self, mixin, audio_tensor):
+        response = mixin.create_audio(
+            CreateAudio(
+                audio_tensor=audio_tensor,
+                sample_rate=24000,
+                output_sample_rate=8000,
+                response_format="wav",
+                speed=1.0,
+                base64_encode=False,
+            )
+        )
+
+        with soundfile.SoundFile(BytesIO(response.audio_data)) as audio_file:
+            assert audio_file.samplerate == 8000
+            assert audio_file.frames == 8000
+
+
+class TestStreamingAudioResampler:
+    @staticmethod
+    def _resample(waveform):
+        resampler = StreamingAudioResampler(24000, 8000)
+        return np.concatenate(
+            (
+                resampler.process(waveform),
+                resampler.process(np.empty(0), final=True),
+            )
+        )
+
+    def test_chunk_boundaries_do_not_change_output(self):
+        waveform = np.sin(np.linspace(0, 200 * np.pi, 24000, endpoint=False)).astype(np.float32)
+
+        expected = self._resample(waveform)
+
+        chunked = StreamingAudioResampler(24000, 8000)
+        pieces = [chunked.process(chunk) for chunk in np.split(waveform, [137, 2048, 9001, 17003])]
+        pieces.append(chunked.process(np.empty(0), final=True))
+        actual = np.concatenate(pieces)
+
+        assert actual.shape == (8000,)
+        np.testing.assert_allclose(actual, expected, atol=1e-6)
+
+        reference = torchaudio.functional.resample(torch.from_numpy(waveform), 24000, 8000).numpy()
+        np.testing.assert_allclose(actual[100:-100], reference[100:-100], atol=2e-3, rtol=1e-3)
+
+    def test_preserves_passband_signal(self):
+        samples = np.arange(24000, dtype=np.float32)
+        waveform = np.sin(2 * np.pi * 3000 * samples / 24000).astype(np.float32)
+
+        output = self._resample(waveform)
+
+        input_rms = np.sqrt(np.mean(waveform**2))
+        output_rms = np.sqrt(np.mean(output[100:-100] ** 2))
+        assert output_rms == pytest.approx(input_rms, rel=0.02)
+
+    @pytest.mark.parametrize("frequency", [6000, 10000])
+    def test_attenuates_aliasing_frequencies(self, frequency):
+        samples = np.arange(24000, dtype=np.float32)
+        waveform = np.sin(2 * np.pi * frequency * samples / 24000).astype(np.float32)
+
+        output = self._resample(waveform)
+
+        input_rms = np.sqrt(np.mean(waveform**2))
+        output_rms = np.sqrt(np.mean(output[100:-100] ** 2))
+        assert output_rms < input_rms * 0.01
+
+    def test_rejects_non_integer_downsampling_ratio(self):
+        with pytest.raises(ValueError, match="integer downsampling ratio"):
+            StreamingAudioResampler(24000, 16000)
+
+    def test_rejects_multichannel_audio(self):
+        resampler = StreamingAudioResampler(24000, 8000)
+
+        with pytest.raises(ValueError, match="only supports mono audio"):
+            resampler.process(np.zeros((2, 240), dtype=np.float32))
+
 
 class TestResolveAudioFormat:
     """Test _resolve_audio_format via the serving chat class."""

tests/entrypoints/openai_api/test_serving_speech.py

@@ -2953,8 +2953,10 @@ class TestStreamingResponse:
             usage_acc=None,
             tts_params=None,
             collect=None,
+            target_sample_rate=None,
         ):
             captured["tts_params"] = tts_params
+            captured["target_sample_rate"] = target_sample_rate
             yield b"\x00\x00"
 
         monkeypatch.setattr(speech_server, "_prepare_speech_generation", prepare)
@@ -2967,6 +2969,7 @@ class TestStreamingResponse:
 
         assert response.status_code == 200
         assert captured["tts_params"] is finalized_tts_params
+        assert captured["target_sample_rate"] is None
 
     def test_sse_streaming(self, streaming_app):
         """stream_format=sse without stream=True returns audio deltas as SSE."""
@@ -3292,6 +3295,28 @@ class TestMergeBatchItem:
 
         assert merged.non_streaming_mode is False
 
+    def test_sample_rate_batch_default_used(self):
+        """Batch-level sample_rate should be used when the item omits it."""
+        batch = BatchSpeechRequest(
+            items=[SpeechBatchItem(input="hi")],
+            sample_rate=8000,
+        )
+
+        merged = OmniOpenAIServingSpeech._merge_batch_item(batch, batch.items[0])
+
+        assert merged.sample_rate == 8000
+
+    def test_sample_rate_item_override_wins(self):
+        """Per-item sample_rate should override the batch-level default."""
+        batch = BatchSpeechRequest(
+            items=[SpeechBatchItem(input="hi", sample_rate=24000)],
+            sample_rate=8000,
+        )
+
+        merged = OmniOpenAIServingSpeech._merge_batch_item(batch, batch.items[0])
+
+        assert merged.sample_rate == 24000
+
 
 def test_streaming_speech_session_config_accepts_non_streaming_mode():
     config = StreamingSpeechSessionConfig(non_streaming_mode=True)
@@ -3299,6 +3324,33 @@ def test_streaming_speech_session_config_accepts_non_streaming_mode():
     assert config.non_streaming_mode is True
 
 
+@pytest.mark.parametrize("sample_rate", [8000, 24000])
+def test_speech_requests_accept_supported_sample_rates(sample_rate):
+    assert OpenAICreateSpeechRequest(input="hello", sample_rate=sample_rate).sample_rate == sample_rate
+
+
+@pytest.mark.parametrize("sample_rate", [0, -8000])
+def test_speech_request_rejects_non_positive_sample_rate(sample_rate):
+    with pytest.raises(ValidationError):
+        OpenAICreateSpeechRequest(input="hello", sample_rate=sample_rate)
+
+
+def test_sample_rate_uses_adapter_capabilities():
+    server = object.__new__(OmniOpenAIServingSpeech)
+    request = OpenAICreateSpeechRequest(input="hello", sample_rate=8000)
+
+    server._tts_model_type = "qwen3_tts"
+    assert server._validate_speech_sample_rate(request) is None
+
+    request.sample_rate = 16000
+    assert server._validate_speech_sample_rate(request) == (
+        "sample_rate=16000 is not supported by the current TTS model; supported rates: 8000, 24000"
+    )
+
+    server._tts_model_type = "fish_speech"
+    assert server._validate_speech_sample_rate(request) == ("sample_rate is not supported by the current TTS model")
+
+
 def test_streaming_speech_session_config_scopes_native_speed_to_http():
     with pytest.raises(
         ValidationError,
@@ -5027,6 +5079,78 @@ class TestTTSAsyncOffloading:
         assert "partial_audio" not in payload["error"]
         assert "action" not in payload["error"]
 
+    @pytest.mark.asyncio
+    async def test_generate_audio_chunks_resamples_pcm_to_8khz(self, qwen3_tts_server):
+        async def pcm_generator():
+            output = create_mock_audio_output_for_test()
+            output.multimodal_output.update(
+                audio=torch.linspace(-0.5, 0.5, 2400, dtype=torch.float32),
+                sr=24000,
+            )
+            yield output
+
+        chunks = [
+            chunk
+            async for chunk in qwen3_tts_server._generate_audio_chunks(
+                pcm_generator(),
+                "req-8khz",
+                target_sample_rate=8000,
+            )
+        ]
+
+        assert len(chunks) >= 2
+        assert len(b"".join(chunks)) == 800 * 2
+
+    @pytest.mark.asyncio
+    async def test_generate_audio_chunks_wav_header_uses_target_sample_rate(self, qwen3_tts_server):
+        async def pcm_generator():
+            output = create_mock_audio_output_for_test()
+            output.multimodal_output.update(
+                audio=torch.zeros(2400, dtype=torch.float32),
+                sr=24000,
+            )
+            yield output
+
+        chunks = [
+            chunk
+            async for chunk in qwen3_tts_server._generate_audio_chunks(
+                pcm_generator(),
+                "req-8khz-wav",
+                response_format="wav",
+                target_sample_rate=8000,
+            )
+        ]
+
+        assert struct.unpack("<I", chunks[0][24:28])[0] == 8000
+
+    @pytest.mark.asyncio
+    async def test_streaming_resampler_retains_first_chunk_sample_rate_metadata(self, qwen3_tts_server):
+        async def pcm_generator():
+            first = create_mock_audio_output_for_test()
+            first.multimodal_output.update(
+                audio=torch.zeros(10, dtype=torch.float32),
+                sr=24000,
+            )
+            yield first
+
+            second = create_mock_audio_output_for_test()
+            second.multimodal_output.update(audio=torch.zeros(230, dtype=torch.float32))
+            second.multimodal_output.pop("sr", None)
+            yield second
+
+        chunks = [
+            chunk
+            async for chunk in qwen3_tts_server._generate_audio_chunks(
+                pcm_generator(),
+                "req-retained-sr",
+                response_format="wav",
+                target_sample_rate=8000,
+            )
+        ]
+
+        assert chunks[0][:4] == b"RIFF"
+        assert struct.unpack("<I", chunks[0][24:28])[0] == 8000
+
     @pytest.mark.asyncio
     async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server):
         async def pcm_generator():