export_onnx.py

#!/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()