[Edge][WP5] Add codec-token streaming mode with on-device decoding
* qwen3_tts_talker_only pipeline (single terminal LLM_AR stage, final_output_type "latent") + vllm_omni/deploy/qwen3_tts_talker_only.yaml; a terminal talker publishes its per-step codes.audio rows to the client (omni_client_multimodal_output_keys). * /v1/audio/speech/stream: session.config.output_mode="codec_tokens" streams codec.start (JSON) + binary frames (<IHB seq/n_frames/flags + int16 rows, vllm_omni/entrypoints/openai/codec_stream.py) + codec.done; new input.cancel message aborts an in-flight generation for both codec and PCM sessions (single reader task + session inbox replaces direct socket reads); optional session.config.seed; audio.done gains "cancelled" only when it happened. * serving_speech._generate_codec_chunks aligns the cumulative code tensor to output token ids (drops prefill placeholder/EOS rows) and emits deltas. * TTS adapters: supports_codec_stream / codec_stream_spec (Qwen3-TTS: 16 codebooks, 2048 ids, 12.5 Hz, 24 kHz, decoder qwen3_tts_tokenizer_v2_12hz). * examples/online_serving/text_to_speech/qwen3_tts/edge_decoder_client.py decodes locally with Qwen3TTSTokenizerV2Decoder (ramp schedule, 72-frame context, cancel/reset), codec_parity_check.py and benchmarks/tts/codec_stream_e2e.py compare against the two-stage server. * docs/serving/codec_token_streaming.md; CPU unit tests for the wire format, alignment, session config and handler (frames, unsupported adapter, cancel); L3 e2e tests/e2e/online_serving/test_qwen3_tts_codec_stream.py. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
tzhouamandClaude Fable 5.1 committed
1 parent
35e43b9983
commit
49a839ecf9
16 files changed
+1533
-6
No files matched your search
@@ -0,0 +1,183 @@
|
||||
"""Codec-token streaming: wire format, row alignment, and handler behaviour (CPU)."""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from fastapi import FastAPI, WebSocket
|
||||
|
||||
from vllm_omni.entrypoints.openai import codec_stream as cs
|
||||
from vllm_omni.entrypoints.openai.protocol.audio import StreamingSpeechSessionConfig
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_pack_unpack_roundtrip_and_flags():
|
||||
codes = np.random.randint(0, 2048, size=(25, 16)).astype(np.int64)
|
||||
buf = cs.pack_codec_frame(7, torch.from_numpy(codes), eos=False, segment_end=True)
|
||||
assert len(buf) == cs.HEADER_SIZE + 25 * 16 * 2
|
||||
seq, n, flags, arr = cs.unpack_codec_frame(buf, 16)
|
||||
assert seq == 7 and n == 25 and flags == cs.FLAG_SEGMENT_END
|
||||
assert arr.dtype == np.int16 and np.array_equal(arr.astype(np.int64), codes)
|
||||
# EOS with zero frames
|
||||
buf = cs.pack_codec_frame(8, None, eos=True, codebooks=16)
|
||||
seq, n, flags, arr = cs.unpack_codec_frame(buf, 16)
|
||||
assert (seq, n, flags) == (8, 0, cs.FLAG_EOS) and arr.shape == (0, 16)
|
||||
buf = cs.pack_codec_frame(2**32 + 5, np.zeros((1, 16), dtype=np.int64), reset=True)
|
||||
assert cs.unpack_codec_frame(buf, 16)[0] == 5
|
||||
assert cs.unpack_codec_frame(buf, 16)[2] == cs.FLAG_RESET
|
||||
|
||||
|
||||
def test_pack_rejects_bad_shapes_and_ranges():
|
||||
with pytest.raises(ValueError):
|
||||
cs.pack_codec_frame(0, np.zeros((2, 2, 2), dtype=np.int64))
|
||||
with pytest.raises(ValueError):
|
||||
cs.pack_codec_frame(0, np.full((1, 16), 40000, dtype=np.int64))
|
||||
with pytest.raises(ValueError):
|
||||
cs.unpack_codec_frame(b"\x00\x01", 16)
|
||||
with pytest.raises(ValueError):
|
||||
cs.unpack_codec_frame(cs.pack_codec_frame(0, np.zeros((1, 16), dtype=np.int64))[:-2], 16)
|
||||
|
||||
|
||||
def test_align_codec_rows_drops_prefill_placeholder_and_eos():
|
||||
# 1 prefill placeholder row (zeros, emitted before any token) + 3 real frames + EOS row
|
||||
codes = torch.tensor([[0] * 16, [5] * 16, [6] * 16, [7] * 16, [0] * 16])
|
||||
token_ids = [5, 6, 7, 2150] # codebook-0 ids then EOS
|
||||
rows = cs.align_codec_rows(codes, token_ids, 2048)
|
||||
assert rows.shape == (3, 16) and rows[:, 0].tolist() == [5, 6, 7]
|
||||
# trailing -1 placeholders in token ids are ignored
|
||||
rows = cs.align_codec_rows(codes, token_ids + [-1, -1], 2048)
|
||||
assert rows.shape == (3, 16)
|
||||
# out-of-range rows are dropped
|
||||
bad = torch.tensor([[5] * 16, [3000] * 16])
|
||||
assert cs.align_codec_rows(bad, [5, 6], 2048).shape == (1, 16)
|
||||
assert cs.align_codec_rows(None, [], 2048).shape == (0, 0)
|
||||
assert cs.align_codec_rows(torch.zeros(0, 16, dtype=torch.long), [], 2048).shape == (0, 16)
|
||||
|
||||
|
||||
def test_session_config_codec_mode_validation():
|
||||
cfg = StreamingSpeechSessionConfig(output_mode="codec_tokens", stream_audio=True, response_format="pcm")
|
||||
assert cfg.output_mode == "codec_tokens"
|
||||
with pytest.raises(ValueError):
|
||||
StreamingSpeechSessionConfig(output_mode="codec_tokens", stream_audio=False, response_format="pcm")
|
||||
with pytest.raises(ValueError):
|
||||
StreamingSpeechSessionConfig(
|
||||
output_mode="codec_tokens", stream_audio=True, response_format="pcm", word_timestamps=True
|
||||
)
|
||||
assert StreamingSpeechSessionConfig().output_mode == "pcm"
|
||||
|
||||
|
||||
def test_codec_start_message_fields():
|
||||
spec = {"codebooks": 16, "codebook_size": 2048, "frame_rate_hz": 12.5, "sample_rate": 24000, "decoder_id": "x"}
|
||||
msg = cs.codec_start_message(utterance_index=1, sentence_index=0, sentence_text="hi", spec=spec, model="m")
|
||||
assert msg["type"] == "codec.start" and msg["codebooks"] == 16 and msg["header_format"] == cs.HEADER_FMT
|
||||
assert msg["decoder_id"] == "x" and msg["version"] == cs.CODEC_STREAM_VERSION
|
||||
|
||||
|
||||
# --------------------------------------------------------------------- handler
|
||||
def _codec_app(mocker, *, supports: bool = True, cancel_after: int | None = None):
|
||||
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
|
||||
from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler
|
||||
|
||||
speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
|
||||
speech_service.forced_aligner_enabled = False
|
||||
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {}))
|
||||
adapter = mocker.MagicMock()
|
||||
adapter.supports_codec_stream = supports
|
||||
adapter.codec_stream_spec = {
|
||||
"codebooks": 16,
|
||||
"codebook_size": 2048,
|
||||
"frame_rate_hz": 12.5,
|
||||
"sample_rate": 24000,
|
||||
"decoder_id": "d",
|
||||
}
|
||||
speech_service._get_tts_adapter = mocker.MagicMock(return_value=adapter)
|
||||
speech_service.engine_client = mocker.MagicMock()
|
||||
speech_service.engine_client.abort = mocker.AsyncMock()
|
||||
|
||||
async def mock_codec_chunks(_generator, _request_id, *, tts_params=None, min_frames=1):
|
||||
import asyncio
|
||||
|
||||
for i in range(3):
|
||||
await asyncio.sleep(0.01)
|
||||
yield np.full((1 + i, 16), i, dtype=np.int64), i == 2
|
||||
|
||||
speech_service._generate_codec_chunks = mock_codec_chunks
|
||||
handler = OmniStreamingSpeechHandler(speech_service=speech_service, idle_timeout=5.0, config_timeout=5.0)
|
||||
app = FastAPI()
|
||||
|
||||
@app.websocket("/v1/audio/speech/stream")
|
||||
async def ws_endpoint(websocket: WebSocket):
|
||||
# NOTE: no ``from __future__ import annotations`` in this module: FastAPI
|
||||
# must resolve the WebSocket annotation, or it treats it as a query param.
|
||||
await handler.handle_session(websocket)
|
||||
|
||||
return app, speech_service
|
||||
|
||||
|
||||
def test_codec_stream_handler_frames_and_done(mocker):
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
app, svc = _codec_app(mocker)
|
||||
with TestClient(app) as client, client.websocket_connect("/v1/audio/speech/stream") as ws:
|
||||
ws.send_json(
|
||||
{"type": "session.config", "output_mode": "codec_tokens", "stream_audio": True, "response_format": "pcm"}
|
||||
)
|
||||
ws.send_json({"type": "input.text", "text": "hello"})
|
||||
ws.send_json({"type": "input.done"})
|
||||
start = ws.receive_json()
|
||||
assert start["type"] == "codec.start" and start["codebooks"] == 16
|
||||
frames = []
|
||||
while True:
|
||||
msg = ws.receive()
|
||||
if "bytes" in msg and msg["bytes"] is not None:
|
||||
frames.append(cs.unpack_codec_frame(msg["bytes"], 16))
|
||||
continue
|
||||
done = __import__("json").loads(msg["text"])
|
||||
break
|
||||
assert [f[1] for f in frames] == [1, 2, 3]
|
||||
assert [f[0] for f in frames] == [0, 1, 2]
|
||||
assert frames[-1][2] & cs.FLAG_EOS
|
||||
assert done["type"] == "codec.done" and done["total_frames"] == 6 and done["cancelled"] is False
|
||||
assert ws.receive_json()["type"] == "session.done"
|
||||
svc.engine_client.abort.assert_not_awaited()
|
||||
|
||||
|
||||
def test_codec_stream_rejected_when_adapter_lacks_support(mocker):
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
app, _ = _codec_app(mocker, supports=False)
|
||||
with TestClient(app) as client, client.websocket_connect("/v1/audio/speech/stream") as ws:
|
||||
ws.send_json(
|
||||
{"type": "session.config", "output_mode": "codec_tokens", "stream_audio": True, "response_format": "pcm"}
|
||||
)
|
||||
ws.send_json({"type": "input.text", "text": "hello"})
|
||||
ws.send_json({"type": "input.done"})
|
||||
err = ws.receive_json()
|
||||
assert err["type"] == "error" and "not supported" in err["message"]
|
||||
|
||||
|
||||
def test_input_cancel_mid_stream_aborts_request(mocker):
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
app, svc = _codec_app(mocker)
|
||||
with TestClient(app) as client, client.websocket_connect("/v1/audio/speech/stream") as ws:
|
||||
ws.send_json(
|
||||
{"type": "session.config", "output_mode": "codec_tokens", "stream_audio": True, "response_format": "pcm"}
|
||||
)
|
||||
ws.send_json({"type": "input.text", "text": "hello"})
|
||||
ws.send_json({"type": "input.done"})
|
||||
assert ws.receive_json()["type"] == "codec.start"
|
||||
ws.send_json({"type": "input.cancel"})
|
||||
types = []
|
||||
while True:
|
||||
msg = ws.receive()
|
||||
if "bytes" in msg and msg["bytes"] is not None:
|
||||
continue
|
||||
payload = __import__("json").loads(msg["text"])
|
||||
types.append(payload["type"])
|
||||
if payload["type"] == "session.done":
|
||||
break
|
||||
if payload["type"] == "codec.done":
|
||||
assert payload["cancelled"] is True
|
||||
assert "codec.done" in types and types[-1] == "session.done"
|
||||
svc.engine_client.abort.assert_awaited()
|
||||
Reference in new issue
Block a user