File size: 4,187 Bytes
f34bd6d | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """Reference streaming TTS runner over the split v2-Stream ONNX models — the exact
orchestration a sherpa-onnx C++ runner should mirror.
enc.onnx : (x,tone,lang,x_lengths,noise_scale,length_scale) -> z[1,192,T] (once)
dec.onnx : z[1,192,Tc] -> wav[1,1,Tc*256] (per chunk)
Streaming = run enc once, then decode z in CHUNK-frame steps via OVERLAP-SAVE:
for chunk frames [a,b) decode z[:, :, a-LEFT : b+RIGHT] and keep the middle
(b-a)*256 samples. Bit-exact vs the monolithic model (validated: cos 1.000000,
maxerr ~1e-6). First audio arrives after enc + one chunk instead of the whole
utterance.
Usage:
python -m streaming.onnx_stream --enc <enc.onnx> --dec <dec.onnx> --ids <parity_inputs.json> [--i 0]
(or --text "..." with the g2pw frontend available)
"""
from __future__ import annotations
import argparse, json, time
import numpy as np
import onnxruntime as ort
C, HOP, CHUNK, LEFT, RIGHT = 192, 256, 24, 64, 4 # validated streaming params
def _blank(seq):
o = [0] * (2 * len(seq) + 1)
o[1::2] = seq
return np.array([o], np.int64)
class StreamingTTS:
def __init__(self, enc_path, dec_path, threads=2):
so = ort.SessionOptions(); so.intra_op_num_threads = threads; so.inter_op_num_threads = 1
self.enc = ort.InferenceSession(enc_path, so, providers=["CPUExecutionProvider"])
self.dec = ort.InferenceSession(dec_path, so, providers=["CPUExecutionProvider"])
def encode(self, phone_ids, tone_ids, lang_ids, noise_scale=0.667, length_scale=1.0):
x, tn, lg = _blank(phone_ids), _blank(tone_ids), _blank(lang_ids)
return self.enc.run(None, {
"x": x, "tone": tn, "lang": lg,
"x_lengths": np.array([x.shape[1]], np.int64),
"noise_scale": np.array([noise_scale], np.float32),
"length_scale": np.array([length_scale], np.float32)})[0] # [1,192,T]
def stream(self, z):
"""Yield audio chunks (np.float32) as z is decoded chunk-by-chunk."""
T = z.shape[2]
for a in range(0, T, CHUNK):
b = min(a + CHUNK, T); s0 = max(0, a - LEFT); e = min(T, b + RIGHT)
w = self.dec.run(None, {"z": z[:, :, s0:e]})[0].reshape(-1)
off = (a - s0) * HOP; keep = (b - a) * HOP
yield w[off:off + keep]
def synth(self, phone_ids, tone_ids, lang_ids, **kw):
z = self.encode(phone_ids, tone_ids, lang_ids, **kw)
return np.concatenate(list(self.stream(z))) if z.shape[2] else np.zeros(0, np.float32)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--enc", default="/home/luigi/mbvits_run/v2stream_split/v2stream_enc.onnx")
ap.add_argument("--dec", default="/home/luigi/mbvits_run/v2stream_split/v2stream_dec.onnx")
ap.add_argument("--ids", default="/home/luigi/mbvits_run/parity_inputs.json")
ap.add_argument("--i", type=int, default=0)
ap.add_argument("--text", default=None)
ap.add_argument("--out", default="/home/luigi/mbvits_run/onnx_stream_demo.wav")
ap.add_argument("--threads", type=int, default=2)
a = ap.parse_args()
tts = StreamingTTS(a.enc, a.dec, a.threads)
if a.text:
import sys; sys.path.insert(0, "/home/luigi/primetts-space")
import frontend_bopomofo as F
o = F.text_to_ids(a.text); p, t, l = o["phone_ids"], o["tone_ids"], o["lang_ids"]
else:
r = json.load(open(a.ids))["rows"][a.i]; p, t, l = r["phone_ids"], r["tone_ids"], r["lang_ids"]
t0 = time.perf_counter(); z = tts.encode(p, t, l); t_enc = time.perf_counter() - t0
chunks = []; tfirst = None
for c in tts.stream(z):
chunks.append(c)
if tfirst is None: tfirst = time.perf_counter() - t0
total = time.perf_counter() - t0
wav = np.concatenate(chunks); audio_s = len(wav) / 16000
pk = np.max(np.abs(wav)); wav = wav * (0.97 / pk) if pk > 1e-6 else wav
import soundfile as sf; sf.write(a.out, wav.astype(np.float32), 16000)
print(f"frames={z.shape[2]} audio={audio_s:.2f}s enc={t_enc*1e3:.0f}ms "
f"first-audio={tfirst*1e3:.0f}ms total={total*1e3:.0f}ms RTF={total/audio_s:.3f} -> {a.out}")
if __name__ == "__main__":
main()
|