diff --git a/benchmarks/tts/codec_parity_diag.py b/benchmarks/tts/codec_parity_diag.py new file mode 100644 index 000000000..2cd3f8ea0 --- /dev/null +++ b/benchmarks/tts/codec_parity_diag.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +"""Offline parity diagnosis for the codec-token stream (CPU only). + +Given a codec_stream_e2e.py output directory (client_.codes.npy, client_.wav, +full_.wav) decode the *received* codes with the exact (non-chunked) decoder and +compare (a) against the client's chunked decode (tests the client's chunking) and +(b) against the full server's PCM (tests token equality + the server's decode). +Reports lag-aligned SNR and correlation. + + python benchmarks/tts/codec_parity_diag.py --dir benchmarks/tts/edge_results/codec_stream/ +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import soundfile as sf +import torch + + +def aligned_snr(ref: np.ndarray, test: np.ndarray, max_lag: int = 4000) -> tuple[float, float, int]: + n = min(ref.size, test.size) + ref, test = ref[:n].astype(np.float64), test[:n].astype(np.float64) + if n < 3 * max_lag: + lag = 0 + else: + c = np.correlate(ref[max_lag:-max_lag], test, mode="valid") + lag = int(np.argmax(c)) - max_lag + if lag >= 0: + a, b = test[lag:], ref[: n - lag] + else: + a, b = test[: n + lag], ref[-lag:] + m = min(a.size, b.size) + a, b = a[:m], b[:m] + snr = 10 * np.log10((b**2).sum() / (((b - a) ** 2).sum() + 1e-12)) + corr = float(np.corrcoef(a, b)[0, 1]) if m > 1 else 0.0 + return float(snr), corr, lag + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--dir", required=True) + p.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice") + args = p.parse_args(argv) + from vllm_omni.edge.decoder_export import load_code2wav_decoder + + dec = load_code2wav_decoder(args.model) + d = Path(args.dir) + report = [] + for codes_path in sorted(d.glob("client_*.codes.npy")): + i = codes_path.name.split("_")[1].split(".")[0] + codes = np.load(codes_path) + if codes.size == 0: + continue + with torch.inference_mode(): + exact = dec._forward_exact(torch.as_tensor(codes.T[None], dtype=torch.long)).reshape(-1).numpy() + client, _ = sf.read(d / f"client_{i}.wav", dtype="float32") + full, _ = sf.read(d / f"full_{i}.wav", dtype="float32") + s1, c1, l1 = aligned_snr(exact, client.reshape(-1)) + s2, c2, l2 = aligned_snr(full.reshape(-1), exact) + row = { + "i": int(i), + "frames": int(codes.shape[0]), + "exact_vs_client_chunked": {"snr_db": s1, "corr": c1, "lag": l1}, + "full_server_vs_exact": {"snr_db": s2, "corr": c2, "lag": l2}, + } + report.append(row) + print(json.dumps(row)) + (d / "parity_diag.json").write_text(json.dumps(report, indent=1)) + 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 index 8ca938bf0..ec7db0d77 100644 --- 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 @@ -199,6 +199,7 @@ async def run(args: argparse.Namespace) -> dict[str, Any]: if val is not None: config[key] = val pcm_parts: list[np.ndarray] = [] + codes_all: list[np.ndarray] = [] chunk_log: list[dict[str, Any]] = [] pending = np.zeros((0, decoder.num_quantizers), dtype=np.int64) chunk_idx = 0 @@ -222,6 +223,7 @@ async def run(args: argparse.Namespace) -> dict[str, Any]: total_frames += n if n: pending = np.concatenate([pending, arr.astype(np.int64)], axis=0) + codes_all.append(arr.astype(np.int64)) 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): @@ -283,6 +285,10 @@ async def run(args: argparse.Namespace) -> dict[str, Any]: sf.write(args.out, audio, sr, format="WAV") result["wav"] = args.out + # Raw received codec rows (frames x codebooks) for offline parity checks. + codes_path = str(Path(args.out).with_suffix(".codes.npy")) + np.save(codes_path, np.concatenate(codes_all, axis=0) if codes_all else np.zeros((0, 0), dtype=np.int64)) + result["codes"] = codes_path 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)) diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 83a585e4f..fc80d3b11 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -3293,8 +3293,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): token_ids: list[int] = [] sent = 0 finished = False + n_outputs = 0 + n_with_codes = 0 + first_codes_at = -1 try: async for res in generator: + n_outputs += 1 if usage_acc is not None: usage_acc.observe(res) outputs = getattr(res, "outputs", None) @@ -3308,6 +3312,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): codes_cum = self._extract_codec_output(res) done = bool(getattr(res, "finished", False)) if codes_cum is not None: + n_with_codes += 1 + if first_codes_at < 0: + first_codes_at = n_outputs 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): @@ -3319,6 +3326,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): finished = True yield None, True finally: + logger.info( + "[codec_stream] req=%s engine outputs=%d with codes=%d first codes at output #%d rows sent=%d", + request_id, + n_outputs, + n_with_codes, + first_codes_at, + sent, + ) if usage_acc is not None and tts_params is not None and finished: self._validate_tts_generation(tts_params, usage_acc)