| """Evaluate exported VibeVoice ONNX sub-parts against their PyTorch references.
|
|
|
| Two stages:
|
| A) COMPONENT parity — for every exported component (acoustic/semantic encoders & decoder,
|
| diffusion_head, connectors, projector) load the PyTorch reference (user_script loaders) and
|
| the ONNX session on the SAME inputs, then report cosine + max|Δ| with a PASS/FAIL threshold.
|
| The LLM decoder is checked structurally (int4 MatMulNBits + GQA + inputs_embeds→hidden).
|
| B) WHOLE pipeline — an end-to-end integrity chain of the ONNX pieces:
|
| • TTS codec round-trip (audio → acoustic_encoder → acoustic_decoder → audio′): waveform
|
| correlation + SNR — the standard neural-codec reconstruction metric.
|
| • Full ONNX chain smoke: latent → acoustic_connector → (cond) → diffusion_head → 1 denoise
|
| step, and acoustic_decoder → waveform — every exported block run in sequence, shapes +
|
| finiteness verified. (The autoregressive LLM + multi-step DDPM loop need the genai
|
| runtime and are out of scope here — reported as SKIPPED, not silently passed.)
|
|
|
| Model is the FINAL positional (a key {1.5b,asr,asr-hf,realtime} or a checkpoint dir path).
|
|
|
| Usage:
|
| uv run eval.py 1.5b
|
| uv run eval.py --device cuda --precision fp16 ./asr-hf
|
| uv run eval.py --audio sample.wav 1.5b
|
| uv run eval.py --components acoustic_encoder diffusion_head realtime
|
| """
|
| import argparse
|
| import sys
|
| import os
|
| from pathlib import Path
|
|
|
| import numpy as np
|
|
|
| HERE = Path(__file__).parent
|
| sys.path.insert(0, str(HERE))
|
| for _s in (sys.stdout, sys.stderr):
|
| try: _s.reconfigure(encoding="utf-8", errors="replace")
|
| except Exception: pass
|
|
|
| import optimize as O
|
| import user_script as us
|
|
|
| COS_TOL = 0.99
|
|
|
|
|
| def _np(o):
|
| """Best-effort tensor → float32 numpy from a torch tensor or an output object."""
|
| import torch
|
| if torch.is_tensor(o):
|
| return o.detach().float().cpu().numpy()
|
| for a in ("latents", "mean", "sample", "audio", "last_hidden_state"):
|
| if hasattr(o, a):
|
| v = getattr(o, a); v = v() if callable(v) else v
|
| if torch.is_tensor(v):
|
| return v.detach().float().cpu().numpy()
|
| if isinstance(o, (tuple, list)):
|
| return _np(o[0])
|
| raise TypeError(f"cannot convert output of type {type(o).__name__}")
|
|
|
|
|
| def _cos(a, b):
|
| a, b = a.ravel().astype(np.float64), b.ravel().astype(np.float64)
|
| n = np.linalg.norm(a) * np.linalg.norm(b)
|
| return float(a @ b / (n + 1e-9))
|
|
|
|
|
| def _session(path):
|
| import onnxruntime as ort
|
| return ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
|
|
|
|
|
| def parity_component(model_key, src, comp, onnx_path):
|
| """Run PyTorch ref + ONNX on identical dummy inputs; return (cos, maxd, ok)."""
|
| import torch
|
| loader_name, io_name, dummy_name, env = O.MODELS[model_key]["olive"][comp]
|
| os.environ.update(env)
|
| pyt = getattr(us, loader_name)(str(src))
|
| if hasattr(pyt, "eval"):
|
| pyt = pyt.eval()
|
| io = getattr(us, io_name)()
|
| names = io["input_names"]
|
| inputs = getattr(us, dummy_name)()
|
| with torch.no_grad():
|
| ref = _np(pyt(*[inputs[n] for n in names]))
|
|
|
|
|
| sess = _session(onnx_path)
|
| _ORT2NP = {"tensor(float)": np.float32, "tensor(float16)": np.float16,
|
| "tensor(int64)": np.int64, "tensor(int32)": np.int32, "tensor(bool)": np.bool_}
|
| want = {i.name: _ORT2NP.get(i.type) for i in sess.get_inputs()}
|
| feed = {n: (inputs[n].numpy().astype(want[n]) if want.get(n) else inputs[n].numpy()) for n in names}
|
| got = sess.run(None, feed)[0]
|
| m = min(ref.size, got.size)
|
| a, b = got.ravel()[:m].astype(np.float64), ref.ravel()[:m].astype(np.float64)
|
| cos = _cos(a, b)
|
| maxd = float(np.abs(a - b).max())
|
| rel = maxd / (float(np.abs(b).max()) + 1e-12)
|
|
|
|
|
| ok = (cos >= COS_TOL) or (maxd < 1e-4) or (rel < 1e-3)
|
| return cos, maxd, ok
|
|
|
|
|
| def check_llm(onnx_path):
|
| """Structural check of the genai-built decoder (no full runtime inference)."""
|
| import onnx
|
| g = onnx.load(str(onnx_path), load_external_data=False).graph
|
| ops = [n.op_type for n in g.node]
|
| innames = [i.name for i in g.input]
|
| nb, gqa = ops.count("MatMulNBits"), ops.count("GroupQueryAttention")
|
| embeds_in = any("inputs_embeds" in n for n in innames)
|
| ok = (nb > 0 or "MatMul" in ops) and embeds_in
|
| return f"MatMulNBits={nb} GQA={gqa} inputs_embeds={embeds_in}", ok
|
|
|
|
|
|
|
|
|
| def _load_audio(path, samples, sr=24000):
|
| import soundfile as sf, librosa
|
| wav, in_sr = sf.read(path, dtype="float32", always_2d=False)
|
| if wav.ndim > 1:
|
| wav = wav.mean(1)
|
| if in_sr != sr:
|
| wav = librosa.resample(wav, orig_sr=in_sr, target_sr=sr)
|
| if len(wav) < samples:
|
| wav = np.pad(wav, (0, samples - len(wav)))
|
| return wav[:samples].astype(np.float32)[None, None, :]
|
|
|
|
|
| def _synth_audio(samples, sr=24000):
|
| t = np.arange(samples) / sr
|
| wav = 0.3 * (np.sin(2 * np.pi * 220 * t) + 0.5 * np.sin(2 * np.pi * 440 * t))
|
| return wav.astype(np.float32)[None, None, :]
|
|
|
|
|
| _ORT2NP_FEED = {"tensor(float)": np.float32, "tensor(float16)": np.float16,
|
| "tensor(int64)": np.int64, "tensor(int32)": np.int32, "tensor(bool)": np.bool_}
|
|
|
|
|
| def _feed(sess, d):
|
| """Cast each fed array to the session's DECLARED input dtype (fp16 graphs want float16;
|
| float32-pinned inputs like diffusion timesteps stay float32 — matches the graph either way)."""
|
| want = {i.name: _ORT2NP_FEED.get(i.type) for i in sess.get_inputs()}
|
| return {k: (v.astype(want[k]) if want.get(k) is not None and hasattr(v, "astype") else v)
|
| for k, v in d.items()}
|
|
|
|
|
| def whole_pipeline(model_key, models_dir, audio_path):
|
| """End-to-end ONNX integrity: codec round-trip + full-chain smoke. Returns list of rows."""
|
| rows = []
|
| present = {c: models_dir / f"{c}.onnx" for c in
|
| ("acoustic_encoder", "acoustic_decoder", "acoustic_connector", "diffusion_head",
|
| "semantic_encoder", "multi_modal_projector")
|
| if (models_dir / f"{c}.onnx").exists()}
|
|
|
|
|
| if "acoustic_encoder" in present and "acoustic_decoder" in present:
|
| enc, dec = _session(present["acoustic_encoder"]), _session(present["acoustic_decoder"])
|
| samples = enc.get_inputs()[0].shape[2]
|
| samples = 24000 if not isinstance(samples, int) else samples
|
| wav = _load_audio(audio_path, samples) if audio_path else _synth_audio(samples)
|
| lat = enc.run(None, _feed(enc, {"audio": wav}))[0]
|
| dec_frames = dec.get_inputs()[0].shape[1]
|
| if isinstance(dec_frames, int) and lat.shape[1] != dec_frames:
|
| f = min(lat.shape[1], dec_frames)
|
| lat = lat[:, :f, :] if lat.shape[1] > f else np.pad(lat, ((0, 0), (0, dec_frames - lat.shape[1]), (0, 0)))
|
| recon = dec.run(None, _feed(dec, {"latents": lat.astype(np.float32)}))[0].ravel()
|
| n = min(len(recon), wav.size)
|
| x = wav.ravel()[:n]; y = recon[:n]
|
| corr = float(np.corrcoef(x, y)[0, 1]) if n > 1 else float("nan")
|
| noise = x - y
|
| snr = 10 * np.log10((np.sum(x**2) + 1e-12) / (np.sum(noise**2) + 1e-12))
|
| finite = bool(np.isfinite(recon).all())
|
| rows.append(("codec round-trip", f"corr={corr:+.3f} SNR={snr:+.1f}dB finite={finite} "
|
| f"in={wav.size} out={len(recon)}", finite))
|
| else:
|
| rows.append(("codec round-trip", "SKIPPED (needs acoustic_encoder+decoder)", None))
|
|
|
|
|
| if "acoustic_connector" in present and "diffusion_head" in present:
|
| conn = _session(present["acoustic_connector"]); dh = _session(present["diffusion_head"])
|
| rng = np.random.default_rng(0)
|
| T = 8
|
| lat = rng.standard_normal((1, T, 64)).astype(np.float32)
|
| cond = conn.run(None, _feed(conn, {"features": lat}))[0]
|
| H = cond.shape[-1]
|
| ni = rng.standard_normal((T, 64)).astype(np.float32)
|
| ts = (rng.random(T) * 1000).astype(np.float32)
|
| pred = dh.run(None, _feed(dh, {"noisy_images": ni, "timesteps": ts,
|
| "condition": cond.reshape(T, H).astype(np.float32)}))[0]
|
| ok = pred.shape == (T, 64) and np.isfinite(pred).all() and np.isfinite(cond).all()
|
| rows.append(("tts chain smoke", f"connector→[H={H}]→diffusion_head pred={pred.shape} finite={ok}", ok))
|
| else:
|
| rows.append(("tts chain smoke", "SKIPPED (needs acoustic_connector+diffusion_head)", None))
|
|
|
|
|
| if "multi_modal_projector" in present and "acoustic_encoder" in present and "semantic_encoder" in present:
|
| ae, se, proj = (_session(present[k]) for k in ("acoustic_encoder", "semantic_encoder", "multi_modal_projector"))
|
| wav = _synth_audio(24000)
|
| al = ae.run(None, {"audio": wav})[0]; sl = se.run(None, {"audio": wav})[0]
|
| f = min(al.shape[1], sl.shape[1])
|
| fused = proj.run(None, {"acoustic_latents": al[:, :f].astype(np.float32),
|
| "semantic_latents": sl[:, :f].astype(np.float32)})[0]
|
| ok = np.isfinite(fused).all()
|
| rows.append(("asr fusion chain", f"ac+sem→projector fused={fused.shape} finite={ok}", ok))
|
|
|
| rows.append(("autoregressive LLM + DDPM loop", "SKIPPED (needs genai runtime; out of scope)", None))
|
| return rows
|
|
|
|
|
| def main():
|
| ap = argparse.ArgumentParser(description="Evaluate VibeVoice ONNX sub-parts",
|
| formatter_class=argparse.RawDescriptionHelpFormatter)
|
| global COS_TOL
|
| ap.add_argument("--device", choices=["cpu", "cuda", "gpu"], default="cpu")
|
| ap.add_argument("--precision", choices=["int4", "fp16", "fp32"], default="int4")
|
| ap.add_argument("--components", help="comma-separated subset (default: all present)")
|
| ap.add_argument("--output-dir", help="models dir (default onnx/{model}/{device}_{precision})")
|
| ap.add_argument("--audio", help="wav for codec round-trip (default: synthetic tone)")
|
| ap.add_argument("--tol", type=float, default=COS_TOL, help=f"cosine PASS threshold (default {COS_TOL})")
|
| ap.add_argument("model", help="keyword {1.5b,asr,asr-hf,realtime} or checkpoint dir path")
|
| args = ap.parse_args()
|
| COS_TOL = args.tol
|
| device = "cuda" if args.device == "gpu" else args.device
|
| model_key, src = O.resolve_target(args.model)
|
| models_dir = Path(args.output_dir) if args.output_dir else (HERE / "onnx" / model_key / f"{device}_{args.precision}")
|
| if not models_dir.exists():
|
| sys.exit(f"no exported models at {models_dir} — run optimize.py first")
|
|
|
| print(f"=== eval {model_key} src={src} models={models_dir} tol={COS_TOL} ===\n")
|
|
|
|
|
| comps = [c.strip() for c in args.components.split(",")] if args.components else O.all_components(model_key)
|
| print("A) COMPONENT PARITY (ONNX vs PyTorch)")
|
| results = []
|
| for comp in comps:
|
| onnx_path = models_dir / (f"{comp}.onnx" if comp != "llm" else "llm_decoder.onnx")
|
| if not onnx_path.exists():
|
| print(f" {comp:22s} — not exported (skip)"); continue
|
| try:
|
| if comp == "llm":
|
| info, ok = check_llm(onnx_path)
|
| print(f" {comp:22s} {'PASS' if ok else 'FAIL'} {info}"); results.append(ok)
|
| else:
|
| cos, maxd, ok = parity_component(model_key, src, comp, onnx_path)
|
| print(f" {comp:22s} {'PASS' if ok else 'FAIL'} cos={cos:.6f} maxd={maxd:.2e}")
|
| results.append(ok)
|
| except Exception as e:
|
| print(f" {comp:22s} ERROR {type(e).__name__}: {str(e)[:120]}"); results.append(False)
|
|
|
|
|
| print("\nB) WHOLE PIPELINE (end-to-end ONNX integrity)")
|
| for name, info, ok in whole_pipeline(model_key, models_dir, args.audio):
|
| tag = "SKIP" if ok is None else ("PASS" if ok else "FAIL")
|
| print(f" {name:32s} {tag} {info}")
|
| if ok is not None:
|
| results.append(ok)
|
|
|
| npass = sum(1 for r in results if r); ntot = len(results)
|
| print(f"\n=== {npass}/{ntot} checks passed ===")
|
| sys.exit(0 if npass == ntot else 1)
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
