[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,237 @@
|
||||
#!/usr/bin/env python3
|
||||
"""End-to-end codec-token streaming check on one device.
|
||||
|
||||
1. Launch a talker-only server (``qwen3_tts_talker_only.yaml``), run the
|
||||
edge decoder client for each prompt (decoder on ``--client-device``),
|
||||
record TTFA / playback-start / RTF / decode time and save WAVs.
|
||||
2. Stop it, launch the full two-stage server, request the same prompts with
|
||||
the same ``tts_local_seed`` through ``/v1/audio/speech`` and compute the
|
||||
SNR between server-decoded and client-decoded audio.
|
||||
|
||||
Both servers run one after another on the same device, so this fits a single
|
||||
``gpu run --gpus 1`` window (or a CPU host with ``VLLM_TARGET_DEVICE=cpu``).
|
||||
Results: ``benchmarks/tts/edge_results/codec_stream/<timestamp>/summary.json``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import datetime as _dt
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
BENCH_DIR = Path(__file__).resolve().parent
|
||||
REPO_ROOT = BENCH_DIR.parent.parent
|
||||
CLIENT_DIR = REPO_ROOT / "examples" / "online_serving" / "text_to_speech" / "qwen3_tts"
|
||||
sys.path.insert(0, str(CLIENT_DIR))
|
||||
sys.path.insert(0, str(BENCH_DIR))
|
||||
|
||||
|
||||
def _free_port() -> int:
|
||||
with socket.socket() as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return int(s.getsockname()[1])
|
||||
|
||||
|
||||
class Server:
|
||||
def __init__(
|
||||
self, model: str, deploy_config: str | None, port: int, log_path: Path, extra: list[str], env: dict[str, str]
|
||||
):
|
||||
cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"vllm_omni.entrypoints.cli.main",
|
||||
"serve",
|
||||
model,
|
||||
"--omni",
|
||||
"--port",
|
||||
str(port),
|
||||
"--trust-remote-code",
|
||||
]
|
||||
if deploy_config:
|
||||
cmd += ["--deploy-config", deploy_config]
|
||||
cmd += extra
|
||||
self.cmd, self.port, self.log_path = cmd, port, log_path
|
||||
self.env = {**os.environ, **env}
|
||||
self.proc: subprocess.Popen | None = None
|
||||
|
||||
def start(self, timeout_s: int) -> float:
|
||||
self.log = open(self.log_path, "w")
|
||||
self.proc = subprocess.Popen(
|
||||
self.cmd,
|
||||
stdout=self.log,
|
||||
stderr=subprocess.STDOUT,
|
||||
env=self.env,
|
||||
cwd=str(REPO_ROOT),
|
||||
start_new_session=True,
|
||||
)
|
||||
t0 = time.perf_counter()
|
||||
import urllib.request
|
||||
|
||||
while time.perf_counter() - t0 < timeout_s:
|
||||
if self.proc.poll() is not None:
|
||||
raise RuntimeError(f"server exited early (rc={self.proc.returncode}); see {self.log_path}")
|
||||
try:
|
||||
with urllib.request.urlopen(f"http://127.0.0.1:{self.port}/health", timeout=2) as r:
|
||||
if r.status == 200:
|
||||
return time.perf_counter() - t0
|
||||
except Exception:
|
||||
time.sleep(1.0)
|
||||
raise TimeoutError(f"server on port {self.port} not healthy after {timeout_s}s")
|
||||
|
||||
def stop(self) -> None:
|
||||
if self.proc and self.proc.poll() is None:
|
||||
os.killpg(self.proc.pid, signal.SIGTERM)
|
||||
try:
|
||||
self.proc.wait(timeout=60)
|
||||
except subprocess.TimeoutExpired:
|
||||
os.killpg(self.proc.pid, signal.SIGKILL)
|
||||
if getattr(self, "log", None):
|
||||
self.log.close()
|
||||
|
||||
|
||||
def snr_db(ref: np.ndarray, test: np.ndarray) -> float:
|
||||
n = min(ref.size, test.size)
|
||||
if n == 0:
|
||||
return float("nan")
|
||||
ref, test = ref[:n].astype(np.float64), test[:n].astype(np.float64)
|
||||
return 10.0 * np.log10((np.sum(ref**2) + 1e-12) / (np.sum((ref - test) ** 2) + 1e-12))
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
p.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice")
|
||||
p.add_argument("--prompts", default=str(BENCH_DIR / "prompts_12.txt"))
|
||||
p.add_argument("--n-prompts", type=int, default=4)
|
||||
p.add_argument("--seed", type=int, default=1234)
|
||||
p.add_argument("--voice", default="vivian")
|
||||
p.add_argument("--client-device", default="cpu")
|
||||
p.add_argument("--talker-deploy", default=str(REPO_ROOT / "vllm_omni" / "deploy" / "qwen3_tts_talker_only.yaml"))
|
||||
p.add_argument("--full-deploy", default=None, help="deploy YAML for the 2-stage server (default: model default)")
|
||||
p.add_argument("--server-timeout-s", type=int, default=1500)
|
||||
p.add_argument("--output-dir", default=str(BENCH_DIR / "edge_results" / "codec_stream"))
|
||||
p.add_argument("--skip-parity", action="store_true")
|
||||
args = p.parse_args(argv)
|
||||
|
||||
import edge_decoder_client as client
|
||||
|
||||
prompts = [line.strip() for line in open(args.prompts) if line.strip()][: args.n_prompts]
|
||||
ts = _dt.datetime.now().strftime("%Y%m%d-%H%M%S")
|
||||
out_dir = Path(args.output_dir) / ts
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
stage0_eager = ["--stage-overrides", json.dumps({"0": {"enforce_eager": True}})]
|
||||
summary: dict[str, Any] = {
|
||||
"model": args.model,
|
||||
"timestamp": ts,
|
||||
"prompts": prompts,
|
||||
"client_device": args.client_device,
|
||||
"rows": [],
|
||||
}
|
||||
|
||||
# --- talker-only server + client decoding
|
||||
port = _free_port()
|
||||
talker = Server(args.model, args.talker_deploy, port, out_dir / "talker_server.log", stage0_eager, {})
|
||||
try:
|
||||
summary["talker_server_start_s"] = talker.start(args.server_timeout_s)
|
||||
for i, text in enumerate(prompts):
|
||||
cargs = argparse.Namespace(
|
||||
host="127.0.0.1",
|
||||
port=port,
|
||||
model=args.model,
|
||||
text=text,
|
||||
voice=args.voice,
|
||||
task_type="CustomVoice",
|
||||
language=None,
|
||||
instructions=None,
|
||||
max_new_tokens=None,
|
||||
device=args.client_device,
|
||||
left_context=72,
|
||||
ramp=None,
|
||||
out=str(out_dir / f"client_{i}.wav"),
|
||||
json=str(out_dir / f"client_{i}.json"),
|
||||
seed=args.seed,
|
||||
)
|
||||
res = asyncio.run(client.run(cargs))
|
||||
summary["rows"].append(
|
||||
{
|
||||
"i": i,
|
||||
"text": text,
|
||||
"frames": res["total_frames"],
|
||||
"client_ttfa_ms": res["ttfa_ms"],
|
||||
"client_playback_start_ms": res["playback_start_ms"],
|
||||
"client_rtf": res["rtf"],
|
||||
"client_decode_ms": res["decode_ms_total"],
|
||||
"client_audio_s": res["audio_s"],
|
||||
}
|
||||
)
|
||||
print(
|
||||
f"[codec_e2e] talker-only prompt {i}: frames={res['total_frames']} "
|
||||
f"ttfa={res['ttfa_ms']} rtf={res['rtf']}",
|
||||
flush=True,
|
||||
)
|
||||
finally:
|
||||
talker.stop()
|
||||
|
||||
# --- full server for parity
|
||||
if not args.skip_parity:
|
||||
import soundfile as sf
|
||||
|
||||
port = _free_port()
|
||||
full = Server(args.model, args.full_deploy, port, out_dir / "full_server.log", stage0_eager, {})
|
||||
try:
|
||||
summary["full_server_start_s"] = full.start(args.server_timeout_s)
|
||||
import urllib.request
|
||||
|
||||
for row in summary["rows"]:
|
||||
body = json.dumps(
|
||||
{
|
||||
"model": args.model,
|
||||
"input": row["text"],
|
||||
"voice": args.voice,
|
||||
"task_type": "CustomVoice",
|
||||
"response_format": "wav",
|
||||
"seed": args.seed,
|
||||
}
|
||||
).encode()
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{port}/v1/audio/speech", data=body, headers={"Content-Type": "application/json"}
|
||||
)
|
||||
t0 = time.perf_counter()
|
||||
with urllib.request.urlopen(req, timeout=600) as r:
|
||||
wav = r.read()
|
||||
row["full_server_s"] = time.perf_counter() - t0
|
||||
ref, sr = sf.read(io.BytesIO(wav), dtype="float32")
|
||||
ref = ref.reshape(-1)
|
||||
sf.write(str(out_dir / f"full_{row['i']}.wav"), ref, sr, format="WAV")
|
||||
test, _ = sf.read(str(out_dir / f"client_{row['i']}.wav"), dtype="float32")
|
||||
row["full_audio_s"] = ref.size / sr
|
||||
row["snr_db"] = snr_db(ref, test.reshape(-1))
|
||||
row["length_ratio"] = (test.size / ref.size) if ref.size else None
|
||||
print(
|
||||
f"[codec_e2e] parity prompt {row['i']}: snr_db={row['snr_db']:.1f} len_ratio={row['length_ratio']}",
|
||||
flush=True,
|
||||
)
|
||||
finally:
|
||||
full.stop()
|
||||
|
||||
(out_dir / "summary.json").write_text(json.dumps(summary, indent=1))
|
||||
print(json.dumps({k: v for k, v in summary.items() if k != "rows"}, indent=1))
|
||||
for row in summary["rows"]:
|
||||
print(json.dumps(row))
|
||||
print(f"[codec_e2e] wrote {out_dir / 'summary.json'}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,65 @@
|
||||
# Codec-token streaming (on-device decoding)
|
||||
|
||||
For edge deployments the expensive part of a TTS pipeline that must run near
|
||||
the user is the audio decoder, while the autoregressive talker can stay on a
|
||||
server. Codec-token streaming serves the **talker only** and streams its codec
|
||||
frames to a client that decodes them locally (Qwen3-TTS: 16 codebooks at
|
||||
12.5 Hz, about 400 B/s of payload before framing).
|
||||
|
||||
## Server
|
||||
|
||||
Deploy the talker-only pipeline:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice --omni --port 8001 \
|
||||
--deploy-config qwen3_tts_talker_only.yaml --trust-remote-code
|
||||
```
|
||||
|
||||
`qwen3_tts_talker_only.yaml` selects the registered `qwen3_tts_talker_only`
|
||||
pipeline (stage 0 only, `final_output_type: latent`). The talker publishes its
|
||||
per-step `codes.audio` rows to the client (`omni_client_multimodal_output_keys`),
|
||||
and the WebSocket handler emits them as binary frames.
|
||||
|
||||
## Protocol (`/v1/audio/speech/stream`)
|
||||
|
||||
Client → server: the usual `session.config` with
|
||||
|
||||
```json
|
||||
{"type": "session.config", "output_mode": "codec_tokens", "stream_audio": true,
|
||||
"response_format": "pcm", "voice": "vivian", "task_type": "CustomVoice"}
|
||||
```
|
||||
|
||||
then `input.text` / `input.done`. `input.cancel` aborts the in-flight
|
||||
generation (works for PCM sessions too).
|
||||
|
||||
Server → client per utterance:
|
||||
|
||||
| message | content |
|
||||
|---|---|
|
||||
| `codec.start` (JSON) | `codebooks`, `codebook_size`, `frame_rate_hz`, `sample_rate`, `decoder_id`, `header_format` |
|
||||
| binary frames | header `<IHB` = `seq:uint32, n_frames:uint16, flags:uint8`, then `int16[n_frames × codebooks]`; flags `EOS=1`, `SEGMENT_END=2`, `RESET=4` |
|
||||
| `codec.done` (JSON) | `total_frames`, `seq_last`, `error`, `cancelled` |
|
||||
| `session.done` | as for PCM sessions |
|
||||
|
||||
Frames are emitted as they are produced (one talker step = one frame); the
|
||||
chunking policy for the decoder (first small chunk, ramp to 25 frames, 72-frame
|
||||
left context) is the client's choice. The wire format lives in
|
||||
`vllm_omni/entrypoints/openai/codec_stream.py`.
|
||||
|
||||
## Client
|
||||
|
||||
`examples/online_serving/text_to_speech/qwen3_tts/edge_decoder_client.py`
|
||||
loads `Qwen3TTSTokenizerV2Decoder` from the checkpoint's `speech_tokenizer/`
|
||||
folder (CPU or CUDA), applies the same chunk schedule and state carry-over as
|
||||
the server's Code2Wav stage, writes a WAV and reports TTFA, playback-start
|
||||
latency and RTF. `codec_parity_check.py` compares the client-decoded audio with
|
||||
the full two-stage server's output (same `tts_local_seed`) and reports SNR.
|
||||
|
||||
## Limits
|
||||
|
||||
- Only adapters with `supports_codec_stream = True` (Qwen3-TTS today) accept
|
||||
`output_mode: codec_tokens`; others return an `error` message.
|
||||
- The talker's per-step output carries the cumulative code tensor, so a long
|
||||
utterance costs O(N²) bytes between engine and API server (about 9 MB for
|
||||
30 s of audio); the wire stream itself is incremental.
|
||||
- `word_timestamps` and `speed != 1.0` are not available in this mode.
|
||||
@@ -0,0 +1,123 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Parity check: server-decoded PCM vs client-decoded PCM from codec tokens.
|
||||
|
||||
Run the full 2-stage server and the talker-only server one after another (or
|
||||
on two ports) with ``enforce_eager`` on stage 0 so per-request seeds are
|
||||
reproducible, then:
|
||||
|
||||
python codec_parity_check.py --model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \\
|
||||
--full-port 8000 --talker-port 8001 --seed 1234 --text "..."
|
||||
|
||||
Reports the length-trimmed SNR (dB) between the two waveforms and the codec
|
||||
frame count. Both requests use the same ``seed``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import io
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(HERE))
|
||||
|
||||
|
||||
def snr_db(ref: np.ndarray, test: np.ndarray) -> float:
|
||||
n = min(ref.size, test.size)
|
||||
if n == 0:
|
||||
return float("nan")
|
||||
ref, test = ref[:n].astype(np.float64), test[:n].astype(np.float64)
|
||||
noise = ref - test
|
||||
p_sig = float(np.sum(ref**2)) + 1e-12
|
||||
p_noise = float(np.sum(noise**2)) + 1e-12
|
||||
return 10.0 * np.log10(p_sig / p_noise)
|
||||
|
||||
|
||||
def fetch_full_pcm(host: str, port: int, model: str, text: str, voice: str, seed: int) -> np.ndarray:
|
||||
import requests
|
||||
|
||||
body = {
|
||||
"model": model,
|
||||
"input": text,
|
||||
"voice": voice,
|
||||
"response_format": "wav",
|
||||
"seed": seed,
|
||||
}
|
||||
r = requests.post(f"http://{host}:{port}/v1/audio/speech", json=body, timeout=600)
|
||||
r.raise_for_status()
|
||||
import soundfile as sf
|
||||
|
||||
data, sr = sf.read(io.BytesIO(r.content), dtype="float32")
|
||||
return data.reshape(-1)
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
p.add_argument("--model", required=True)
|
||||
p.add_argument("--host", default="localhost")
|
||||
p.add_argument("--full-port", type=int, default=8000)
|
||||
p.add_argument("--talker-port", type=int, default=8001)
|
||||
p.add_argument("--voice", default="vivian")
|
||||
p.add_argument("--text", default="Parity check between server decoding and on-device decoding.")
|
||||
p.add_argument("--seed", type=int, default=1234)
|
||||
p.add_argument("--device", default="cpu")
|
||||
p.add_argument("--full-wav", default=None, help="use a previously saved full-server WAV instead of requesting one")
|
||||
p.add_argument("--out-dir", default=None)
|
||||
args = p.parse_args(argv)
|
||||
|
||||
import edge_decoder_client as client
|
||||
|
||||
out_dir = Path(args.out_dir) if args.out_dir else None
|
||||
if out_dir:
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
if args.full_wav:
|
||||
import soundfile as sf
|
||||
|
||||
ref, _ = sf.read(args.full_wav, dtype="float32")
|
||||
ref = ref.reshape(-1)
|
||||
else:
|
||||
ref = fetch_full_pcm(args.host, args.full_port, args.model, args.text, args.voice, args.seed)
|
||||
client_args = argparse.Namespace(
|
||||
host=args.host,
|
||||
port=args.talker_port,
|
||||
model=args.model,
|
||||
text=args.text,
|
||||
voice=args.voice,
|
||||
task_type="CustomVoice",
|
||||
language=None,
|
||||
instructions=None,
|
||||
max_new_tokens=None,
|
||||
device=args.device,
|
||||
left_context=72,
|
||||
ramp=None,
|
||||
seed=args.seed,
|
||||
out=str(out_dir / "client.wav") if out_dir else None,
|
||||
json=str(out_dir / "client.json") if out_dir else None,
|
||||
)
|
||||
result = asyncio.run(client.run(client_args))
|
||||
import soundfile as sf
|
||||
|
||||
test, _ = sf.read(client_args.out, dtype="float32") if client_args.out else (np.zeros(0, dtype=np.float32), 0)
|
||||
report = {
|
||||
"ref_samples": int(ref.size),
|
||||
"client_samples": int(test.size),
|
||||
"snr_db": snr_db(ref, test),
|
||||
"length_ratio": (test.size / ref.size) if ref.size else None,
|
||||
"client_ttfa_ms": result.get("ttfa_ms"),
|
||||
"client_rtf": result.get("rtf"),
|
||||
"frames": result.get("total_frames"),
|
||||
}
|
||||
print(json.dumps(report, indent=1))
|
||||
if out_dir:
|
||||
(out_dir / "parity.json").write_text(json.dumps(report, indent=1))
|
||||
sf.write(str(out_dir / "full.wav"), ref, 24000, format="WAV")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,248 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Codec-token streaming client with on-device decoding (Qwen3-TTS).
|
||||
|
||||
Connects to ``/v1/audio/speech/stream`` of a **talker-only** deployment
|
||||
(``--deploy-config qwen3_tts_talker_only.yaml``), requests
|
||||
``output_mode="codec_tokens"``, receives 16-codebook frames at 12.5 Hz and
|
||||
decodes them locally with ``Qwen3TTSTokenizerV2Decoder`` (from the model's
|
||||
``speech_tokenizer/`` folder) using the same chunk schedule / 72-frame left
|
||||
context as the server-side Code2Wav stage. Measures TTFA (first PCM out),
|
||||
per-chunk decode time, playback-start latency and RTF, and writes a WAV.
|
||||
|
||||
Example:
|
||||
python edge_decoder_client.py --host localhost --port 8000 \\
|
||||
--model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice --voice vivian \\
|
||||
--text "Hello from the edge." --device cpu --out out.wav
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
sys.path.insert(0, str(REPO_ROOT / "benchmarks" / "tts"))
|
||||
|
||||
from vllm_omni.entrypoints.openai.codec_stream import FLAG_EOS, FLAG_RESET, unpack_codec_frame # noqa: E402
|
||||
from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp # noqa: E402
|
||||
|
||||
DEFAULT_RAMP = [2, 4, 8, 16, 25]
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ decoder
|
||||
class LocalDecoder:
|
||||
"""Stateful chunked Qwen3-TTS decoder mirroring the Code2Wav stage."""
|
||||
|
||||
def __init__(
|
||||
self, model: str, device: str = "cpu", left_context: int = 72, chunk_frames: int = 25, dtype: str = "float32"
|
||||
):
|
||||
from safetensors.torch import load_file
|
||||
|
||||
from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (
|
||||
Qwen3TTSTokenizerV2Config,
|
||||
)
|
||||
from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
|
||||
Qwen3TTSTokenizerV2Decoder,
|
||||
)
|
||||
|
||||
model_dir = self._resolve_dir(model)
|
||||
cfg = Qwen3TTSTokenizerV2Config.from_pretrained(str(model_dir), subfolder="speech_tokenizer")
|
||||
self.decoder = Qwen3TTSTokenizerV2Decoder._from_config(cfg.decoder_config)
|
||||
state = load_file(str(model_dir / "speech_tokenizer" / "model.safetensors"))
|
||||
dec_state = {k[len("decoder.") :]: v for k, v in state.items() if k.startswith("decoder.")}
|
||||
missing, unexpected = self.decoder.load_state_dict(dec_state, strict=False)
|
||||
if unexpected:
|
||||
raise RuntimeError(f"unexpected decoder keys: {unexpected[:5]}")
|
||||
if missing:
|
||||
print(f"[decoder] warning: {len(missing)} missing keys (first: {missing[:3]})", file=sys.stderr)
|
||||
self.decoder.eval().to(device=device, dtype=getattr(torch, dtype))
|
||||
if hasattr(self.decoder, "precompute_snake_caches"):
|
||||
self.decoder.precompute_snake_caches()
|
||||
if hasattr(self.decoder, "_incremental_chunk_frames"):
|
||||
self.decoder._incremental_chunk_frames = chunk_frames
|
||||
self.device = device
|
||||
self.num_quantizers = int(cfg.decoder_config.num_quantizers)
|
||||
self.sample_rate = int(cfg.output_sample_rate)
|
||||
self.left_context = left_context
|
||||
self.chunk_frames = chunk_frames
|
||||
self.reset()
|
||||
|
||||
@staticmethod
|
||||
def _resolve_dir(model: str) -> Path:
|
||||
p = Path(model)
|
||||
if p.is_dir():
|
||||
return p
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
return Path(snapshot_download(model, local_files_only=True))
|
||||
|
||||
def reset(self) -> None:
|
||||
self.state: dict[str, Any] = {"prefix_frames": 0}
|
||||
self.frames_decoded = 0
|
||||
|
||||
@torch.inference_mode()
|
||||
def decode(self, codes: np.ndarray) -> np.ndarray:
|
||||
"""Decode ``[n_frames, Q]`` new frames, carrying state across calls; returns float32 PCM."""
|
||||
if codes.shape[0] == 0:
|
||||
return np.zeros(0, dtype=np.float32)
|
||||
qf = torch.as_tensor(codes.T, dtype=torch.long, device=self.device) # [Q, F]
|
||||
wavs = self.decoder.batched_chunked_decode(
|
||||
[qf],
|
||||
[int(qf.shape[-1])],
|
||||
caches=[self.state],
|
||||
chunk_size=self.chunk_frames,
|
||||
left_context_size=self.left_context,
|
||||
max_batch_size=1,
|
||||
)
|
||||
wav = wavs[0]
|
||||
if wav.dim() == 2 and wav.shape[0] == 1:
|
||||
wav = wav[0]
|
||||
self.frames_decoded += int(codes.shape[0])
|
||||
return wav.float().cpu().numpy().reshape(-1)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- client
|
||||
async def run(args: argparse.Namespace) -> dict[str, Any]:
|
||||
import websockets
|
||||
|
||||
ramp = args.ramp or DEFAULT_RAMP
|
||||
schedule = parse_chunk_ramp({"codec_chunk_ramp": ramp}, steady=ramp[-1]) or ramp
|
||||
steady = schedule[-1]
|
||||
decoder = LocalDecoder(args.model, device=args.device, left_context=args.left_context, chunk_frames=steady)
|
||||
url = f"ws://{args.host}:{args.port}/v1/audio/speech/stream"
|
||||
config: dict[str, Any] = {
|
||||
"type": "session.config",
|
||||
"model": args.model,
|
||||
"output_mode": "codec_tokens",
|
||||
"stream_audio": True,
|
||||
"response_format": "pcm",
|
||||
}
|
||||
for key in ("voice", "task_type", "language", "instructions", "max_new_tokens", "seed"):
|
||||
val = getattr(args, key, None)
|
||||
if val is not None:
|
||||
config[key] = val
|
||||
pcm_parts: list[np.ndarray] = []
|
||||
chunk_log: list[dict[str, Any]] = []
|
||||
pending = np.zeros((0, decoder.num_quantizers), dtype=np.int64)
|
||||
chunk_idx = 0
|
||||
t0 = None
|
||||
t_first_pcm = None
|
||||
codec_start = None
|
||||
total_frames = 0
|
||||
async with websockets.connect(url, max_size=64 * 1024 * 1024) as ws:
|
||||
await ws.send(json.dumps(config))
|
||||
await ws.send(json.dumps({"type": "input.text", "text": args.text}))
|
||||
t0 = time.perf_counter()
|
||||
await ws.send(json.dumps({"type": "input.done"}))
|
||||
while True:
|
||||
raw = await ws.recv()
|
||||
now = time.perf_counter()
|
||||
if isinstance(raw, (bytes, bytearray)):
|
||||
seq, n, flags, arr = unpack_codec_frame(bytes(raw), decoder.num_quantizers)
|
||||
if flags & FLAG_RESET:
|
||||
decoder.reset()
|
||||
pending = pending[:0]
|
||||
total_frames += n
|
||||
if n:
|
||||
pending = np.concatenate([pending, arr.astype(np.int64)], axis=0)
|
||||
eos = bool(flags & FLAG_EOS)
|
||||
target = schedule[min(chunk_idx, len(schedule) - 1)]
|
||||
while pending.shape[0] >= target or (eos and pending.shape[0] > 0):
|
||||
take = pending[: min(target, steady)] if not eos else pending[:steady]
|
||||
pending = pending[take.shape[0] :]
|
||||
td = time.perf_counter()
|
||||
pcm = decoder.decode(take)
|
||||
t_dec = time.perf_counter() - td
|
||||
if t_first_pcm is None and pcm.size:
|
||||
t_first_pcm = time.perf_counter()
|
||||
pcm_parts.append(pcm)
|
||||
chunk_log.append(
|
||||
{
|
||||
"idx": chunk_idx,
|
||||
"t_ms": (time.perf_counter() - t0) * 1e3,
|
||||
"arrival_ms": (now - t0) * 1e3,
|
||||
"frames": int(take.shape[0]),
|
||||
"samples": int(pcm.size),
|
||||
"decode_ms": t_dec * 1e3,
|
||||
}
|
||||
)
|
||||
chunk_idx += 1
|
||||
target = schedule[min(chunk_idx, len(schedule) - 1)]
|
||||
if eos:
|
||||
continue
|
||||
else:
|
||||
msg = json.loads(raw)
|
||||
mtype = msg.get("type")
|
||||
if mtype == "codec.start":
|
||||
codec_start = msg
|
||||
elif mtype in ("codec.done", "audio.done"):
|
||||
pass
|
||||
elif mtype == "session.done":
|
||||
break
|
||||
elif mtype == "error":
|
||||
raise RuntimeError(f"server error: {msg.get('message')}")
|
||||
await ws.send(json.dumps({"type": "session.close"}))
|
||||
import stream_latency_metrics as metrics
|
||||
|
||||
sr = decoder.sample_rate
|
||||
audio = np.concatenate(pcm_parts) if pcm_parts else np.zeros(0, dtype=np.float32)
|
||||
total_s = (time.perf_counter() - t0) if t0 else 0.0
|
||||
result = {
|
||||
"codec_start": codec_start,
|
||||
"schedule": schedule,
|
||||
"total_frames": total_frames,
|
||||
"ttfa_ms": ((t_first_pcm - t0) * 1e3) if t_first_pcm else None,
|
||||
"playback_start_ms": metrics.playback_start_ms(chunk_log, sr),
|
||||
"stall_at_ttfa_ms": metrics.stall_at_ttfa_ms(chunk_log, sr),
|
||||
"audio_s": audio.size / sr,
|
||||
"total_s": total_s,
|
||||
"rtf": (total_s / (audio.size / sr)) if audio.size else None,
|
||||
"decode_ms_total": sum(c["decode_ms"] for c in chunk_log),
|
||||
"chunks": chunk_log,
|
||||
"device": args.device,
|
||||
}
|
||||
if args.out:
|
||||
import soundfile as sf
|
||||
|
||||
sf.write(args.out, audio, sr, format="WAV")
|
||||
result["wav"] = args.out
|
||||
if args.json:
|
||||
Path(args.json).write_text(json.dumps(result, indent=1))
|
||||
print(json.dumps({k: v for k, v in result.items() if k != "chunks"}, indent=1))
|
||||
return result
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
p.add_argument("--host", default="localhost")
|
||||
p.add_argument("--port", type=int, default=8000)
|
||||
p.add_argument(
|
||||
"--model", required=True, help="HF id or local dir (decoder weights are read from speech_tokenizer/)"
|
||||
)
|
||||
p.add_argument("--text", default="Hello, this is a codec token streaming test from the edge decoder client.")
|
||||
p.add_argument("--voice", default="vivian")
|
||||
p.add_argument("--task-type", dest="task_type", default="CustomVoice")
|
||||
p.add_argument("--language", default=None)
|
||||
p.add_argument("--instructions", default=None)
|
||||
p.add_argument("--max-new-tokens", dest="max_new_tokens", type=int, default=None)
|
||||
p.add_argument("--seed", type=int, default=None)
|
||||
p.add_argument("--device", default="cpu")
|
||||
p.add_argument("--left-context", type=int, default=72)
|
||||
p.add_argument("--ramp", type=int, nargs="*", default=None, help="chunk schedule in frames, default 2 4 8 16 25")
|
||||
p.add_argument("--out", default=None)
|
||||
p.add_argument("--json", default=None)
|
||||
args = p.parse_args(argv)
|
||||
asyncio.run(run(args))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,38 @@
|
||||
"""Talker-only Qwen3-TTS pipeline registration (CPU)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
|
||||
from vllm_omni.config.stage_config import _DEPLOY_DIR, StageExecutionType, load_deploy_config, merge_pipeline_deploy
|
||||
from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE, QWEN3_TTS_TALKER_ONLY_PIPELINE
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_registered_single_terminal_stage():
|
||||
assert OMNI_PIPELINES["qwen3_tts_talker_only"] is QWEN3_TTS_TALKER_ONLY_PIPELINE
|
||||
p = QWEN3_TTS_TALKER_ONLY_PIPELINE
|
||||
assert len(p.stages) == 1
|
||||
s = p.stages[0]
|
||||
assert s.execution_type is StageExecutionType.LLM_AR
|
||||
assert s.final_output is True and s.final_output_type == "latent" and s.engine_output_type == "latent"
|
||||
assert s.model_stage == QWEN3_TTS_PIPELINE.stages[0].model_stage == "qwen3_tts"
|
||||
assert s.custom_process_next_stage_input_func is None
|
||||
assert s.async_chunk_process_next_stage_input_func is None
|
||||
assert s.sampling_constraints["stop_token_ids"] == [2150]
|
||||
|
||||
|
||||
def test_deploy_yaml_merges_to_one_stage_config():
|
||||
deploy_path = _DEPLOY_DIR / "qwen3_tts_talker_only.yaml"
|
||||
assert deploy_path.exists()
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
assert len(deploy.stages) == 1 and deploy.stages[0].devices == "0"
|
||||
merged = merge_pipeline_deploy(QWEN3_TTS_TALKER_ONLY_PIPELINE, deploy)
|
||||
stage_configs = merged if isinstance(merged, list) else getattr(merged, "stages", merged)
|
||||
assert len(stage_configs) == 1
|
||||
cfg = stage_configs[0]
|
||||
engine_args = getattr(cfg, "engine_args", cfg)
|
||||
assert bool(getattr(engine_args, "async_chunk", False)) is False
|
||||
assert getattr(engine_args, "custom_process_next_stage_input_func", None) is None
|
||||
@@ -0,0 +1,107 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""E2E: codec-token streaming from a talker-only Qwen3-TTS deployment.
|
||||
|
||||
Deploy Setting: qwen3_tts_talker_only.yaml (single LLM_AR stage)
|
||||
Input Modal: text (WebSocket /v1/audio/speech/stream, output_mode=codec_tokens)
|
||||
Output Modal: codec frames decoded on the client with Qwen3TTSTokenizerV2Decoder
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServerParams
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni.entrypoints.openai.codec_stream import FLAG_EOS, unpack_codec_frame
|
||||
|
||||
MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice"
|
||||
_STAGE_CONFIG = get_deploy_config_path("qwen3_tts_talker_only.yaml")
|
||||
|
||||
server_params = [
|
||||
pytest.param(
|
||||
OmniServerParams(model=MODEL, stage_config_path=_STAGE_CONFIG, server_args=["--trust-remote-code"]),
|
||||
id="talker_only",
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def _stream_codec(port: int, text: str) -> tuple[dict, list[tuple[int, int, int, np.ndarray]], dict]:
|
||||
from websockets.sync.client import connect
|
||||
|
||||
frames = []
|
||||
with connect(f"ws://localhost:{port}/v1/audio/speech/stream", max_size=64 * 1024 * 1024) as ws:
|
||||
ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.config",
|
||||
"model": MODEL,
|
||||
"output_mode": "codec_tokens",
|
||||
"stream_audio": True,
|
||||
"response_format": "pcm",
|
||||
"voice": "vivian",
|
||||
"task_type": "CustomVoice",
|
||||
}
|
||||
)
|
||||
)
|
||||
ws.send(json.dumps({"type": "input.text", "text": text}))
|
||||
ws.send(json.dumps({"type": "input.done"}))
|
||||
start = json.loads(ws.recv())
|
||||
assert start["type"] == "codec.start", start
|
||||
done = None
|
||||
while True:
|
||||
raw = ws.recv()
|
||||
if isinstance(raw, (bytes, bytearray)):
|
||||
frames.append(unpack_codec_frame(bytes(raw), start["codebooks"]))
|
||||
continue
|
||||
msg = json.loads(raw)
|
||||
if msg["type"] == "codec.done":
|
||||
done = msg
|
||||
elif msg["type"] == "session.done":
|
||||
break
|
||||
elif msg["type"] == "error":
|
||||
raise AssertionError(msg)
|
||||
ws.send(json.dumps({"type": "session.close"}))
|
||||
return start, frames, done
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", server_params, indirect=True)
|
||||
def test_codec_token_stream_and_client_decode(omni_server) -> None:
|
||||
start, frames, done = _stream_codec(omni_server.port, "Codec token streaming end to end test.")
|
||||
assert start["codebooks"] == 16 and start["frame_rate_hz"] == 12.5 and start["sample_rate"] == 24000
|
||||
assert frames, "no codec frames received"
|
||||
seqs = [f[0] for f in frames]
|
||||
assert seqs == list(range(len(seqs)))
|
||||
assert frames[-1][2] & FLAG_EOS
|
||||
total = sum(f[1] for f in frames)
|
||||
assert total >= 10, f"too few frames: {total}"
|
||||
assert done is not None and done["total_frames"] == total and done["error"] is False
|
||||
codes = np.concatenate([f[3].astype(np.int64) for f in frames if f[1]], axis=0)
|
||||
assert codes.min() >= 0 and codes.max() < 2048
|
||||
|
||||
import sys
|
||||
|
||||
sys.path.insert(
|
||||
0,
|
||||
str(
|
||||
os.path.join(
|
||||
os.path.dirname(__file__), "..", "..", "..", "examples", "online_serving", "text_to_speech", "qwen3_tts"
|
||||
)
|
||||
),
|
||||
)
|
||||
from edge_decoder_client import LocalDecoder
|
||||
|
||||
dec = LocalDecoder(MODEL, device="cpu")
|
||||
pcm = np.concatenate([dec.decode(codes[i : i + 25]) for i in range(0, codes.shape[0], 25)])
|
||||
expected_s = total / 12.5
|
||||
assert abs(pcm.size / dec.sample_rate - expected_s) < 0.1 * expected_s + 0.2
|
||||
assert float(np.abs(pcm).max()) > 1e-3
|
||||
@@ -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()
|
||||
@@ -107,7 +107,7 @@ from vllm_omni.model_executor.models.qwen3_omni.pipeline import (
|
||||
QWEN3_OMNI_THINKER_ONLY_PIPELINE,
|
||||
resolve_qwen3_omni_pipeline,
|
||||
)
|
||||
from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE
|
||||
from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE, QWEN3_TTS_TALKER_ONLY_PIPELINE
|
||||
from vllm_omni.model_executor.models.step_audio2.pipeline import (
|
||||
STEP_AUDIO2_ASR_PIPELINE,
|
||||
STEP_AUDIO2_PIPELINE,
|
||||
@@ -134,6 +134,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
|
||||
"qwen3_omni_moe": resolve_qwen3_omni_pipeline,
|
||||
"qwen3_omni_moe_thinker_only": QWEN3_OMNI_THINKER_ONLY_PIPELINE,
|
||||
"qwen3_tts": QWEN3_TTS_PIPELINE,
|
||||
"qwen3_tts_talker_only": QWEN3_TTS_TALKER_ONLY_PIPELINE,
|
||||
"step_audio_2": STEP_AUDIO2_PIPELINE,
|
||||
"step_audio_2_asr": STEP_AUDIO2_ASR_PIPELINE,
|
||||
"covo_audio": COVO_AUDIO_PIPELINE,
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
# Qwen3-TTS talker-only deployment: stage 0 is terminal and publishes codec
|
||||
# frames (codes.audio, 16 codebooks @ 12.5 Hz) to the client instead of PCM.
|
||||
# Used by the codec-token streaming mode of /v1/audio/speech/stream
|
||||
# (session.config.output_mode: "codec_tokens") with decoding on the client
|
||||
# device (see examples/online_serving/text_to_speech/qwen3_tts/edge_decoder_client.py).
|
||||
# Stage settings mirror stage 0 of qwen3_tts.yaml.
|
||||
pipeline: qwen3_tts_talker_only
|
||||
async_chunk: false
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_seqs: 64
|
||||
gpu_memory_utilization: 0.3
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: true
|
||||
max_num_batched_tokens: 32768
|
||||
max_model_len: 4096
|
||||
devices: '0'
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
min_tokens: 2
|
||||
repetition_penalty: 1.05
|
||||
silence_ban_frames: 0
|
||||
subtalker_sampling_params:
|
||||
do_sample: true
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
top_p: 1.0
|
||||
@@ -0,0 +1,129 @@
|
||||
"""Wire format for codec-token streaming (server → on-device decoder).
|
||||
|
||||
Binary frame layout (little-endian):
|
||||
|
||||
header ``<IHB`` seq:uint32 n_frames:uint16 flags:uint8 (7 bytes)
|
||||
body int16[n_frames * codebooks] row-major (frame-major)
|
||||
|
||||
Flags: ``FLAG_EOS`` (last frame of the utterance), ``FLAG_SEGMENT_END``
|
||||
(sentence boundary inside a flush), ``FLAG_RESET`` (client must drop its
|
||||
decoder state before consuming this frame). An EOS frame may carry
|
||||
``n_frames == 0``.
|
||||
|
||||
Codebook ids are ``< 32768`` for every supported codec (Qwen3-TTS: 2048), so
|
||||
int16 is lossless. Pure functions only; no engine imports.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import struct
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
HEADER_FMT = "<IHB"
|
||||
HEADER_SIZE = struct.calcsize(HEADER_FMT)
|
||||
FLAG_EOS = 0x1
|
||||
FLAG_SEGMENT_END = 0x2
|
||||
FLAG_RESET = 0x4
|
||||
MAX_FRAMES_PER_PACKET = 0xFFFF
|
||||
CODEC_STREAM_VERSION = 1
|
||||
|
||||
|
||||
def pack_codec_frame(
|
||||
seq: int,
|
||||
codes: Any,
|
||||
*,
|
||||
eos: bool = False,
|
||||
segment_end: bool = False,
|
||||
reset: bool = False,
|
||||
codebooks: int | None = None,
|
||||
) -> bytes:
|
||||
"""Pack ``codes`` (``[n_frames, codebooks]`` ints, or ``None``/empty) into one binary frame."""
|
||||
if codes is None:
|
||||
arr = np.zeros((0, codebooks or 0), dtype=np.int16)
|
||||
else:
|
||||
if hasattr(codes, "detach"):
|
||||
codes = codes.detach().cpu().numpy()
|
||||
arr = np.asarray(codes)
|
||||
if arr.ndim == 1:
|
||||
arr = arr.reshape(1, -1) if arr.size else arr.reshape(0, codebooks or 0)
|
||||
if arr.ndim != 2:
|
||||
raise ValueError(f"codes must be 2-D [n_frames, codebooks], got shape {arr.shape}")
|
||||
if arr.size and (arr.min() < 0 or arr.max() > np.iinfo(np.int16).max):
|
||||
raise ValueError("codec ids must fit in int16")
|
||||
arr = arr.astype("<i2", copy=False)
|
||||
n_frames = int(arr.shape[0])
|
||||
if n_frames > MAX_FRAMES_PER_PACKET:
|
||||
raise ValueError(f"too many frames for one packet: {n_frames}")
|
||||
flags = (FLAG_EOS if eos else 0) | (FLAG_SEGMENT_END if segment_end else 0) | (FLAG_RESET if reset else 0)
|
||||
return struct.pack(HEADER_FMT, int(seq) & 0xFFFFFFFF, n_frames, flags) + arr.tobytes(order="C")
|
||||
|
||||
|
||||
def unpack_codec_frame(buf: bytes, codebooks: int) -> tuple[int, int, int, np.ndarray]:
|
||||
"""Inverse of :func:`pack_codec_frame`: ``(seq, n_frames, flags, int16[n_frames, codebooks])``."""
|
||||
if len(buf) < HEADER_SIZE:
|
||||
raise ValueError("codec frame shorter than its header")
|
||||
seq, n_frames, flags = struct.unpack(HEADER_FMT, buf[:HEADER_SIZE])
|
||||
body = buf[HEADER_SIZE:]
|
||||
expected = n_frames * codebooks * 2
|
||||
if len(body) != expected:
|
||||
raise ValueError(f"codec frame body has {len(body)} bytes, expected {expected} for {n_frames}x{codebooks}")
|
||||
arr = (
|
||||
np.frombuffer(body, dtype="<i2").reshape(n_frames, codebooks)
|
||||
if n_frames
|
||||
else np.zeros((0, codebooks), dtype=np.int16)
|
||||
)
|
||||
return seq, n_frames, flags, arr
|
||||
|
||||
|
||||
def align_codec_rows(codes_cum: Any, token_ids_cum: list[int], codebook_size: int) -> np.ndarray:
|
||||
"""Keep the codec rows that correspond to real codebook-0 tokens.
|
||||
|
||||
Mirrors the talker's full-payload filter: rows are tail-aligned to the
|
||||
output token ids, the prefill placeholder row (and any row whose token is
|
||||
not a codebook id, e.g. EOS) is dropped, and rows with out-of-range ids
|
||||
are dropped. Returns ``int64[n, codebooks]`` (possibly empty).
|
||||
"""
|
||||
if codes_cum is None:
|
||||
return np.zeros((0, 0), dtype=np.int64)
|
||||
if hasattr(codes_cum, "detach"):
|
||||
codes_cum = codes_cum.detach().cpu().numpy()
|
||||
arr = np.asarray(codes_cum)
|
||||
if arr.ndim != 2 or arr.size == 0:
|
||||
return np.zeros((0, arr.shape[1] if arr.ndim == 2 else 0), dtype=np.int64)
|
||||
ids = list(token_ids_cum)
|
||||
while ids and ids[-1] == -1:
|
||||
ids.pop()
|
||||
n = min(arr.shape[0], len(ids))
|
||||
if n <= 0:
|
||||
return arr[:0].astype(np.int64)
|
||||
rows = arr[-n:]
|
||||
toks = np.asarray(ids[-n:])
|
||||
tok_ok = (toks >= 0) & (toks < codebook_size)
|
||||
row_ok = (rows.min(axis=1) >= 0) & (rows.max(axis=1) < codebook_size)
|
||||
return rows[tok_ok & row_ok].astype(np.int64)
|
||||
|
||||
|
||||
def codec_start_message(
|
||||
*,
|
||||
utterance_index: int,
|
||||
sentence_index: int,
|
||||
sentence_text: str,
|
||||
spec: dict[str, Any],
|
||||
model: str | None,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"type": "codec.start",
|
||||
"version": CODEC_STREAM_VERSION,
|
||||
"utterance_index": utterance_index,
|
||||
"sentence_index": sentence_index,
|
||||
"sentence_text": sentence_text,
|
||||
"model": model,
|
||||
"codebooks": int(spec["codebooks"]),
|
||||
"codebook_size": int(spec["codebook_size"]),
|
||||
"frame_rate_hz": float(spec["frame_rate_hz"]),
|
||||
"sample_rate": int(spec["sample_rate"]),
|
||||
"decoder_id": spec.get("decoder_id"),
|
||||
"header_format": HEADER_FMT,
|
||||
}
|
||||
@@ -578,9 +578,26 @@ class StreamingSpeechSessionConfig(BaseModel):
|
||||
"frames (existing behavior)."
|
||||
),
|
||||
)
|
||||
seed: int | None = Field(default=None, description="Random seed forwarded to each utterance request.")
|
||||
output_mode: Literal["pcm", "codec_tokens"] = Field(
|
||||
default="pcm",
|
||||
description=(
|
||||
"'pcm' streams decoded audio. 'codec_tokens' streams the talker's codec frames "
|
||||
"(binary frames described by a 'codec.start' message) for decoding on the client "
|
||||
"device; requires a talker-only deployment (e.g. deploy config qwen3_tts_talker_only.yaml), "
|
||||
"stream_audio=true, response_format='pcm' and word_timestamps=false."
|
||||
),
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_streaming_constraints(self) -> "StreamingSpeechSessionConfig":
|
||||
if self.output_mode == "codec_tokens":
|
||||
if not self.stream_audio:
|
||||
raise ValueError("output_mode='codec_tokens' requires stream_audio=true.")
|
||||
if self.word_timestamps:
|
||||
raise ValueError("output_mode='codec_tokens' does not support word_timestamps.")
|
||||
if self.response_format != "pcm":
|
||||
raise ValueError("output_mode='codec_tokens' requires response_format='pcm'.")
|
||||
if self.stream_audio:
|
||||
if self.response_format != "pcm":
|
||||
raise ValueError(
|
||||
|
||||
@@ -3251,6 +3251,77 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
):
|
||||
yield chunk
|
||||
|
||||
@staticmethod
|
||||
def _extract_codec_output(res) -> Any:
|
||||
"""Return the cumulative ``codes.audio`` tensor from an engine output, or ``None``."""
|
||||
outputs = getattr(res, "outputs", None)
|
||||
if not outputs:
|
||||
return None
|
||||
mm = getattr(outputs[0], "multimodal_output", None)
|
||||
if mm is None:
|
||||
return None
|
||||
getter = mm.get if hasattr(mm, "get") else None
|
||||
codes = getter("codes") if getter else None
|
||||
if codes is None:
|
||||
return None
|
||||
if hasattr(codes, "get"):
|
||||
return codes.get("audio")
|
||||
return None
|
||||
|
||||
async def _generate_codec_chunks(
|
||||
self,
|
||||
generator,
|
||||
request_id: str,
|
||||
*,
|
||||
tts_params: dict[str, Any] | None = None,
|
||||
min_frames: int = 1,
|
||||
):
|
||||
"""Yield ``(codes[int64, n, codebooks], eos)`` deltas from a talker-only pipeline.
|
||||
|
||||
The AR stage publishes the *cumulative* ``codes.audio`` tensor on every
|
||||
step (``omni_client_multimodal_output_keys``); this generator aligns the
|
||||
rows to the output token ids (dropping the prefill placeholder and the
|
||||
EOS row) and emits only the new rows. The final item always carries
|
||||
``eos=True`` (possibly with zero frames).
|
||||
"""
|
||||
from vllm_omni.entrypoints.openai.codec_stream import align_codec_rows
|
||||
|
||||
adapter = self._get_tts_adapter()
|
||||
spec = getattr(adapter, "codec_stream_spec", None) or {}
|
||||
codebook_size = int(spec.get("codebook_size", 2048))
|
||||
usage_acc = SpeechOutputTokenCounter() if (adapter is not None and adapter.validates_generation) else None
|
||||
token_ids: list[int] = []
|
||||
sent = 0
|
||||
finished = False
|
||||
try:
|
||||
async for res in generator:
|
||||
if usage_acc is not None:
|
||||
usage_acc.observe(res)
|
||||
outputs = getattr(res, "outputs", None)
|
||||
if outputs:
|
||||
new_ids = getattr(outputs[0], "token_ids", None) or []
|
||||
# DELTA outputs carry per-step ids; CUMULATIVE ones carry the whole prefix.
|
||||
if len(new_ids) > len(token_ids) and list(new_ids[: len(token_ids)]) == token_ids:
|
||||
token_ids = list(new_ids)
|
||||
else:
|
||||
token_ids.extend(int(t) for t in new_ids)
|
||||
codes_cum = self._extract_codec_output(res)
|
||||
done = bool(getattr(res, "finished", False))
|
||||
if codes_cum is not None:
|
||||
aligned = align_codec_rows(codes_cum, token_ids, codebook_size)
|
||||
delta = aligned[sent:]
|
||||
if delta.shape[0] >= min_frames or (done and delta.shape[0] > 0):
|
||||
sent = int(aligned.shape[0])
|
||||
yield delta, done
|
||||
if done:
|
||||
finished = True
|
||||
if done and not finished:
|
||||
finished = True
|
||||
yield None, True
|
||||
finally:
|
||||
if usage_acc is not None and tts_params is not None and finished:
|
||||
self._validate_tts_generation(tts_params, usage_acc)
|
||||
|
||||
async def _iter_pcm_audio_bytes(self, request: OpenAICreateSpeechRequest):
|
||||
"""Yield raw PCM bytes for a speech request as soon as chunks are decoded."""
|
||||
request_id, generator, tts_params = await self._prepare_speech_generation(request)
|
||||
|
||||
@@ -53,6 +53,7 @@ Protocol:
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
from collections import deque
|
||||
from contextlib import aclosing
|
||||
|
||||
from fastapi import WebSocket, WebSocketDisconnect
|
||||
@@ -76,6 +77,54 @@ _MAX_CONFIG_MESSAGE_SIZE = 4 * 1024 * 1024 # allow large ref_audio payloads
|
||||
_MAX_INPUT_TEXT_MESSAGE_SIZE = 128 * 1024
|
||||
|
||||
|
||||
_DISCONNECT = object()
|
||||
|
||||
|
||||
class _SessionInbox:
|
||||
"""Single-reader message queue for one WebSocket session.
|
||||
|
||||
``pump`` is the only coroutine that calls ``websocket.receive_text()``;
|
||||
messages (or a disconnect marker) are queued and consumed by the session
|
||||
loop or, during a generation, by ``_iter_with_cancel``.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._queue: asyncio.Queue[object] = asyncio.Queue()
|
||||
self._replay: deque[str] = deque()
|
||||
|
||||
async def pump(self, websocket: WebSocket) -> None:
|
||||
try:
|
||||
while True:
|
||||
raw = await websocket.receive_text()
|
||||
await self._queue.put(raw)
|
||||
except (WebSocketDisconnect, RuntimeError):
|
||||
await self._queue.put(_DISCONNECT)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
|
||||
async def get(self, timeout: float | None) -> str:
|
||||
if self._replay:
|
||||
return self._replay.popleft()
|
||||
item = await asyncio.wait_for(self._queue.get(), timeout=timeout)
|
||||
if item is _DISCONNECT:
|
||||
raise WebSocketDisconnect(code=1006)
|
||||
return item # type: ignore[return-value]
|
||||
|
||||
async def peek_control(self) -> str:
|
||||
"""Wait for the next message; classify cancel/disconnect, replay the rest later."""
|
||||
item = await self._queue.get()
|
||||
if item is _DISCONNECT:
|
||||
return "disconnect"
|
||||
try:
|
||||
msg = json.loads(item) if isinstance(item, str) and item else {}
|
||||
except json.JSONDecodeError:
|
||||
msg = {}
|
||||
if isinstance(msg, dict) and msg.get("type") == "input.cancel":
|
||||
return "cancel"
|
||||
self._replay.append(item) # type: ignore[arg-type]
|
||||
return "other"
|
||||
|
||||
|
||||
class OmniStreamingSpeechHandler:
|
||||
"""Handles WebSocket sessions for streaming text-input TTS.
|
||||
|
||||
@@ -116,14 +165,17 @@ class OmniStreamingSpeechHandler:
|
||||
config: StreamingSpeechSessionConfig | None = None
|
||||
text_parts: list[str] = []
|
||||
utterance_index = 0
|
||||
# One reader task owns ``websocket.receive_text()`` for the whole
|
||||
# session; every consumer reads from the inbox instead. This is what
|
||||
# lets a generation loop notice ``input.cancel`` (or a disconnect)
|
||||
# while it streams, without cancelling a socket receive mid-flight.
|
||||
inbox = _SessionInbox()
|
||||
reader = asyncio.ensure_future(inbox.pump(websocket))
|
||||
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
raw = await asyncio.wait_for(
|
||||
websocket.receive_text(),
|
||||
timeout=self._config_timeout if config is None else self._idle_timeout,
|
||||
)
|
||||
raw = await inbox.get(timeout=self._config_timeout if config is None else self._idle_timeout)
|
||||
except asyncio.TimeoutError:
|
||||
if config is None:
|
||||
await self._send_error(websocket, "Timeout waiting for session.config")
|
||||
@@ -180,6 +232,7 @@ class OmniStreamingSpeechHandler:
|
||||
full_text,
|
||||
utterance_index=utterance_index,
|
||||
sentence_index=0,
|
||||
inbox=inbox,
|
||||
)
|
||||
total_sentences = 1
|
||||
|
||||
@@ -192,6 +245,12 @@ class OmniStreamingSpeechHandler:
|
||||
)
|
||||
utterance_index += 1
|
||||
|
||||
elif msg_type == "input.cancel":
|
||||
# Nothing is streaming at this point (cancels that arrive
|
||||
# during generation are handled in ``_iter_with_cancel``).
|
||||
text_parts.clear()
|
||||
await websocket.send_json({"type": "input.cancelled", "utterance_index": utterance_index})
|
||||
|
||||
elif msg_type == "session.close":
|
||||
await websocket.close()
|
||||
return
|
||||
@@ -210,6 +269,9 @@ class OmniStreamingSpeechHandler:
|
||||
await self._send_error(websocket, f"Internal error: {e}")
|
||||
except Exception:
|
||||
logger.debug("Failed to send error to streaming speech client", exc_info=True)
|
||||
finally:
|
||||
if not reader.done():
|
||||
reader.cancel()
|
||||
|
||||
async def _parse_message(self, websocket: WebSocket, raw: str) -> dict | None:
|
||||
"""Decode one client message, or report why it was rejected.
|
||||
@@ -257,6 +319,58 @@ class OmniStreamingSpeechHandler:
|
||||
|
||||
return config
|
||||
|
||||
async def _iter_with_cancel(
|
||||
self,
|
||||
websocket: WebSocket,
|
||||
stream,
|
||||
request_id: str,
|
||||
inbox: "_SessionInbox | None",
|
||||
):
|
||||
"""Yield items from ``stream`` while watching the session inbox.
|
||||
|
||||
``input.cancel`` aborts the engine request and ends the stream (the
|
||||
``audio.done`` / ``codec.done`` message then carries ``cancelled``);
|
||||
a disconnect raises ``WebSocketDisconnect`` like a plain send would;
|
||||
any other message is left queued for the session loop.
|
||||
"""
|
||||
self._last_cancelled = False
|
||||
if inbox is None:
|
||||
async for item in stream:
|
||||
yield item
|
||||
return
|
||||
next_task = asyncio.ensure_future(stream.__anext__())
|
||||
msg_task = asyncio.ensure_future(inbox.peek_control())
|
||||
try:
|
||||
while True:
|
||||
done, _ = await asyncio.wait({msg_task, next_task}, return_when=asyncio.FIRST_COMPLETED)
|
||||
if msg_task in done:
|
||||
kind = msg_task.result()
|
||||
if kind == "disconnect":
|
||||
next_task.cancel()
|
||||
raise WebSocketDisconnect(code=1006)
|
||||
if kind == "cancel":
|
||||
self._last_cancelled = True
|
||||
try:
|
||||
await self._speech_service.engine_client.abort(request_id)
|
||||
except Exception:
|
||||
logger.debug("Failed to abort request %s on input.cancel", request_id, exc_info=True)
|
||||
next_task.cancel()
|
||||
return
|
||||
msg_task = asyncio.ensure_future(inbox.peek_control())
|
||||
if next_task not in done:
|
||||
continue
|
||||
if next_task in done:
|
||||
try:
|
||||
item = next_task.result()
|
||||
except StopAsyncIteration:
|
||||
return
|
||||
yield item
|
||||
next_task = asyncio.ensure_future(stream.__anext__())
|
||||
finally:
|
||||
for task in (msg_task, next_task):
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
|
||||
async def _generate_and_send(
|
||||
self,
|
||||
websocket: WebSocket,
|
||||
@@ -265,6 +379,7 @@ class OmniStreamingSpeechHandler:
|
||||
*,
|
||||
utterance_index: int,
|
||||
sentence_index: int,
|
||||
inbox: "_SessionInbox | None" = None,
|
||||
) -> None:
|
||||
"""Generate audio for a single sentence and send it over WebSocket.
|
||||
|
||||
@@ -272,6 +387,18 @@ class OmniStreamingSpeechHandler:
|
||||
``sentence_index`` its position inside that flush.
|
||||
"""
|
||||
response_format = config.response_format or "wav"
|
||||
self._last_cancelled = False
|
||||
|
||||
if getattr(config, "output_mode", "pcm") == "codec_tokens":
|
||||
await self._generate_and_send_codec(
|
||||
websocket,
|
||||
config,
|
||||
sentence_text,
|
||||
utterance_index=utterance_index,
|
||||
sentence_index=sentence_index,
|
||||
inbox=inbox,
|
||||
)
|
||||
return
|
||||
|
||||
# Reject unmet word-timestamps preconditions early with a clear reason.
|
||||
if config.word_timestamps:
|
||||
@@ -308,6 +435,7 @@ class OmniStreamingSpeechHandler:
|
||||
ref_text=config.ref_text,
|
||||
x_vector_only_mode=config.x_vector_only_mode,
|
||||
speaker_embedding=config.speaker_embedding,
|
||||
seed=config.seed,
|
||||
stream=config.stream_audio,
|
||||
word_timestamps=config.word_timestamps,
|
||||
)
|
||||
@@ -352,7 +480,7 @@ class OmniStreamingSpeechHandler:
|
||||
tts_params=tts_params,
|
||||
)
|
||||
) as stream:
|
||||
async for chunk in stream:
|
||||
async for chunk in self._iter_with_cancel(websocket, stream, request_id, inbox):
|
||||
total_bytes += len(chunk)
|
||||
await websocket.send_bytes(chunk)
|
||||
else:
|
||||
@@ -388,11 +516,116 @@ class OmniStreamingSpeechHandler:
|
||||
"sentence_index": sentence_index,
|
||||
"total_bytes": total_bytes,
|
||||
"error": generation_failed,
|
||||
# Only present when a client input.cancel stopped the stream
|
||||
# (keeps the existing audio.done contract unchanged otherwise).
|
||||
**({"cancelled": True} if getattr(self, "_last_cancelled", False) else {}),
|
||||
}
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("Failed to send audio.done for sentence %d", sentence_index, exc_info=True)
|
||||
|
||||
async def _generate_and_send_codec(
|
||||
self,
|
||||
websocket: WebSocket,
|
||||
config: StreamingSpeechSessionConfig,
|
||||
sentence_text: str,
|
||||
*,
|
||||
utterance_index: int,
|
||||
sentence_index: int,
|
||||
inbox: "_SessionInbox | None" = None,
|
||||
) -> None:
|
||||
"""Stream the talker's codec frames (no server-side decoding).
|
||||
|
||||
Protocol: ``codec.start`` (JSON, codec description) → binary frames
|
||||
(``codec_stream.pack_codec_frame``) → ``codec.done``. Requires a
|
||||
talker-only deployment and an adapter with ``supports_codec_stream``.
|
||||
"""
|
||||
from vllm_omni.entrypoints.openai.codec_stream import codec_start_message, pack_codec_frame
|
||||
|
||||
adapter = self._speech_service._get_tts_adapter()
|
||||
if adapter is None or not getattr(adapter, "supports_codec_stream", False):
|
||||
await self._send_error(websocket, "output_mode='codec_tokens' is not supported by the loaded TTS model.")
|
||||
return
|
||||
spec = dict(adapter.codec_stream_spec)
|
||||
request = OpenAICreateSpeechRequest(
|
||||
input=sentence_text,
|
||||
model=config.model,
|
||||
voice=config.voice,
|
||||
task_type=config.task_type,
|
||||
language=config.language,
|
||||
instructions=config.instructions,
|
||||
response_format="pcm",
|
||||
speed=config.speed,
|
||||
max_new_tokens=config.max_new_tokens,
|
||||
non_streaming_mode=config.non_streaming_mode,
|
||||
ref_audio=config.ref_audio,
|
||||
ref_text=config.ref_text,
|
||||
x_vector_only_mode=config.x_vector_only_mode,
|
||||
speaker_embedding=config.speaker_embedding,
|
||||
seed=config.seed,
|
||||
stream=True,
|
||||
)
|
||||
await websocket.send_json(
|
||||
codec_start_message(
|
||||
utterance_index=utterance_index,
|
||||
sentence_index=sentence_index,
|
||||
sentence_text=sentence_text,
|
||||
spec=spec,
|
||||
model=config.model,
|
||||
)
|
||||
)
|
||||
total_frames = 0
|
||||
seq = 0
|
||||
generation_failed = False
|
||||
request_id = None
|
||||
eos_sent = False
|
||||
try:
|
||||
request_id, generator, tts_params = await self._speech_service._prepare_speech_generation(request)
|
||||
async with aclosing(
|
||||
self._speech_service._generate_codec_chunks(generator, request_id, tts_params=tts_params)
|
||||
) as stream:
|
||||
async for codes, eos in self._iter_with_cancel(websocket, stream, request_id, inbox):
|
||||
n = 0 if codes is None else int(codes.shape[0])
|
||||
total_frames += n
|
||||
await websocket.send_bytes(pack_codec_frame(seq, codes, eos=eos, codebooks=int(spec["codebooks"])))
|
||||
seq += 1
|
||||
eos_sent = eos_sent or eos
|
||||
except WebSocketDisconnect:
|
||||
if request_id is not None:
|
||||
try:
|
||||
await self._speech_service.engine_client.abort(request_id)
|
||||
except Exception:
|
||||
logger.debug("Failed to abort codec stream request %s", request_id, exc_info=True)
|
||||
raise
|
||||
except Exception as e:
|
||||
generation_failed = True
|
||||
logger.error("Codec stream failed for utterance %d: %s", utterance_index, e)
|
||||
await self._send_error(
|
||||
websocket, f"Codec stream failed for utterance {utterance_index}: {e}", partial_audio=total_frames > 0
|
||||
)
|
||||
finally:
|
||||
cancelled = bool(getattr(self, "_last_cancelled", False))
|
||||
if not eos_sent and not generation_failed:
|
||||
try:
|
||||
await websocket.send_bytes(pack_codec_frame(seq, None, eos=True, codebooks=int(spec["codebooks"])))
|
||||
seq += 1
|
||||
except Exception:
|
||||
logger.debug("Failed to send terminal codec frame", exc_info=True)
|
||||
try:
|
||||
await websocket.send_json(
|
||||
{
|
||||
"type": "codec.done",
|
||||
"utterance_index": utterance_index,
|
||||
"sentence_index": sentence_index,
|
||||
"total_frames": total_frames,
|
||||
"seq_last": seq - 1,
|
||||
"error": generation_failed,
|
||||
"cancelled": cancelled,
|
||||
}
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("Failed to send codec.done", exc_info=True)
|
||||
|
||||
async def _stream_audio_with_alignments(
|
||||
self,
|
||||
*,
|
||||
|
||||
@@ -169,6 +169,13 @@ class TTSModelAdapter(ABC):
|
||||
#: Target sample rates validated for this adapter's output path. An empty
|
||||
#: set means that the adapter does not expose per-request resampling.
|
||||
supported_output_sample_rates: ClassVar[frozenset[int]] = frozenset()
|
||||
#: Whether the model's talker can be served alone and its codec frames
|
||||
#: streamed to a client-side decoder (``output_mode="codec_tokens"``).
|
||||
supports_codec_stream: ClassVar[bool] = False
|
||||
#: Codec description sent in ``codec.start`` when ``supports_codec_stream``:
|
||||
#: ``codebooks``, ``codebook_size``, ``frame_rate_hz``, ``sample_rate``,
|
||||
#: ``decoder_id`` and ``eos_token_id``.
|
||||
codec_stream_spec: ClassVar[dict[str, Any]] = {}
|
||||
|
||||
max_new_tokens_min = 1
|
||||
|
||||
|
||||
@@ -44,6 +44,17 @@ class Qwen3TTSAdapter(ARTTSAdapter):
|
||||
stage_keys = frozenset({"qwen3_tts"})
|
||||
name = "qwen3_tts"
|
||||
supported_output_sample_rates = frozenset({8000, 24000})
|
||||
# Talker-only deployments (qwen3_tts_talker_only.yaml) stream 16-codebook
|
||||
# frames at 12.5 Hz for on-device decoding with Qwen3TTSTokenizerV2Decoder.
|
||||
supports_codec_stream = True
|
||||
codec_stream_spec = {
|
||||
"codebooks": 16,
|
||||
"codebook_size": 2048,
|
||||
"frame_rate_hz": 12.5,
|
||||
"sample_rate": 24000,
|
||||
"decoder_id": "qwen3_tts_tokenizer_v2_12hz",
|
||||
"eos_token_id": 2150,
|
||||
}
|
||||
|
||||
def _get_model_variant(self) -> str | None:
|
||||
"""Return the task supported by the loaded Qwen3-TTS checkpoint.
|
||||
|
||||
@@ -56,3 +56,30 @@ QWEN3_TTS_PIPELINE = PipelineConfig(
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
# Talker-only pipeline: stage 0 is terminal and its codec frames (codes.audio)
|
||||
# are published to the client for on-device decoding (codec-token streaming,
|
||||
# see entrypoints/openai/codec_stream.py). No next-stage processors, so no
|
||||
# omni connector is created; ``final_output_type="latent"`` keeps the AR
|
||||
# output path (audio-specific terminal handling is not applied).
|
||||
QWEN3_TTS_TALKER_ONLY_PIPELINE = PipelineConfig(
|
||||
model_type="qwen3_tts_talker_only",
|
||||
default_deploy_config_name="qwen3_tts_talker_only.yaml",
|
||||
model_arch="Qwen3TTSTalkerForConditionalGeneration",
|
||||
stages=(
|
||||
StagePipelineConfig(
|
||||
stage_id=0,
|
||||
model_stage="qwen3_tts",
|
||||
execution_type=StageExecutionType.LLM_AR,
|
||||
input_sources=(),
|
||||
owns_tokenizer=True,
|
||||
final_output=True,
|
||||
final_output_type="latent",
|
||||
engine_output_type="latent",
|
||||
sampling_constraints={
|
||||
"detokenize": False,
|
||||
"stop_token_ids": [2150],
|
||||
},
|
||||
),
|
||||
),
|
||||
)
|
||||
Reference in new issue
Block a user