File size: 2,919 Bytes
a37967e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | #!/usr/bin/env python3
"""RANK-1 acoustic probe: synth N training clips with FORCED (ground-truth) durations through an
ONNX dir, isolating the phones->mel mapping from the duration predictor + g2p frontend.
Reads phone/tone/lang ids + GT durations directly from m2_align.jsonl (no frontend).
Run in moss-train-venv. Then ASR the wavs (probe_forced_asr via xasr_offline) -> CER.
Pairs the 0.80(forced)-vs-0.145(GT-mel) acoustic gap to a single number per config."""
import argparse, json, sys
from pathlib import Path
import numpy as np, soundfile as sf, onnxruntime as ort
ZT = "/home/luigi/jetson-tts/mossnano/zhtw8k"
sys.path.insert(0, ZT)
from synth_from_text import host_regulate
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--onnx-dir", required=True)
ap.add_argument("--out-dir", required=True)
ap.add_argument("--n", type=int, default=30)
ap.add_argument("--align-jsonl", default=f"{ZT}/m2_align.jsonl")
args = ap.parse_args()
meta = json.load(open(f"{args.onnx_dir}/meta.json"))
so = ort.SessionOptions(); so.intra_op_num_threads = 4
sA = ort.InferenceSession(f"{args.onnx_dir}/acoustic_encoder.onnx", so, providers=["CPUExecutionProvider"])
sB = ort.InferenceSession(f"{args.onnx_dir}/acoustic_decoder.onnx", so, providers=["CPUExecutionProvider"])
sV = ort.InferenceSession(f"{args.onnx_dir}/vocoder.onnx", so, providers=["CPUExecutionProvider"])
bn = ["frames", "frame_meta", "local_ctx_raw", "abs_pos", "pitch_frame", "frame_mask"]
Path(args.out_dir).mkdir(parents=True, exist_ok=True)
rows = [json.loads(l) for l in open(args.align_jsonl) if l.strip()][:args.n]
out = open(f"{args.out_dir}/synth.jsonl", "w")
for i, r in enumerate(rows):
phone = np.array([r["phone_ids"]], np.int64); tone = np.array([r["tone_ids"]], np.int64)
lang = np.array([r["lang_ids"]], np.int64); spk = np.zeros(1, np.int64)
cond, _dur_pred, pitch = sA.run(None, {"phone": phone, "tone": tone, "lang": lang, "speaker": spk})
# substitute GT (forced) durations, rescaled so total ~ predicted length (stable regulator)
df = np.array([r["hifigan_durations"]], np.float32)
df = df * (_dur_pred.sum() / max(1.0, df.sum()))
reg = host_regulate(cond, df, pitch, meta["abs_frame_bins"], meta["max_frames"])
feeds = {n: (reg[n].astype(np.float32) if reg[n].dtype != bool else reg[n]) for n in bn}
feeds["abs_pos"] = reg["abs_pos"].astype(np.int64)
mel = sB.run(None, feeds)[0]
wav = sV.run(None, {"mel": mel.astype(np.float32)})[0].reshape(-1)
wp = f"{args.out_dir}/p{i:03d}.wav"; sf.write(wp, wav, meta["sample_rate"])
out.write(json.dumps({"id": f"p{i:03d}", "text": r["text"], "wav": wp}, ensure_ascii=False) + "\n")
out.close()
print(f"PROBE SYNTH DONE {len(rows)} clips -> {args.out_dir}/synth.jsonl")
if __name__ == "__main__":
main()
|