From 1c71a2c03aab506f114bf5c4458b7b664d82dd5e Mon Sep 17 00:00:00 2001 From: Yueqian Lin Date: Wed, 7 Oct 2026 17:07:49 -0700 Subject: [PATCH] [Bugfix] Report sample rate and channels for streamed speech PCM (#8567) Signed-off-by: Yueqian Lin --- docs/serving/speech_api.md | 18 +- tests/benchmarks/patch/test_patch.py | 43 +++ .../openai_api/test_serving_speech.py | 310 +++++++++++++----- .../openai_api/test_serving_speech_stream.py | 104 ++++++ .../benchmarks/data_modules/seed_tts_eval.py | 6 +- vllm_omni/benchmarks/patch/patch.py | 28 ++ .../entrypoints/openai/serving_speech.py | 103 +++++- .../openai/serving_speech_stream.py | 58 +++- 8 files changed, 577 insertions(+), 93 deletions(-) diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index a9750fd68..2fadb8a47 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -183,6 +183,16 @@ emit these headers. **Raw audio stream** (`stream_format="audio"`). Streams raw audio bytes (PCM or WAV) as they are decoded. +`pcm` is 16-bit signed little-endian, interleaved, at the model's native sample +rate and channel count (24 kHz mono for most models; MOSS-TTS-Local v1.5 emits +48 kHz stereo). Raw streams and non-streaming responses state the format in +response headers: + +| Header | Description | +| --- | --- | +| `x-audio-sample-rate` | Sample rate of the returned audio in Hz (after any `sample_rate` resampling). | +| `x-audio-channels` | Channel count of the returned audio. | + Both streaming modes (`stream_format="audio"` and `"sse"`) require `response_format="pcm"` or `"wav"`, and `speed` must be `1.0` (or omitted). @@ -193,9 +203,13 @@ Each event has an `event:` line and a JSON `data:` line: - `speech.audio.delta` — a base64 audio chunk: ```json - { "type": "speech.audio.delta", "audio": "", "response_format": "pcm" } + { "type": "speech.audio.delta", "audio": "", "response_format": "pcm", + "sample_rate": 24000, "channels": 1 } ``` + `sample_rate` and `channels` extend the OpenAI event and describe the + decoded chunk. + - `speech.audio.done` — terminal event, carrying token `usage`: ```json @@ -340,7 +354,7 @@ Server -> Client: | Message | Description | | --------- | ------------- | -| `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input | +| `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000, "channels": 1}` | Audio generation starting for the buffered input. With `stream_audio=true` it is sent with the first audio chunk, and for `pcm` it carries that audio's `sample_rate` and `channels` (omitted when no audio was produced). | | Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) | | `{"type": "audio.done", "utterance_index": 0, "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input | | `{"type": "session.done", "utterance_index": 0, "total_sentences": N}` | Flushed utterance complete | diff --git a/tests/benchmarks/patch/test_patch.py b/tests/benchmarks/patch/test_patch.py index 6ee5e66cb..7ae27c9df 100644 --- a/tests/benchmarks/patch/test_patch.py +++ b/tests/benchmarks/patch/test_patch.py @@ -35,6 +35,7 @@ from vllm_omni.benchmarks.patch.patch import ( _iter_image_reference_inputs, _iter_video_reference_inputs, _omni_request_timeout_s, + async_request_openai_audio_speech, async_request_openai_chat_omni_completions, async_request_openai_image_edits_omni, async_request_openai_image_generations_omni, @@ -474,6 +475,48 @@ async def test_streaming_error_chunk_marks_request_failed(mocker: MockerFixture) assert output.error == "EngineCore encountered an issue" +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("headers", "env", "expected_duration_s"), + [ + # Server-stated format wins over the env fallback. + ({"X-Audio-Sample-Rate": "48000", "X-Audio-Channels": "2"}, {}, 1.0), + ({"X-Audio-Sample-Rate": "48000", "X-Audio-Channels": "2"}, {"rate": "24000", "channels": "1"}, 1.0), + # Older servers send no headers: the env (default 24 kHz mono) applies. + ({}, {}, 4.0), + ({}, {"rate": "48000", "channels": "2"}, 1.0), + ], +) +async def test_audio_speech_stream_uses_server_pcm_format( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, headers, env, expected_duration_s +): + from vllm_omni.metrics import definitions as defs + + for key, var in (("rate", defs.AUDIO_SAMPLE_RATE_ENV), ("channels", defs.AUDIO_CHANNELS_ENV)): + if key in env: + monkeypatch.setenv(var, env[key]) + else: + monkeypatch.delenv(var, raising=False) + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/audio/speech", + prompt_len=10, + output_len=20, + ) + # 1 s of 48 kHz stereo s16le = 192000 bytes. + mock_response = MockResponse(200, [b"\x00" * 96000, b"\x00" * 96000]) + mock_response.headers = headers + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) + + output = await async_request_openai_audio_speech(request_input, mock_session) + + assert output.success is True + assert output.audio_duration == pytest.approx(expected_duration_s) + + @pytest.mark.asyncio async def test_positive_client_text_timings_take_precedence_over_stage0(mocker: MockerFixture): request_input = RequestFuncInput( diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index f3ad3bf78..d4b6f6801 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -3468,88 +3468,105 @@ class TestStreamingProtocolValidation: assert req.is_raw_audio_stream() is False +def _build_streaming_speech_app( + mocker: MockerFixture, + chunk: torch.Tensor | None = None, + sample_rate: int | None = None, +) -> FastAPI: + """Speech app whose mock engine yields one intermediate then one final chunk. + + ``sample_rate`` is reported as the model's ``sr`` when given; omitted, the + serving layer's 24 kHz default applies (as for most TTS models). + """ + mock_engine_client = mocker.MagicMock() + mock_engine_client.errored = False + if chunk is None: + chunk = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 24000)) + + def _make_output(finished: bool) -> OmniRequestOutput: + mm_output: dict = {"audio": chunk} + if sample_rate is not None: + mm_output["sr"] = torch.tensor(sample_rate) + + class MockCompletionOutput: + def __init__(self, index: int = 0): + self.index = index + self.text = "" + self.token_ids: list[int] = [] + self.finish_reason = "stop" + self.stop_reason = None + self.logprobs = None + + class MockRequestOutput: + def __init__(self, multimodal_output: dict): + self.request_id = "speech-stream-test" + self.outputs = [MockCompletionOutput(index=0)] + self.multimodal_output = multimodal_output + self.finished = finished + self.prompt_token_ids = None + self.encoder_prompt_token_ids = None + self.num_cached_tokens = None + self.prompt_logprobs = None + self.kv_transfer_params = None + + return OmniRequestOutput.from_stage_output( + MockRequestOutput(mm_output), + stage_id=0, + final_output_type="audio", + finished=finished, + ) + + async def mock_generate_streaming(*args, **kwargs): + yield _make_output(finished=False) + yield _make_output(finished=True) + + mock_engine_client.generate = mocker.MagicMock(side_effect=mock_generate_streaming) + mock_engine_client.default_sampling_params_list = [{}] + mock_models = mocker.MagicMock() + mock_models.is_base_model.return_value = True + + speech_server = OmniOpenAIServingSpeech( + engine_client=mock_engine_client, + models=mock_models, + request_logger=mocker.MagicMock(), + ) + speech_server._tts_model_type = None + speech_server._adapter = None + speech_server._diffusion_mode = False + mocker.patch.object(speech_server, "_uses_native_speed_control", return_value=False) + mocker.patch.object( + speech_server, + "create_error_response", + side_effect=lambda message, **_: JSONResponse( + status_code=400, + content={"error": {"message": message}}, + ), + ) + + original_create_speech = speech_server.create_speech + sig = signature(original_create_speech) + new_parameters = [p for name, p in sig.parameters.items() if name != "raw_request"] + new_sig = Signature(parameters=new_parameters, return_annotation=sig.return_annotation) + + async def awaitable_create_speech(*args, **kwargs): + return await original_create_speech(*args, **kwargs) + + awaitable_create_speech.__signature__ = new_sig # type: ignore[attr-defined] + speech_server.create_speech = awaitable_create_speech + + app = FastAPI() + app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None) + app.state.speech_server = speech_server + return app + + class TestStreamingResponse: """Integration tests for the streaming audio response path.""" @pytest.fixture def streaming_app(self, mocker: MockerFixture): """Test app whose mock engine yields one intermediate chunk then a final chunk.""" - mock_engine_client = mocker.MagicMock() - mock_engine_client.errored = False - - def _make_output(finished: bool) -> OmniRequestOutput: - chunk = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 24000)) - - class MockCompletionOutput: - def __init__(self, index: int = 0): - self.index = index - self.text = "" - self.token_ids: list[int] = [] - self.finish_reason = "stop" - self.stop_reason = None - self.logprobs = None - - class MockRequestOutput: - def __init__(self, audio_tensor: torch.Tensor): - self.request_id = "speech-stream-test" - self.outputs = [MockCompletionOutput(index=0)] - self.multimodal_output = {"audio": audio_tensor} - self.finished = finished - self.prompt_token_ids = None - self.encoder_prompt_token_ids = None - self.num_cached_tokens = None - self.prompt_logprobs = None - self.kv_transfer_params = None - - return OmniRequestOutput.from_stage_output( - MockRequestOutput(audio_tensor=chunk), - stage_id=0, - final_output_type="audio", - finished=finished, - ) - - async def mock_generate_streaming(*args, **kwargs): - yield _make_output(finished=False) - yield _make_output(finished=True) - - mock_engine_client.generate = mocker.MagicMock(side_effect=mock_generate_streaming) - mock_engine_client.default_sampling_params_list = [{}] - mock_models = mocker.MagicMock() - mock_models.is_base_model.return_value = True - - speech_server = OmniOpenAIServingSpeech( - engine_client=mock_engine_client, - models=mock_models, - request_logger=mocker.MagicMock(), - ) - speech_server._tts_model_type = None - speech_server._adapter = None - speech_server._diffusion_mode = False - mocker.patch.object(speech_server, "_uses_native_speed_control", return_value=False) - mocker.patch.object( - speech_server, - "create_error_response", - side_effect=lambda message, **_: JSONResponse( - status_code=400, - content={"error": {"message": message}}, - ), - ) - - original_create_speech = speech_server.create_speech - sig = signature(original_create_speech) - new_parameters = [p for name, p in sig.parameters.items() if name != "raw_request"] - new_sig = Signature(parameters=new_parameters, return_annotation=sig.return_annotation) - - async def awaitable_create_speech(*args, **kwargs): - return await original_create_speech(*args, **kwargs) - - awaitable_create_speech.__signature__ = new_sig # type: ignore[attr-defined] - speech_server.create_speech = awaitable_create_speech - - app = FastAPI() - app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None) - app.state.speech_server = speech_server - return app + return _build_streaming_speech_app(mocker) @staticmethod def _assert_sse_audio_response(response, response_format: str = "pcm"): @@ -3603,6 +3620,7 @@ class TestStreamingResponse: tts_params=None, collect=None, target_sample_rate=None, + audio_format=None, ): captured["tts_params"] = tts_params captured["target_sample_rate"] = target_sample_rate @@ -3790,6 +3808,146 @@ class TestStreamingResponse: assert "audio/wav" in response.headers["content-type"] +class TestSpeechAudioFormatMetadata: + """pcm carries no header, so responses must state the model-native format.""" + + # 24 kHz mono (most TTS models) and 48 kHz [channels, samples] stereo + # (e.g. MOSS-TTS-Local v1.5); the serving layer must not care which. + _MONO = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 2400)) + _STEREO = torch.stack([_MONO.repeat(2) * 0.5, -_MONO.repeat(2) * 0.25]) + + _CASES = [ + pytest.param(None, None, 24000, 1, id="default-24k-mono"), + pytest.param(_MONO, 24000, 24000, 1, id="explicit-24k-mono"), + pytest.param(_STEREO, 48000, 48000, 2, id="48k-stereo"), + ] + + @staticmethod + def _expected_pcm(chunk: torch.Tensor | None) -> bytes: + """Two engine outputs, each the chunk as interleaved s16le.""" + # ``None`` is the builder's default chunk. + audio = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 24000)) if chunk is None else chunk + samples = audio.numpy().T if audio.ndim == 2 else audio.numpy() + # libsndfile-compatible PCM_16 conversion used by ``create_audio``. + one = np.clip(np.floor(samples.astype(np.float64) * 32768.0), -32768, 32767).astype("= 400 + assert "audio/pcm" not in response.headers.get("content-type", "") + assert "x-audio-sample-rate" not in response.headers + + def test_primed_stream_closes_inner_generator(self): + """Closing the response body before iteration still closes the engine stream.""" + from vllm_omni.entrypoints.openai.serving_speech import _PrimedAudioStream + + closed: list[bool] = [] + + async def chunks(): + try: + yield b"\x01\x00" + yield b"\x02\x00" + finally: + closed.append(True) + + async def run() -> list[bytes]: + body = await _PrimedAudioStream.start(chunks()) + await body.aclose() + assert closed == [True] + body = await _PrimedAudioStream.start(chunks()) + return [c async for c in body] + + assert asyncio.run(run()) == [b"\x01\x00", b"\x02\x00"] + + class TestSpeechBatchAPI: """Tests for the /v1/audio/speech/batch endpoint.""" diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py index 1d34b6e0b..47c25a323 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_stream.py +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -33,6 +33,12 @@ def _fake_aligner_res(pairs, words): ) +def _fill_audio_format(audio_format: dict | None, sample_rate: int = 24000, channels: int = 1) -> None: + """Mirror ``_generate_audio_chunks``: report the format before the first chunk.""" + if audio_format is not None and not audio_format: + audio_format.update(sample_rate=sample_rate, channels=channels) + + def _build_test_app( speech_service=None, *, @@ -57,8 +63,10 @@ def _build_test_app( tts_params=None, collect=None, cumulative_audio=False, + audio_format=None, ): for chunk in (b"\x01\x02", b"\x03\x04\x05"): + _fill_audio_format(audio_format) yield (chunk, 24000) if include_sample_rate else chunk speech_service._generate_pcm_chunks = mock_generate_pcm_chunks @@ -282,6 +290,7 @@ class TestStreamingSpeechWebSocket: request_arrival_ts=None, include_sample_rate=False, tts_params=None, + audio_format=None, ): assert request_start_s is not None assert request_arrival_ts is not None @@ -289,6 +298,7 @@ class TestStreamingSpeechWebSocket: captured_timing["chunk_arrival"] = request_arrival_ts captured_tts_params.append(tts_params) for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"): + _fill_audio_format(audio_format) yield (chunk, 24000) if include_sample_rate else chunk speech_service._generate_pcm_chunks = mock_generate_pcm_chunks @@ -312,6 +322,7 @@ class TestStreamingSpeechWebSocket: assert start["type"] == "audio.start" assert start["format"] == "pcm" assert start["sample_rate"] == 24000 + assert start["channels"] == 1 assert ws.receive_bytes() == b"\x01\x02" assert ws.receive_bytes() == b"\x03\x04\x05" @@ -337,6 +348,91 @@ class TestStreamingSpeechWebSocket: assert captured_tts_params == [{"_qwen3_tts_effective_max_tokens": [192]}] assert speech_service._generate_audio_bytes.await_count == 0 + @staticmethod + def _stream_service(mocker: MockerFixture, generate_pcm_chunks, *, forced_aligner: bool = False): + speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) + speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) + speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {})) + speech_service._generate_pcm_chunks = generate_pcm_chunks + speech_service.engine_client = mocker.MagicMock() + speech_service.engine_client.abort = mocker.AsyncMock() + speech_service.forced_aligner_enabled = forced_aligner + return speech_service + + @staticmethod + def _send_pcm_session(ws, **config) -> None: + ws.send_json( + {"type": "session.config", "voice": "Vivian", "stream_audio": True, "response_format": "pcm", **config} + ) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) + + def test_streaming_pcm_start_reports_model_native_format(self, mocker: MockerFixture): + """audio.start states the real rate/channels, not a nominal 24 kHz.""" + + async def mock_generate_pcm_chunks(_generator, _request_id, *, audio_format=None, **_kwargs): + _fill_audio_format(audio_format, sample_rate=48000, channels=2) + yield b"\x01\x02\x03\x04" + + app, _ = _build_test_app(self._stream_service(mocker, mock_generate_pcm_chunks)) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + self._send_pcm_session(ws) + start = ws.receive_json() + assert start["type"] == "audio.start" + assert start["sample_rate"] == 48000 + assert start["channels"] == 2 + assert ws.receive_bytes() == b"\x01\x02\x03\x04" + assert ws.receive_json()["type"] == "audio.done" + + def test_streaming_pcm_error_before_audio_keeps_start_done_pairing(self, mocker: MockerFixture): + """No audio means no known format: audio.start omits it rather than guessing.""" + + async def mock_generate_pcm_chunks(_generator, _request_id, *, audio_format=None, **_kwargs): + raise RuntimeError("early boom") + yield # pragma: no cover - generator marker, unreachable + + app, _ = _build_test_app(self._stream_service(mocker, mock_generate_pcm_chunks)) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + self._send_pcm_session(ws) + start = ws.receive_json() + assert start["type"] == "audio.start" + assert "sample_rate" not in start + assert "channels" not in start + assert ws.receive_json() == { + "type": "error", + "message": "Generation failed for utterance 0, sentence 0: early boom", + } + done = ws.receive_json() + assert done["type"] == "audio.done" + assert done["error"] is True + + def test_word_timestamps_chunk_ms_account_for_channels(self, mocker: MockerFixture): + """Stereo s16le is 4 bytes per frame; chunk offsets must not double.""" + + async def mock_generate_pcm_chunks( + _generator, _request_id, *, include_sample_rate=False, collect=None, audio_format=None, **_kwargs + ): + _fill_audio_format(audio_format, sample_rate=1000, channels=2) + yield (b"\x01" * 4000, 1000) if include_sample_rate else b"\x01" * 4000 + + speech_service = self._stream_service(mocker, mock_generate_pcm_chunks, forced_aligner=True) + app, _ = _build_test_app(speech_service) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + self._send_pcm_session(ws, word_timestamps=True) + start = ws.receive_json() + assert (start["sample_rate"], start["channels"]) == (1000, 2) + chunk = ws.receive_json() + assert (chunk["chunk_start_ms"], chunk["chunk_end_ms"]) == (0, 1000) + assert (chunk["sample_rate"], chunk["channels"]) == (1000, 2) + final = ws.receive_json() + assert final["chunk_end_ms"] == 1000 + def test_word_timestamps_requires_configured_aligner(self, mocker: MockerFixture): app, _ = _build_test_app(mocker=mocker) @@ -388,11 +484,13 @@ class TestStreamingSpeechWebSocket: tts_params=None, collect=None, cumulative_audio=False, + audio_format=None, ): assert request_start_s is not None assert request_arrival_ts is not None assert cumulative_audio is True for chunk in (first_chunk, second_chunk): + _fill_audio_format(audio_format, sample_rate=1000) yield (chunk, 1000) if include_sample_rate else chunk if collect is not None: collect["aligner_res"] = _fake_aligner_res([[0, 200], [200, 900]], ["Hello", "world"]) @@ -484,8 +582,10 @@ class TestStreamingSpeechWebSocket: tts_params=None, collect=None, cumulative_audio=False, + audio_format=None, ): chunk = b"\x01" * 1000 + _fill_audio_format(audio_format, sample_rate=1000) yield (chunk, 1000) if include_sample_rate else chunk if collect is not None: collect["aligner_res"] = _fake_aligner_res([[0, 1000], [1000, 1200]], ["Hello", "world"]) @@ -673,7 +773,9 @@ class TestStreamingSpeechWebSocket: request_arrival_ts=None, include_sample_rate=False, tts_params=None, + audio_format=None, ): + _fill_audio_format(audio_format) yield b"\x01\x02" raise RuntimeError("stream boom") @@ -777,7 +879,9 @@ class TestStreamingSpeechWebSocket: request_arrival_ts=None, include_sample_rate=False, tts_params=None, + audio_format=None, ): + _fill_audio_format(audio_format) yield b"\x01\x02" speech_service._generate_pcm_chunks = mock_generate_pcm_chunks diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py index 100933870..61b751dd3 100644 --- a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py +++ b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py @@ -42,8 +42,10 @@ Env: ``SEED_TTS_EVAL_DEVICE`` (e.g. ``cuda:0``, ``cpu``); ``SEED_TTS_HF_WHISPER_ defaults to ``openai/whisper-large-v3`` (override for debugging only). Set ``SEED_TTS_WER_SAVE_AUDIO_DIR`` to save the captured 24 kHz mono WAV used by WER evaluation for each synthesized utterance. -Streaming PCM is decoded using ``VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE`` / -``VLLM_OMNI_BENCH_AUDIO_CHANNELS`` (default: 24 kHz mono). +Streaming PCM is decoded using the server's ``X-Audio-Sample-Rate`` / +``X-Audio-Channels`` response headers; ``VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE`` / +``VLLM_OMNI_BENCH_AUDIO_CHANNELS`` (default: 24 kHz mono) apply when a server +does not send them. """ from __future__ import annotations diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index b0ca826df..c2dd001c2 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -207,6 +207,32 @@ def _pcm_s16le_to_seed_tts_wer_bytes( return (pcm_f32 * 32767).astype(np.int16).tobytes() +def _stream_pcm_format_from_headers( + headers: Mapping[str, str], + sample_rate: int, + channels: int, +) -> tuple[int, int]: + """Prefer the server's ``X-Audio-Sample-Rate``/``X-Audio-Channels`` headers. + + Raw PCM has no self-describing header; the server states the model-native + format. ``sample_rate``/``channels`` (env-configured) remain the fallback. + """ + resolved = [] + for name, fallback in (("X-Audio-Sample-Rate", sample_rate), ("X-Audio-Channels", channels)): + raw = headers.get(name) + value = fallback + if raw: + try: + value = int(raw) + except ValueError: + logger.warning("Ignoring invalid %s header %r; using %d", name, raw, fallback) + else: + if value != fallback: + logger.debug("Streamed PCM %s=%d from server header (configured %d)", name, value, fallback) + resolved.append(max(value, 1)) + return resolved[0], resolved[1] + + get_samples_old = datasets.get_samples _DEFAULT_DAILY_OMNI_REPO = "liarliar/Daily-Omni" @@ -2610,6 +2636,7 @@ async def async_request_openai_audio_speech( output.prompt_len = request_func_input.prompt_len # PCM format: 16-bit signed; sample_rate/channels are model-dependent. + # The env vars are a fallback for servers that do not send format headers. sample_rate, channels = defs.stream_pcm_format_from_env() sample_width = defs.DEFAULT_AUDIO_SAMPLE_WIDTH @@ -2623,6 +2650,7 @@ async def async_request_openai_audio_speech( try: async with session.post(url=api_url, json=payload, headers=headers) as response: if response.status == 200: + sample_rate, channels = _stream_pcm_format_from_headers(response.headers, sample_rate, channels) async for chunk in response.content.iter_any(): if not chunk: continue diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 42afdbeb0..60b15dfe1 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -79,6 +79,10 @@ _SPEECH_USAGE_OUTPUT_TOKENS_HEADER = "X-VLLM-OMNI-OUTPUT-TOKENS" _SPEECH_USAGE_TOTAL_TOKENS_HEADER = "X-VLLM-OMNI-TOTAL-TOKENS" _SPEECH_USAGE_INPUT_TEXT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TEXT-TOKENS" _SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-AUDIO-TOKENS" +# Raw ``pcm`` bodies carry no format header and the rate/channel count is +# model-native, so responses state the s16le layout explicitly. +SPEECH_AUDIO_SAMPLE_RATE_HEADER = "X-Audio-Sample-Rate" +SPEECH_AUDIO_CHANNELS_HEADER = "X-Audio-Channels" def _stage_speech_metadata(stage: Any) -> tuple[str | None, str | None, str | None]: @@ -176,6 +180,64 @@ def _infer_audio_num_channels(audio: np.ndarray) -> int: return 1 +def _audio_format_headers(audio_format: dict[str, int] | None) -> dict[str, str]: + """Response headers describing the emitted audio's rate and channel count.""" + if not audio_format: + return {} + return { + SPEECH_AUDIO_SAMPLE_RATE_HEADER: str(int(audio_format["sample_rate"])), + SPEECH_AUDIO_CHANNELS_HEADER: str(int(audio_format["channels"])), + # Without this, CORS hides both headers from cross-origin browser clients. + "Access-Control-Expose-Headers": f"{SPEECH_AUDIO_SAMPLE_RATE_HEADER}, {SPEECH_AUDIO_CHANNELS_HEADER}", + } + + +class _PrimedAudioStream: + """Async iterator over an audio stream whose first chunk was already pulled. + + Raw audio streams pull the first chunk before the HTTP response starts so + its headers can state the real sample rate and channel count. + """ + + def __init__(self, chunks: Any, first: Any, has_first: bool) -> None: + self._chunks = chunks + self._first = first + self._has_first = has_first + + @classmethod + async def start(cls, chunks: Any) -> "_PrimedAudioStream": + try: + first = await anext(chunks) + except StopAsyncIteration: + return cls(chunks, None, False) + except BaseException: + with anyio.CancelScope(shield=True): + await chunks.aclose() + raise + return cls(chunks, first, True) + + def __aiter__(self) -> "_PrimedAudioStream": + return self + + async def __anext__(self) -> Any: + if self._has_first: + first, self._first, self._has_first = self._first, None, False + return first + return await anext(self._chunks) + + async def aclose(self) -> None: + self._first, self._has_first = None, False + await self._chunks.aclose() + + +def _encoded_audio_format(audio_response: Any, sample_rate: int, audio: np.ndarray) -> dict[str, int]: + """Rate and channel count of an encoded chunk, preferring encoder metadata.""" + metadata = getattr(audio_response, "audio_metadata", None) + if metadata is not None: + return {"sample_rate": int(metadata.sample_rate_hz), "channels": int(metadata.channels)} + return {"sample_rate": int(sample_rate), "channels": _infer_audio_num_channels(np.asarray(audio))} + + def _sanitize_filename(filename: str) -> str: """Sanitize filename to prevent path traversal attacks. @@ -1568,6 +1630,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): collect: dict | None = None, target_sample_rate: int | None = None, cumulative_audio: bool = False, + audio_format: dict[str, int] | None = None, ): """Generate audio chunks for streaming response. @@ -1581,6 +1644,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): generator: Async generator from the engine request_id: Request identifier for logging response_format: Audio format (pcm or wav) + audio_format: Optional out-parameter. Before the first audio + chunk is yielded it receives that chunk's ``sample_rate`` and + ``channels``, so callers can describe a raw PCM stream. Yields: Raw audio bytes for each chunk (with WAV header for first chunk if wav format) @@ -1722,8 +1788,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): speed=1.0, base64_encode=False, ) - audio_bytes = cast(bytes, self.create_audio(audio_obj).audio_data) + audio_response = self.create_audio(audio_obj) + audio_bytes = cast(bytes, audio_response.audio_data) record_audio_chunk(audio_bytes, chunk_np, res) + if audio_format is not None and not audio_format and audio_bytes: + audio_format.update(_encoded_audio_format(audio_response, output_sample_rate, chunk_np)) if wav_header is not None: yield wav_header if include_sample_rate: @@ -1750,8 +1819,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): speed=1.0, base64_encode=False, ) - audio_bytes = cast(bytes, self.create_audio(audio_obj).audio_data) + audio_response = self.create_audio(audio_obj) + audio_bytes = cast(bytes, audio_response.audio_data) record_audio_chunk(audio_bytes, final_chunk, last_audio_result) + if audio_format is not None and not audio_format and audio_bytes: + audio_format.update(_encoded_audio_format(audio_response, output_sample_rate, final_chunk)) if wav_header is not None: yield wav_header if include_sample_rate: @@ -1880,6 +1952,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): """ usage_acc = SpeechOutputTokenCounter() emitted_audio = False + audio_format: dict[str, int] = {} try: async with aclosing( self._generate_audio_chunks( @@ -1892,6 +1965,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): usage_acc=usage_acc, tts_params=tts_params, target_sample_rate=request.sample_rate if request is not None else None, + audio_format=audio_format, ) ) as chunks: async for chunk in chunks: @@ -1900,6 +1974,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): "audio": base64.b64encode(chunk).decode("ascii"), "response_format": response_format, } + if audio_format: + # Additive to the OpenAI schema: pcm deltas are + # model-native s16le, not always 24 kHz mono. + payload["sample_rate"] = audio_format["sample_rate"] + payload["channels"] = audio_format["channels"] data = json.dumps(payload, separators=(",", ":")) emitted_audio = True yield f"event: speech.audio.delta\ndata: {data}\n\n" @@ -2134,6 +2213,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): collect: dict | None = None, target_sample_rate: int | None = None, cumulative_audio: bool = False, + audio_format: dict[str, int] | None = None, ): """Yield raw PCM byte chunks from the engine generator. @@ -2154,6 +2234,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): collect=collect, target_sample_rate=target_sample_rate, cumulative_audio=cumulative_audio, + audio_format=audio_format, ) ) as chunks: async for chunk in chunks: @@ -2355,6 +2436,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): base64_encode=base64_encode, ) audio_response: AudioResponse = self.create_audio(audio_obj) + if collect is not None: + collect["audio_format"] = _encoded_audio_format( + audio_response, request.sample_rate or sample_rate, audio_tensor + ) self._mark_ref_audio_artifact_ready_for_request(request_id) artifact_ready = True if usage_out is not None: @@ -2634,7 +2719,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id=request_id, arrival_time=request_arrival_ts, ) - return _SpeechStreamingResponse( + audio_format: dict[str, int] = {} + # Pull the first chunk before sending headers: the body has no + # self-describing format for pcm and the rate/channel count is + # model-native (not always 24 kHz mono), so the headers state it. + body = await _PrimedAudioStream.start( self._generate_audio_chunks( generator, request_id, @@ -2644,8 +2733,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_arrival_ts=request_arrival_ts, tts_params=raw_tts_params, target_sample_rate=request.sample_rate, - ), + audio_format=audio_format, + ) + ) + return _SpeechStreamingResponse( + body, media_type=media_type, + headers=_audio_format_headers(audio_format), ) if request.is_sse_stream(): @@ -2735,6 +2829,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ) if collect.get("audio_truncated") is not None: headers["X-Audio-Truncated"] = "true" if collect["audio_truncated"] else "false" + headers.update(_audio_format_headers(collect.get("audio_format"))) return Response(content=audio_bytes, media_type=media_type, headers=headers) except asyncio.CancelledError: diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index 45bdd4e90..2b3a4eb98 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -30,6 +30,9 @@ Protocol: Server -> Client (default, word_timestamps=false): {"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "wav"} + # With stream_audio=true, audio.start is sent with the first audio + # chunk; for format "pcm" it also carries that audio's "sample_rate" + # and "channels" (s16le, interleaved, model-native). ... {"type": "audio.done", "utterance_index": 0, "sentence_index": 0} @@ -56,6 +59,7 @@ import asyncio import base64 import json import time +from collections.abc import Awaitable, Callable from contextlib import aclosing from fastapi import WebSocket, WebSocketDisconnect @@ -75,7 +79,7 @@ logger = init_logger(__name__) _DEFAULT_IDLE_TIMEOUT = 30.0 # seconds _DEFAULT_CONFIG_TIMEOUT = 10.0 # seconds _PCM_SAMPLE_RATE = 24000 -_BYTES_PER_SAMPLE = 2 # 16-bit mono PCM +_BYTES_PER_SAMPLE = 2 # 16-bit PCM, per channel _MAX_CONFIG_MESSAGE_SIZE = 4 * 1024 * 1024 # allow large ref_audio payloads _MAX_INPUT_TEXT_MESSAGE_SIZE = 128 * 1024 @@ -336,13 +340,28 @@ class OmniStreamingSpeechHandler: "sentence_text": sentence_text, "format": response_format, } - if config.stream_audio and response_format == "pcm": - # Nominal stream rate; each audio.chunk carries the authoritative - # per-chunk sample_rate. - start_payload["sample_rate"] = _PCM_SAMPLE_RATE if config.word_timestamps: start_payload["word_timestamps"] = True - await websocket.send_json(start_payload) + # Streamed pcm is model-native s16le. Its rate and channel count are + # known once the first chunk is encoded, so for pcm audio.start is + # deferred to that point and states them; it is sent bare if no audio + # arrives. Non-streamed responses still send it before generation. + describe_pcm = config.stream_audio and response_format == "pcm" + audio_format: dict[str, int] = {} + start_sent = False + + async def send_start() -> None: + nonlocal start_sent + if start_sent: + return + start_sent = True + if describe_pcm and audio_format: + start_payload["sample_rate"] = audio_format["sample_rate"] + start_payload["channels"] = audio_format["channels"] + await websocket.send_json(start_payload) + + if not describe_pcm: + await send_start() total_bytes = 0 generation_failed = False @@ -365,6 +384,8 @@ class OmniStreamingSpeechHandler: request_start_s=request_start_s, request_arrival_ts=request_arrival_ts, tts_params=tts_params, + audio_format=audio_format, + send_start=send_start, ) else: async with aclosing( @@ -374,9 +395,11 @@ class OmniStreamingSpeechHandler: request_start_s=request_start_s, request_arrival_ts=request_arrival_ts, tts_params=tts_params, + audio_format=audio_format, ) ) as stream: async for chunk in stream: + await send_start() total_bytes += len(chunk) await websocket.send_bytes(chunk) else: @@ -401,6 +424,10 @@ class OmniStreamingSpeechHandler: sentence_index, e, ) + try: + await send_start() + except Exception: + logger.debug("Failed to send audio.start for sentence %d", sentence_index, exc_info=True) await self._send_error( websocket, f"Generation failed for utterance {utterance_index}, sentence {sentence_index}: {e}", @@ -408,6 +435,9 @@ class OmniStreamingSpeechHandler: ) finally: try: + if not start_sent: + # Keep audio.start/audio.done paired when no audio was produced. + await send_start() await websocket.send_json( { "type": "audio.done", @@ -433,6 +463,8 @@ class OmniStreamingSpeechHandler: request_arrival_ts: float, language: str | None = None, tts_params: dict | None = None, + audio_format: dict[str, int] | None = None, + send_start: Callable[[], Awaitable[None]] | None = None, ) -> int: """Stream PCM as JSON ``audio.chunk`` frames, aligned per sentence. @@ -450,6 +482,10 @@ class OmniStreamingSpeechHandler: chunk_id = 0 # Receives the aligner stage's pooling output from the generator. collect: dict = {} + audio_format = audio_format if audio_format is not None else {} + + def frame_bytes() -> int: + return _BYTES_PER_SAMPLE * int(audio_format.get("channels", 1)) async def send_chunk( chunk: bytes, @@ -459,6 +495,8 @@ class OmniStreamingSpeechHandler: chunk_end_ms: int, ) -> None: nonlocal chunk_id + if send_start is not None: + await send_start() await websocket.send_json( { "type": "audio.chunk", @@ -468,6 +506,7 @@ class OmniStreamingSpeechHandler: "chunk_start_ms": chunk_start_ms, "chunk_end_ms": chunk_end_ms, "sample_rate": chunk_sample_rate, + "channels": int(audio_format.get("channels", 1)), "audio_b64": base64.b64encode(chunk).decode("ascii"), "timestamps": timestamps_payload, } @@ -484,13 +523,14 @@ class OmniStreamingSpeechHandler: tts_params=tts_params, collect=collect, cumulative_audio=True, + audio_format=audio_format, ) ) as stream: async for chunk, chunk_sample_rate in stream: sample_rate = chunk_sample_rate - chunk_start_ms = int(round((audio_bytes_seen / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) + chunk_start_ms = int(round((audio_bytes_seen / frame_bytes() / sample_rate) * 1000.0)) audio_bytes_seen += len(chunk) - chunk_end_ms = int(round((audio_bytes_seen / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) + chunk_end_ms = int(round((audio_bytes_seen / frame_bytes() / sample_rate) * 1000.0)) total_bytes += len(chunk) # Audio first, timestamps after the whole sentence is aligned. await send_chunk(chunk, chunk_sample_rate, None, chunk_start_ms, chunk_end_ms) @@ -502,7 +542,7 @@ class OmniStreamingSpeechHandler: timestamps_payload = ( extract_word_timestamps(aligner_res, sentence_text, language) if aligner_res is not None else None ) - sentence_end_ms = int(round((audio_bytes_seen / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) + sentence_end_ms = int(round((audio_bytes_seen / frame_bytes() / sample_rate) * 1000.0)) await send_chunk(b"", sample_rate, timestamps_payload, 0, sentence_end_ms) return total_bytes