#!/usr/bin/env python3 from __future__ import annotations import argparse import json import random import sys import time import wave from pathlib import Path from typing import Any import numpy as np import soundfile as sf REPO_ROOT = Path(__file__).resolve().parents[2] DEFAULT_REFERENCE_ROOT = REPO_ROOT / "reference" / "neutts" DEFAULT_CASES = REPO_ROOT / "tests" / "neutts" / "neutts_2e_warm_bench_cases.json" DEFAULT_BACKBONE = REPO_ROOT / "models" / "NeuTTS-2E" DEFAULT_OUT_ROOT = REPO_ROOT / "build" / "logs" / "neutts" / "python_baseline" FAMILY = "neutts" SAMPLE_RATE = 24_000 def load_cases(path: Path) -> list[dict[str, Any]]: root = json.loads(path.read_text(encoding="utf-8")) cases = root.get("cases") if not isinstance(cases, list) or not cases: raise RuntimeError(f"case file has no cases: {path}") return cases def select_cases(cases: list[dict[str, Any]], only: set[str]) -> list[dict[str, Any]]: selected = [] for case in cases: if case.get("family") != FAMILY: continue if only and str(case.get("id")) not in only: continue selected.append(case) if only: found = {str(case.get("id")) for case in selected} missing = sorted(only - found) if missing: raise RuntimeError(f"unknown {FAMILY} case id(s): {', '.join(missing)}") if not selected: raise RuntimeError(f"no {FAMILY} cases selected") return selected def sync_cuda(torch_mod: Any, device: str) -> None: if device.startswith("cuda"): torch_mod.cuda.synchronize(torch_mod.device(device)) def reset_peak(torch_mod: Any, device: str) -> None: if device.startswith("cuda"): torch_mod.cuda.reset_peak_memory_stats(torch_mod.device(device)) def memory_snapshot(torch_mod: Any, device: str) -> dict[str, float]: out: dict[str, float] = {} try: import psutil out["rss_mb"] = psutil.Process().memory_info().rss / (1024.0 * 1024.0) except Exception: pass if device.startswith("cuda"): dev = torch_mod.device(device) out["cuda_alloc_mb"] = torch_mod.cuda.memory_allocated(dev) / (1024.0 * 1024.0) out["cuda_peak_mb"] = torch_mod.cuda.max_memory_allocated(dev) / (1024.0 * 1024.0) out["cuda_reserved_mb"] = torch_mod.cuda.memory_reserved(dev) / (1024.0 * 1024.0) return out def configure_torch_runtime(torch_mod: Any, device: str) -> None: if device.startswith("cuda"): torch_mod.backends.cuda.matmul.allow_tf32 = False torch_mod.backends.cudnn.allow_tf32 = False torch_mod.backends.cudnn.benchmark = False torch_mod.use_deterministic_algorithms(True, warn_only=True) def seed_runtime(torch_mod: Any, device: str, seed: int) -> None: random.seed(seed) np.random.seed(seed & 0xFFFF_FFFF) torch_mod.manual_seed(seed) if device.startswith("cuda"): torch_mod.cuda.manual_seed_all(seed) def wav_duration_seconds(path: Path) -> float: with wave.open(str(path), "rb") as wav: return wav.getnframes() / float(wav.getframerate()) def audio_summary(audio: Any, sample_rate: int) -> dict[str, Any]: values = np.asarray(audio, dtype=np.float32) if values.ndim == 2 and values.shape[0] == 1: values = values[0] if values.ndim not in (1, 2): raise RuntimeError(f"NeuTTS warmbench expected 1D or 2D audio, got shape {values.shape}") flat = values.reshape(-1) if flat.size == 0: raise RuntimeError("NeuTTS warmbench received empty audio") frames = int(values.shape[0]) channels = 1 if values.ndim == 1 else int(values.shape[1]) return { "sample_rate": int(sample_rate), "channels": channels, "samples": int(flat.size), "frames": frames, "duration_sec": float(frames / sample_rate), "sum": float(np.sum(flat, dtype=np.float64)), "mean_abs": float(np.mean(np.abs(flat), dtype=np.float64)), "rms": float(np.sqrt(np.mean(np.square(flat, dtype=np.float64)))), "min": float(np.min(flat)), "max": float(np.max(flat)), } def case_text(request: dict[str, Any], key: str, fallback: str) -> str: value = request.get(key, fallback) if not isinstance(value, str) or not value: raise RuntimeError(f"request requires non-empty string field '{key}'") return value def case_int(request: dict[str, Any], key: str, fallback: int) -> int: return int(request[key]) if key in request and request[key] is not None else fallback def case_float(request: dict[str, Any], key: str, fallback: float) -> float: return float(request[key]) if key in request and request[key] is not None else fallback def run_request(tts: Any, torch_mod: Any, device: str, request: dict[str, Any], output_path: Path) -> dict[str, Any]: request_id = case_text(request, "id", "request") speaker = case_text(request, "speaker", "emily") emotion = case_text(request, "emotion", "neutral") text = case_text(request, "text", "") seed = case_int(request, "seed", 1234) temperature = case_float(request, "temperature", 1.0) top_k = case_int(request, "top_k", 50) # NeuTTS keeps the seed on the session object. Update it per request so the # model remains loaded while each request stays reproducible. tts._seed = seed seed_runtime(torch_mod, device, seed) reset_peak(torch_mod, device) sync_cuda(torch_mod, device) start = time.perf_counter() wav = tts.infer(text, speaker=speaker, emotion=emotion, temperature=temperature, top_k=top_k) sync_cuda(torch_mod, device) inference_ms = (time.perf_counter() - start) * 1000.0 output_path.parent.mkdir(parents=True, exist_ok=True) sf.write(output_path, wav, SAMPLE_RATE) audio_seconds = wav_duration_seconds(output_path) rtf = (inference_ms / 1000.0) / audio_seconds if audio_seconds > 0 else 0.0 result = memory_snapshot(torch_mod, device) result.update( { "request_id": request_id, "speaker": speaker, "emotion": emotion, "seed": seed, "temperature": temperature, "top_k": top_k, "inference_ms": inference_ms, "audio_seconds": audio_seconds, "rtf": rtf, "x_realtime": (1.0 / rtf) if rtf > 0.0 else 0.0, "audio": str(output_path), "summary": audio_summary(wav, SAMPLE_RATE), } ) return result def main() -> int: parser = argparse.ArgumentParser(description="NeuTTS-2E Python reference warmbench.") parser.add_argument("--cases", type=Path, default=DEFAULT_CASES) parser.add_argument("--reference-root", type=Path, default=DEFAULT_REFERENCE_ROOT) parser.add_argument("--backbone", default=str(DEFAULT_BACKBONE)) parser.add_argument("--codec", default="neuphonic/neucodec") parser.add_argument("--backend", default="cuda") parser.add_argument("--device", default="cuda") parser.add_argument("--codec-device", default="cuda") parser.add_argument("--only", action="append", default=[], help="case id(s), comma-separated allowed") parser.add_argument("--out-root", type=Path, default=DEFAULT_OUT_ROOT) parser.add_argument("--enable-watermark", action="store_true") args = parser.parse_args() reference_root = args.reference_root.resolve() if not (reference_root / "neutts").is_dir(): raise RuntimeError(f"missing NeuTTS reference checkout: {reference_root}") sys.path.insert(0, str(reference_root)) import torch from neutts import NeuTTS2E if args.device.startswith("cuda") and not torch.cuda.is_available(): raise RuntimeError("NeuTTS warmbench requested CUDA, but CUDA is not available") configure_torch_runtime(torch, args.device) only = {item.strip() for raw in args.only for item in raw.split(",") if item.strip()} cases = select_cases(load_cases(args.cases), only) out_root = args.out_root.resolve() out_root.mkdir(parents=True, exist_ok=True) tts = NeuTTS2E( backbone_repo=args.backbone, backbone_device=args.device, codec_repo=args.codec, codec_device=args.codec_device, seed=None, ) if not args.enable_watermark: tts.watermarker = None sync_cuda(torch, args.device) for case in cases: case_dir = out_root / str(case["id"]) outputs_dir = case_dir / "outputs" outputs_dir.mkdir(parents=True, exist_ok=True) stdout_lines = [ f"family={FAMILY}", f"task={case.get('task', 'tts')}", f"mode={case.get('mode', 'offline')}", "torch_tf32=0", "deterministic_algorithms=warn_only", ] memory: dict[str, Any] = { "case": case["id"], "device": args.device, "codec_device": args.codec_device, "requests": [], } steps: list[dict[str, Any]] = [] warmup = case.get("warmup") if warmup is not None: warm_path = outputs_dir / f"{case_text(warmup, 'id', 'warmup')}.wav" warm = run_request(tts, torch, args.device, warmup, warm_path) stdout_lines.append(f"warmup_id={warm['request_id']}") stdout_lines.append(f"warmup_inference_ms={warm['inference_ms']:.6f}") memory["warmup"] = warm total_inference_ms = 0.0 for index, request in enumerate(case["requests"]): request_id = case_text(request, "id", f"request_{index}") result = run_request(tts, torch, args.device, request, outputs_dir / f"{request_id}.wav") total_inference_ms += float(result["inference_ms"]) memory["requests"].append(result) stdout_lines.extend( [ f"request_index={index}", f"request_id={request_id}", f"text[{index}]={case_text(request, 'text', '')}", f"[TIMING] request.{request_id}.inference_ms {result['inference_ms']:.6f}", f"audio_seconds={result['audio_seconds']:.6f}", f"rtf={result['rtf']:.6f}", f"x_realtime={result['x_realtime']:.6f}", f"audio_out[{index}]={result['audio']}", ] ) step = { "request_index": index, "id": request_id, "speaker": result["speaker"], "emotion": result["emotion"], "seed": result["seed"], "temperature": result["temperature"], "top_k": result["top_k"], "text_length": len(case_text(request, "text", "")), "stems": [ { "name": "audio", "audio": result["audio"], "summary": result["summary"], } ], "metrics": { "wall_ms": result["inference_ms"], "inference_ms": result["inference_ms"], "rtf": result["rtf"], "x_realtime": result["x_realtime"], }, } steps.append(step) stdout_lines.append(f"summary_json[{index}]=" + json.dumps(result["summary"], ensure_ascii=False, separators=(",", ":"))) print( f"[REF] {case['id']}/{request_id}: {result['inference_ms']:.2f} ms, " f"{result['audio_seconds']:.2f}s audio, RTF={result['rtf']:.3f}, " f"{result['x_realtime']:.2f}x", flush=True, ) stdout_lines.append(f"[TIMING] session.inference_ms {total_inference_ms:.6f}") summary = { "family": FAMILY, "task": case.get("task", "tts"), "mode": case.get("mode", "offline"), "device": args.device, "codec_device": args.codec_device, "case_name": case["id"], "timing_boundary": "per-request inference only; model load and wav write are excluded", "sequence_steps": steps, } stdout_lines.append("summary_json=" + json.dumps(summary, ensure_ascii=False, separators=(",", ":"))) (case_dir / "command.json").write_text( json.dumps( { "reference": "neutts_2e_python_warm_bench.py", "case": case["id"], "cases": str(args.cases), "reference_root": str(reference_root), "backbone": args.backbone, "codec": args.codec, "device": args.device, "codec_device": args.codec_device, "enable_watermark": args.enable_watermark, }, indent=2, ) + "\n", encoding="utf-8", ) (case_dir / "stdout.log").write_text("\n".join(stdout_lines) + "\n", encoding="utf-8") (case_dir / "memory.json").write_text(json.dumps(memory, indent=2) + "\n", encoding="utf-8") (case_dir / "summary.json").write_text( json.dumps(summary, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) print("summary_json=" + json.dumps(summary, ensure_ascii=False, separators=(",", ":")), flush=True) return 0 if __name__ == "__main__": raise SystemExit(main())