diff --git a/benchmarks/tts/codec_stream_e2e.py b/benchmarks/tts/codec_stream_e2e.py new file mode 100644 index 000000000..48ddd95f5 --- /dev/null +++ b/benchmarks/tts/codec_stream_e2e.py @@ -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//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()) diff --git a/docs/serving/codec_token_streaming.md b/docs/serving/codec_token_streaming.md new file mode 100644 index 000000000..974aed25d --- /dev/null +++ b/docs/serving/codec_token_streaming.md @@ -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 ` 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()) diff --git a/examples/online_serving/text_to_speech/qwen3_tts/edge_decoder_client.py b/examples/online_serving/text_to_speech/qwen3_tts/edge_decoder_client.py new file mode 100644 index 000000000..0d7043e69 --- /dev/null +++ b/examples/online_serving/text_to_speech/qwen3_tts/edge_decoder_client.py @@ -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()) diff --git a/tests/config/test_talker_only_pipeline.py b/tests/config/test_talker_only_pipeline.py new file mode 100644 index 000000000..64c7406c2 --- /dev/null +++ b/tests/config/test_talker_only_pipeline.py @@ -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 diff --git a/tests/e2e/online_serving/test_qwen3_tts_codec_stream.py b/tests/e2e/online_serving/test_qwen3_tts_codec_stream.py new file mode 100644 index 000000000..bd6d514fb --- /dev/null +++ b/tests/e2e/online_serving/test_qwen3_tts_codec_stream.py @@ -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 diff --git a/tests/entrypoints/openai_api/test_codec_stream_protocol.py b/tests/entrypoints/openai_api/test_codec_stream_protocol.py new file mode 100644 index 000000000..bc914a474 --- /dev/null +++ b/tests/entrypoints/openai_api/test_codec_stream_protocol.py @@ -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() diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 5ac47dddc..d9751ae5f 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -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, diff --git a/vllm_omni/deploy/qwen3_tts_talker_only.yaml b/vllm_omni/deploy/qwen3_tts_talker_only.yaml new file mode 100644 index 000000000..25857312b --- /dev/null +++ b/vllm_omni/deploy/qwen3_tts_talker_only.yaml @@ -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 diff --git a/vllm_omni/entrypoints/openai/codec_stream.py b/vllm_omni/entrypoints/openai/codec_stream.py new file mode 100644 index 000000000..7ccb97221 --- /dev/null +++ b/vllm_omni/entrypoints/openai/codec_stream.py @@ -0,0 +1,129 @@ +"""Wire format for codec-token streaming (server → on-device decoder). + +Binary frame layout (little-endian): + + header `` 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(" 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=" 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, + } diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index f204ae209..1d199adc3 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -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( diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 813d916fe..83a585e4f 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -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) diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index 0f883f179..2d3d011f2 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -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, *, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index d4316a2fe..7c82a297c 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py @@ -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 diff --git a/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py index 48d019eeb..e558dffbe 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py @@ -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. diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index 6dfe8736d..3b25dabb6 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py @@ -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], + }, + ), + ), +)