78 lines
2.8 KiB
Python
78 lines
2.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Offline parity diagnosis for the codec-token stream (CPU only).
|
|
|
|
Given a codec_stream_e2e.py output directory (client_<i>.codes.npy, client_<i>.wav,
|
|
full_<i>.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/<ts>
|
|
"""
|
|
|
|
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())
|