anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Bugfix] Report sample rate and channels for streamed speech PCM (#8567)

Browse Source 
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
This commit is contained in:
Yueqian Lin authored and GitHub committed 2026-10-07 17:07:49 -07:00
1 parent c32a26a845
commit 1c71a2c03a
8 files changed
+577 -93
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/serving/speech_api.md
+16 -2
View File
@@ -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": "<base64>", "response_format": "pcm" }
{ "type": "speech.audio.delta", "audio": "<base64>", "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 |
tests/benchmarks/patch/test_patch.py
+43
View File
@@ -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(
tests/entrypoints/openai_api/test_serving_speech.py
+234 -76
View File
@@ -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("<i2").tobytes()
return one * 2
@pytest.mark.parametrize(("chunk", "model_sr", "rate", "channels"), _CASES)
def test_raw_pcm_stream_headers(self, mocker: MockerFixture, chunk, model_sr, rate, channels):
app = _build_streaming_speech_app(mocker, chunk=chunk, sample_rate=model_sr)
response = TestClient(app).post(
"/v1/audio/speech",
json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"},
)
assert response.status_code == 200
assert response.headers["content-type"].startswith("audio/pcm")
assert response.headers["x-audio-sample-rate"] == str(rate)
assert response.headers["x-audio-channels"] == str(channels)
exposed = {h.strip().lower() for h in response.headers["access-control-expose-headers"].split(",")}
assert {"x-audio-sample-rate", "x-audio-channels"} <= exposed
# The body stays the model-native interleaved s16le stream, byte for byte.
assert response.content == self._expected_pcm(chunk)
# WAV streaming requires the model to report ``sr`` on its first chunk.
@pytest.mark.parametrize(("chunk", "model_sr", "rate", "channels"), _CASES[1:])
def test_raw_wav_stream_header_matches_metadata(self, mocker: MockerFixture, chunk, model_sr, rate, channels):
app = _build_streaming_speech_app(mocker, chunk=chunk, sample_rate=model_sr)
response = TestClient(app).post(
"/v1/audio/speech",
json={"input": "Hello", "stream_format": "audio", "response_format": "wav"},
)
assert response.status_code == 200
assert response.headers["x-audio-sample-rate"] == str(rate)
assert response.headers["x-audio-channels"] == str(channels)
num_channels, sample_rate = struct.unpack("<HI", response.content[22:28])
assert (sample_rate, num_channels) == (rate, channels)
@pytest.mark.parametrize(("chunk", "model_sr", "rate", "channels"), _CASES)
def test_sse_deltas_carry_format(self, mocker: MockerFixture, chunk, model_sr, rate, channels):
app = _build_streaming_speech_app(mocker, chunk=chunk, sample_rate=model_sr)
response = TestClient(app).post(
"/v1/audio/speech",
json={"input": "Hello", "stream": True, "response_format": "pcm"},
)
deltas = [
json.loads(line.removeprefix("data: "))
for line in response.text.splitlines()
if line.startswith("data: ") and '"speech.audio.delta"' in line
]
assert deltas
assert all((d["sample_rate"], d["channels"]) == (rate, channels) for d in deltas)
@pytest.mark.parametrize(("chunk", "model_sr", "rate", "channels"), _CASES)
def test_non_streaming_pcm_headers(self, mocker: MockerFixture, chunk, model_sr, rate, channels):
app = _build_streaming_speech_app(mocker, chunk=chunk, sample_rate=model_sr)
response = TestClient(app).post("/v1/audio/speech", json={"input": "Hello", "response_format": "pcm"})
assert response.status_code == 200
assert response.headers["x-audio-sample-rate"] == str(rate)
assert response.headers["x-audio-channels"] == str(channels)
assert len(response.content) % (2 * channels) == 0
def test_requested_sample_rate_is_reported(self, mocker: MockerFixture):
"""The headers describe the bytes sent, i.e. after request.sample_rate resampling."""
app = _build_streaming_speech_app(mocker)
speech_server = app.state.speech_server
mocker.patch.object(speech_server, "_validate_speech_sample_rate", return_value=None)
response = TestClient(app).post(
"/v1/audio/speech",
json={"input": "Hello", "stream_format": "audio", "response_format": "pcm", "sample_rate": 8000},
)
assert response.status_code == 200
assert response.headers["x-audio-sample-rate"] == "8000"
assert response.headers["x-audio-channels"] == "1"
# Two 1 s chunks at 8 kHz mono s16le, within resampler edge tolerance.
assert abs(len(response.content) - 2 * 8000 * 2) <= 2 * 64
def test_raw_stream_error_before_audio_is_an_http_error(self, mocker: MockerFixture):
"""Priming the first chunk lets a pre-audio failure return a real status code."""
app = _build_streaming_speech_app(mocker)
speech_server = app.state.speech_server
async def failing_generate(*args, **kwargs):
raise RuntimeError("boom before audio")
yield # pragma: no cover - generator marker, unreachable
speech_server.engine_client.generate = mocker.MagicMock(side_effect=failing_generate)
response = TestClient(app).post(
"/v1/audio/speech",
json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"},
)
assert response.status_code >= 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."""
tests/entrypoints/openai_api/test_serving_speech_stream.py
+104
View File
@@ -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
vllm_omni/benchmarks/data_modules/seed_tts_eval.py
+4 -2
View File
@@ -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
vllm_omni/benchmarks/patch/patch.py
+28
View File
@@ -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
vllm_omni/entrypoints/openai/serving_speech.py
+99 -4
View File
@@ -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:
vllm_omni/entrypoints/openai/serving_speech_stream.py
+49 -9
View File
@@ -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).
<binary frame: audio bytes>
...
{"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