Initial release

tools/onnx-export/export_sa3_same.py

@@ -0,0 +1,137 @@
+#!/usr/bin/env python3
+"""Export Stable Audio 3 SAME-L autoencoder (encoder + decoder) to ONNX.
+
+Part of the SA3 post-processing refiner port. Exports fixed-size chunk graphs;
+the C++ engine tiles with overlap-trim exactly like the Oobleck tiled decode.
+
+Tensor specs (fp32, static shapes — TRT-friendly):
+  Encoder: "audio"   [1, 2, 524288]  (stereo 44.1kHz, 128-latent chunk)
+        -> "latents" [1, 256, 128]
+  Decoder: "latents" [1, 256, 128]
+        -> "audio"   [1, 2, 524288]
+
+Stochastic decode paths (bottleneck noise_regularize, resampler mask_noise) are
+zeroed for determinism — quality impact to be validated by listening test.
+
+Runs in the StableAudio3 uv venv (NOT hot-step-9000):
+  cd d:/Ace-Step-Latest/StableAudio3
+  uv run --with onnx --with onnxruntime python \
+      d:/Ace-Step-Latest/hot-step-cpp/tools/onnx-export/export_sa3_same.py
+"""
+
+import argparse
+import os
+import sys
+import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+sys.path.insert(0, r"d:/Ace-Step-Latest/StableAudio3")
+
+# Force the plain/chunked-halo SDPA attention tiers — flash and flex_attention
+# do not trace to ONNX.
+import stable_audio_3.models.transformer as sat
+sat.flash_attn_func = None
+sat.flash_attn_kvpacked_func = None
+sat.flex_attention_available = False
+sat.flex_attention_compiled = None
+
+from stable_audio_3.model_configs import ae_models
+from stable_audio_3.loading_utils import load_autoencoder
+
+CHUNK_LATENTS = 128
+DOWNSAMPLING = 4096
+CHUNK_SAMPLES = CHUNK_LATENTS * DOWNSAMPLING  # 524288
+
+
+def zero_stochastic_paths(ae):
+    ae.bottleneck.noise_regularize = False
+    for m in ae.modules():
+        if hasattr(m, "mask_noise"):
+            m.mask_noise = 0
+
+
+class EncoderWrapper(nn.Module):
+    def __init__(self, ae):
+        super().__init__()
+        self.ae = ae
+
+    def forward(self, audio):
+        return self.ae.encode(audio)
+
+
+class DecoderWrapper(nn.Module):
+    def __init__(self, ae):
+        super().__init__()
+        self.ae = ae
+
+    def forward(self, latents):
+        return self.ae.decode(latents)
+
+
+def parity(name, onnx_path, feed_name, feed, ref):
+    import onnxruntime as ort
+    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
+    t0 = time.time()
+    out = sess.run(None, {feed_name: feed.cpu().numpy()})[0]
+    ref_np = ref.cpu().numpy()
+    max_abs = np.abs(out - ref_np).max()
+    denom = np.linalg.norm(out.ravel()) * np.linalg.norm(ref_np.ravel())
+    cos = float(np.dot(out.ravel(), ref_np.ravel()) / denom) if denom > 0 else 0.0
+    print(f"  [{name}] ORT-CPU {time.time()-t0:.1f}s  max_abs_diff={max_abs:.3e}  cosine={cos:.6f}")
+    return cos
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--output-dir", default=r"d:/Ace-Step-Latest/hot-step-cpp/models/onnx/sa3")
+    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    args = ap.parse_args()
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    print("Loading SAME-L...")
+    cfg_path, ckpt_path = ae_models["same-l"].resolve()
+    ae = load_autoencoder(cfg_path, ckpt_path, device=args.device).eval().requires_grad_(False)
+    zero_stochastic_paths(ae)
+
+    torch.manual_seed(0)
+    audio = (torch.randn(1, 2, CHUNK_SAMPLES, device=args.device) * 0.1).clamp(-1, 1)
+
+    enc = EncoderWrapper(ae)
+    with torch.no_grad():
+        ref_latents = enc(audio)
+    print(f"Encoder reference: {tuple(ref_latents.shape)}")
+
+    enc_path = os.path.join(args.output_dir, "sa3-same_encoder.onnx")
+    torch.onnx.export(
+        enc, (audio,), enc_path,
+        input_names=["audio"], output_names=["latents"],
+        opset_version=18, dynamo=False,
+    )
+    print(f"Exported {enc_path} ({os.path.getsize(enc_path)/1e9:.2f} GB)")
+
+    dec = DecoderWrapper(ae)
+    with torch.no_grad():
+        ref_audio = dec(ref_latents)
+    print(f"Decoder reference: {tuple(ref_audio.shape)}")
+
+    dec_path = os.path.join(args.output_dir, "sa3-same_decoder.onnx")
+    torch.onnx.export(
+        dec, (ref_latents,), dec_path,
+        input_names=["latents"], output_names=["audio"],
+        opset_version=18, dynamo=False,
+    )
+    print(f"Exported {dec_path} ({os.path.getsize(dec_path)/1e9:.2f} GB)")
+
+    print("Parity vs PyTorch (fp32):")
+    c1 = parity("encoder", enc_path, "audio", audio, ref_latents)
+    c2 = parity("decoder", dec_path, "latents", ref_latents, ref_audio)
+    ok = c1 > 0.999 and c2 > 0.999
+    print("PARITY OK" if ok else "PARITY FAILED")
+    return 0 if ok else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())