Initial release

tools/onnx-export/dump_sa3_goldens.py

@@ -0,0 +1,132 @@
+#!/usr/bin/env python3
+"""Dump golden input/output tensor pairs for the StableStep GGML port.
+
+Runs the validated ONNX graphs (the numerical reference — cosine 0.99995 vs
+PyTorch) on fixed-seed inputs and writes raw little-endian f32/i64 .bin files
+plus a manifest.json describing shapes. The C++ GGML modules replay these in
+unit tests: load input.bin -> forward -> compare against output.bin
+(target cosine > 0.999 for BF16 weights).
+
+Components dumped:
+  text_enc:  input_ids [1,256] i64, attention_mask [1,256] u8 -> embeddings [1,256,768]
+  seconds:   seconds [1] f32 -> embed [1,768]
+  same_enc:  audio [1,2,524288] f32 -> latents [1,256,128]
+  same_dec:  latents [1,256,128] f32 -> audio [1,2,524288]
+  dit:       x [1,256,64], t [1], cross [1,257,768], glob [1,768],
+             local [1,257,64], pad [1,64] -> v [1,256,64]   (small T=64 for speed)
+
+Runs in the StableAudio3 uv venv:
+  uv run --with onnx --with onnxruntime python dump_sa3_goldens.py \
+      --onnx-dir <dir> --out-dir <dir>
+"""
+
+import argparse
+import json
+import os
+
+import numpy as np
+
+
+def save(out_dir, name, arr):
+    path = os.path.join(out_dir, name + ".bin")
+    np.ascontiguousarray(arr).tofile(path)
+    return {"file": name + ".bin", "shape": list(arr.shape), "dtype": str(arr.dtype)}
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--onnx-dir", required=True)
+    ap.add_argument("--out-dir", required=True)
+    args = ap.parse_args()
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    import onnxruntime as ort
+    load = lambda n: ort.InferenceSession(
+        os.path.join(args.onnx_dir, n), providers=["CPUExecutionProvider"])
+
+    rng = np.random.default_rng(42)
+    manifest = {}
+
+    # ── text encoder ────────────────────────────────────────────────────
+    # Realistic ids: the validation prompt's 26 tokens + pad, from tokens_csv
+    # if present, else synthetic small ids.
+    ids = np.zeros((1, 256), dtype=np.int64)
+    n_tok = 26
+    tok_csv = os.path.join(os.path.dirname(args.out_dir), "tokens_csv.txt")
+    if os.path.exists(tok_csv):
+        vals = [int(x) for x in open(tok_csv).read().strip().split(",")]
+        ids[0, :len(vals)] = vals[:256]
+        n_tok = sum(1 for v in vals if v != 0) or 26
+    else:
+        ids[0, :n_tok] = rng.integers(3, 50000, n_tok)
+    mask = np.zeros((1, 256), dtype=np.bool_)
+    mask[0, :n_tok] = True
+    s = load("sa3-text_encoder.onnx")
+    emb = s.run(None, {"input_ids": ids, "attention_mask": mask})[0]
+    manifest["text_enc"] = {
+        "inputs": {"input_ids": save(args.out_dir, "text_enc.input_ids", ids),
+                   "attention_mask": save(args.out_dir, "text_enc.attention_mask",
+                                          mask.astype(np.uint8))},
+        "outputs": {"embeddings": save(args.out_dir, "text_enc.embeddings", emb)},
+        "n_tokens": n_tok,
+    }
+    del s
+
+    # ── seconds embedder ────────────────────────────────────────────────
+    sec = np.array([203.8], dtype=np.float32)
+    s = load("sa3-seconds_embedder.onnx")
+    sec_emb = s.run(None, {"seconds": sec})[0]
+    manifest["seconds"] = {
+        "inputs": {"seconds": save(args.out_dir, "seconds.in", sec)},
+        "outputs": {"embed": save(args.out_dir, "seconds.embed", sec_emb)},
+    }
+    del s
+
+    # ── SAME encoder / decoder (one static chunk each) ──────────────────
+    audio = (rng.standard_normal((1, 2, 524288)) * 0.1).astype(np.float32)
+    s = load("sa3-same_encoder.onnx")
+    latents = s.run(None, {"audio": audio})[0]
+    manifest["same_enc"] = {
+        "inputs": {"audio": save(args.out_dir, "same_enc.audio", audio)},
+        "outputs": {"latents": save(args.out_dir, "same_enc.latents", latents)},
+    }
+    del s
+
+    s = load("sa3-same_decoder.onnx")
+    dec_audio = s.run(None, {"latents": latents})[0]
+    manifest["same_dec"] = {
+        "inputs": {"latents": save(args.out_dir, "same_dec.latents", latents)},
+        "outputs": {"audio": save(args.out_dir, "same_dec.audio", dec_audio)},
+    }
+    del s
+
+    # ── DiT single forward at small T ───────────────────────────────────
+    T = 64
+    x = rng.standard_normal((1, 256, T)).astype(np.float32)
+    t = np.array([0.3], dtype=np.float32)
+    cross = np.concatenate([emb, sec_emb[:, None, :]], axis=1).astype(np.float32)
+    glob = sec_emb.astype(np.float32)
+    local = np.zeros((1, 257, T), dtype=np.float32)
+    pad = np.ones((1, T), dtype=np.bool_)
+    pad[0, 48:] = False  # exercise the padding-mask path
+    s = load("sa3-dit.onnx")
+    v = s.run(None, {"x": x, "t": t, "cross_attn_cond": cross,
+                     "global_embed": glob, "local_add_cond": local,
+                     "padding_mask": pad})[0]
+    manifest["dit"] = {
+        "inputs": {"x": save(args.out_dir, "dit.x", x),
+                   "t": save(args.out_dir, "dit.t", t),
+                   "cross_attn_cond": save(args.out_dir, "dit.cross", cross),
+                   "global_embed": save(args.out_dir, "dit.glob", glob),
+                   "local_add_cond": save(args.out_dir, "dit.local", local),
+                   "padding_mask": save(args.out_dir, "dit.pad", pad.astype(np.uint8))},
+        "outputs": {"v": save(args.out_dir, "dit.v", v)},
+    }
+
+    with open(os.path.join(args.out_dir, "manifest.json"), "w") as f:
+        json.dump(manifest, f, indent=2)
+    print(f"Goldens written to {args.out_dir}")
+
+
+if __name__ == "__main__":
+    main()