#!/usr/bin/env python3 """ export_onnx.py — regenerate susurro.onnx from susurro.pth (reproducibility). Traces the exact voicepack inference path (see infer.py) into a single ONNX graph inputs : input_ids [1, T] int64 (already wrapped [0, *tokens, 0]), ref_s [1, 256] output : audio [N] float32 (24 kHz) with three swaps that make it ONNX-exportable & deterministic: * TorchSTFT (torch.stft/istft, complex) -> ONNXSTFT (conv1d / conv_transpose1d) * SineGen randomness (phase + noise) -> dropped (deterministic, sub-perceptual) * duration->alignment python loop -> vectorized cumsum/compare (dynamic frames) * InstanceNorm / LSTM packing / neg-perm -> export-safe equivalents (batch=1) pip install -r requirements-raw.txt onnx onnxruntime python export_onnx.py # writes susurro.onnx + prints ONNX-vs-PyTorch parity """ from __future__ import annotations import sys import types from pathlib import Path import numpy as np HERE = Path(__file__).resolve().parent REPO = HERE / "styletts2" sys.path.insert(0, str(HERE)) sys.path.insert(0, str(REPO)) from infer import load_model, g2p_en, _patch_torch_load # noqa: E402 _patch_torch_load() import torch # noqa: E402 import torch.nn as nn # noqa: E402 import torch.nn.functional as F # noqa: E402 from onnx_stft import ONNXSTFT # noqa: E402 from Modules.istftnet import TorchSTFT, SineGen # noqa: E402 DEVICE = "cpu" def det_f02sine(self, f0_values): rad = (f0_values / self.sampling_rate) % 1 rad = F.interpolate(rad.transpose(1, 2), scale_factor=1 / self.upsample_scale, mode="linear").transpose(1, 2) phase = torch.cumsum(rad, dim=1) * 2 * np.pi phase = F.interpolate(phase.transpose(1, 2) * self.upsample_scale, scale_factor=self.upsample_scale, mode="linear").transpose(1, 2) two_pi = 2 * np.pi phase = phase - two_pi * torch.floor(phase / two_pi) # wrap: float32 sin(large) drifts return torch.sin(phase) def det_sinegen_forward(self, f0): fn = torch.multiply( f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device)) sine = self._f02sine(fn) * self.sine_amp uv = self._f02uv(f0) return sine * uv, uv, torch.zeros_like(uv) # drop additive noise class INorm(nn.Module): def __init__(self, ref, dims): super().__init__() self.eps, self.dims, self.affine = ref.eps, dims, ref.affine if ref.affine: self.weight, self.bias = ref.weight, ref.bias def forward(self, x): mean = x.mean(dim=self.dims, keepdim=True) var = ((x - mean) ** 2).mean(dim=self.dims, keepdim=True) y = (x - mean) / torch.sqrt(var + self.eps) if self.affine: shp = [1, -1] + [1] * len(self.dims) y = y * self.weight.view(*shp) + self.bias.view(*shp) return y def te_forward(self, x, input_lengths, m): x = self.embedding(x).transpose(1, 2) m = m.unsqueeze(1) x = x.masked_fill(m, 0.0) for c in self.cnn: x = c(x).masked_fill(m, 0.0) x = x.transpose(1, 2) x, _ = self.lstm(x) return x.transpose(-1, -2).masked_fill(m, 0.0) def de_forward(self, x, style, text_lengths, m): from models import AdaLayerNorm x = x.permute(2, 0, 1) s = style.expand(x.shape[0], x.shape[1], -1) x = torch.cat([x, s], axis=-1) x = x.masked_fill(m.unsqueeze(-1).transpose(0, 1), 0.0).transpose(0, 1).transpose(-1, -2) for block in self.lstms: if isinstance(block, AdaLayerNorm): x = block(x.transpose(-1, -2), style).transpose(-1, -2) x = torch.cat([x, s.permute(1, 2, 0)], axis=1) x = x.masked_fill(m.unsqueeze(-1).transpose(-1, -2), 0.0) else: x = x.transpose(-1, -2) x, _ = block(x) x = x.transpose(-1, -2) return x.transpose(-1, -2) def patch(model): def swap_stft(mod): for cn, ch in list(mod.named_children()): if isinstance(ch, TorchSTFT): setattr(mod, cn, ONNXSTFT(ch.filter_length, ch.hop_length, ch.win_length)) else: swap_stft(ch) swap_stft(model.decoder) for m in model.decoder.modules(): if isinstance(m, SineGen): m._f02sine = types.MethodType(det_f02sine, m) m.forward = types.MethodType(det_sinegen_forward, m) model.text_encoder.forward = types.MethodType(te_forward, model.text_encoder) model.predictor.text_encoder.forward = types.MethodType( de_forward, model.predictor.text_encoder) def swap_in(root): for cn, ch in list(root.named_children()): if isinstance(ch, nn.InstanceNorm1d): setattr(root, cn, INorm(ch, (2,))) elif isinstance(ch, nn.InstanceNorm2d): setattr(root, cn, INorm(ch, (2, 3))) else: swap_in(ch) for m in (model.bert_encoder, model.predictor, model.text_encoder, model.decoder): swap_in(m) class SusurroONNX(nn.Module): def __init__(self, model): super().__init__() self.bert, self.bert_encoder = model.bert, model.bert_encoder self.predictor, self.text_encoder = model.predictor, model.text_encoder self.decoder = model.decoder def forward(self, input_ids, ref_s): ref_acoustic, ref_prosodic = ref_s[:, :128], ref_s[:, 128:] L = input_ids.shape[1] input_lengths = torch.tensor([L], dtype=torch.long, device=input_ids.device) text_mask = torch.gt(torch.arange(L, device=input_ids.device).unsqueeze(0) + 1, input_lengths.unsqueeze(1)) bert_dur = self.bert(input_ids, attention_mask=(~text_mask).int()) d_en = self.bert_encoder(bert_dur).transpose(-1, -2) d = self.predictor.text_encoder(d_en, ref_prosodic, input_lengths, text_mask) x, _ = self.predictor.lstm(d) duration = torch.sigmoid(self.predictor.duration_proj(x)).sum(dim=-1) pred_dur = torch.round(duration.squeeze(0)).clamp(min=1).long() cum = torch.cumsum(pred_dur, dim=0) starts = cum - pred_dur t = torch.arange(cum[-1], device=input_ids.device) aln = ((t.unsqueeze(0) >= starts.unsqueeze(1)) & (t.unsqueeze(0) < cum.unsqueeze(1))).float().unsqueeze(0) en = d.transpose(-1, -2) @ aln F0_pred, N_pred = self.predictor.F0Ntrain(en, ref_prosodic) t_en = self.text_encoder(input_ids, input_lengths, text_mask) asr = t_en @ aln return self.decoder(asr, F0_pred, N_pred, ref_acoustic).squeeze() def main(): model, tc = load_model(str(HERE / "config.yml"), str(HERE / "susurro.pth"), DEVICE) patch(model) net = SusurroONNX(model).eval() ids = torch.LongTensor([[0, *tc(g2p_en("Hey, I wasn't expecting you tonight.")), 0]]) ref = torch.from_numpy( np.load(HERE / "voicepacks.npz")["voice_a__neutral"]).reshape(1, 256) with torch.no_grad(): ref_audio = net(ids, ref).cpu().numpy() dur, peak = len(ref_audio) / 24000, float(np.abs(ref_audio).max()) print(f"[export] torch output: {dur:.2f}s peak={peak:.3f}") if not (1.0 < dur < 8.0 and peak < 5.0): raise SystemExit("ABORT: patched output insane — weights/patches broken") out = str(HERE / "susurro.onnx") torch.onnx.export(net, (ids, ref), out, input_names=["input_ids", "ref_s"], output_names=["audio"], opset_version=17, do_constant_folding=True, dynamic_axes={"input_ids": {1: "tokens"}, "audio": {0: "samples"}}) import onnxruntime as ort sess = ort.InferenceSession(out, providers=["CPUExecutionProvider"]) oa = sess.run(None, {"input_ids": ids.numpy(), "ref_s": ref.numpy()})[0] L = min(len(ref_audio), len(oa)) a, b = ref_audio[:L], oa[:L] print(f"[export] wrote {out}") print(f"[export] PARITY corr={np.corrcoef(a, b)[0,1]:.5f} " f"max_abs_err={np.abs(a-b).max():.2e} rms_err={np.sqrt(np.mean((a-b)**2)):.2e}") if __name__ == "__main__": main()