Initial release

tools/onnx-export/export_sa3_conditioners.py

@@ -0,0 +1,185 @@
+#!/usr/bin/env python3
+"""Export the Stable Audio 3 conditioners to ONNX: T5Gemma text encoder + seconds_total embedder.
+
+Part of the SA3 post-processing refiner port.
+
+Tensor specs (fp32):
+  Text encoder (tokenization stays outside — HF tokenizer.json, 256 max length,
+  pad to max with learned-padding substitution baked into the graph):
+    "input_ids"      [1, 256] int64
+    "attention_mask" [1, 256] bool
+    -> "embeddings"  [1, 256, 768]
+
+  Seconds embedder (replaces hand-porting Fourier-feature math to C++):
+    "seconds"     [1] float32   (clamped/normalized inside the graph, max 384s)
+    -> "embed"    [1, 768]      (used both as global_embed and, unsqueezed, as
+                                 the extra cross-attention token after the prompt)
+
+Runs in the StableAudio3 uv venv:
+  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_conditioners.py
+"""
+
+import argparse
+import json
+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")
+
+from safetensors import safe_open
+from stable_audio_3.model_configs import models
+from stable_audio_3.factory import create_multi_conditioner_from_conditioning_config
+
+# transformers v5 mask construction (vmap-based) doesn't trace to ONNX. The
+# encoder is bidirectional with plain padding, so a broadcast bool keep-mask
+# [B,1,Q,K] is equivalent — create_bidirectional_mask's contract accepts a
+# prepared 4D mask. Patch at the t5gemma module level (direct name import).
+import transformers.models.t5gemma.modeling_t5gemma as t5g_mod
+
+
+def _trace_friendly_bidirectional_mask(config=None, inputs_embeds=None,
+                                       attention_mask=None, **kwargs):
+    if attention_mask is None:
+        return None
+    q = inputs_embeds.shape[1]
+    return attention_mask.to(torch.bool)[:, None, None, :].expand(
+        attention_mask.shape[0], 1, q, attention_mask.shape[-1]
+    )
+
+
+def _trace_friendly_sliding_window_mask(config=None, inputs_embeds=None,
+                                        attention_mask=None, **kwargs):
+    window = getattr(config, "sliding_window", None) or 4096
+    q = inputs_embeds.shape[1]
+    idx = torch.arange(q, device=inputs_embeds.device)
+    band = (idx[None, :] - idx[:, None]).abs() < window
+    mask = band[None, None, :, :]
+    if attention_mask is not None:
+        mask = mask & attention_mask.to(torch.bool)[:, None, None, :]
+    return mask.expand(inputs_embeds.shape[0], 1, q, q)
+
+
+t5g_mod.create_bidirectional_mask = _trace_friendly_bidirectional_mask
+t5g_mod.create_bidirectional_sliding_window_mask = _trace_friendly_sliding_window_mask
+
+
+class TextEncWrapper(nn.Module):
+    def __init__(self, cond):
+        super().__init__()
+        self.cond = cond
+
+    def forward(self, input_ids, attention_mask):
+        emb = self.cond.model(input_ids=input_ids, attention_mask=attention_mask)["last_hidden_state"]
+        emb = self.cond.proj_out(emb)
+        emb = self.cond.apply_padding(emb, attention_mask)
+        return emb
+
+
+class SecondsWrapper(nn.Module):
+    def __init__(self, cond):
+        super().__init__()
+        self.cond = cond
+
+    def forward(self, seconds):
+        x = seconds.clamp(self.cond.min_val, self.cond.max_val)
+        x = (x - self.cond.min_val) / (self.cond.max_val - self.cond.min_val)
+        return self.cond.embedder(x)
+
+
+def parity(tag, path, feeds, ref):
+    import onnxruntime as ort
+    sess = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
+    out = sess.run(None, feeds)[0]
+    ref_np = ref.float().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"  [{tag}] 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")
+    args = ap.parse_args()
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    cfg_path, ckpt_path = models["medium"].resolve()
+    with open(cfg_path) as f:
+        config = json.load(f)
+
+    print("Building conditioners (T5Gemma from HF subfolder)...")
+    conditioner = create_multi_conditioner_from_conditioning_config(
+        config["model"]["conditioning"]
+    )
+
+    # Learned-padding embeddings etc. live in the main checkpoint under conditioner.*
+    with safe_open(ckpt_path, framework="pt", device="cpu") as f:
+        cond_sd = {
+            k[len("conditioner."):]: f.get_tensor(k)
+            for k in f.keys() if k.startswith("conditioner.")
+        }
+    missing, unexpected = conditioner.load_state_dict(cond_sd, strict=False)
+    print(f"  conditioner tensors loaded: {len(cond_sd)} (missing={len(missing)}, unexpected={len(unexpected)})")
+
+    prompt_cond = conditioner.conditioners["prompt"]
+    prompt_cond.model.float().eval().requires_grad_(False)
+    prompt_cond.proj_out.float()
+    seconds_cond = conditioner.conditioners["seconds_total"].float().eval().requires_grad_(False)
+
+    # --- Text encoder ------------------------------------------------------
+    text = "Instrumental punk rock with distorted electric guitars. BPM: 160. Length: 200 seconds."
+    enc = prompt_cond.tokenizer(
+        [text], truncation=True, max_length=prompt_cond.max_length,
+        padding="max_length", return_tensors="pt",
+    )
+    input_ids = enc["input_ids"]
+    attention_mask = enc["attention_mask"].to(torch.bool)
+
+    wrapper = TextEncWrapper(prompt_cond).eval()
+    # Conditioner stores the HF model outside nn.Module registration (enable_grad
+    # False path uses __dict__) — reattach for export.
+    wrapper.cond.model.eval()
+    with torch.no_grad():
+        ref_emb = wrapper(input_ids, attention_mask)
+    print(f"Text encoder reference: {tuple(ref_emb.shape)}")
+
+    text_path = os.path.join(args.output_dir, "sa3-text_encoder.onnx")
+    torch.onnx.export(
+        wrapper, (input_ids, attention_mask), text_path,
+        input_names=["input_ids", "attention_mask"], output_names=["embeddings"],
+        opset_version=18, dynamo=False,
+    )
+    print(f"Exported {text_path}")
+    c1 = parity("text-enc", text_path,
+                {"input_ids": input_ids.numpy(), "attention_mask": attention_mask.numpy()},
+                ref_emb)
+
+    # --- Seconds embedder --------------------------------------------------
+    sw = SecondsWrapper(seconds_cond).eval()
+    seconds = torch.tensor([203.8], dtype=torch.float32)
+    with torch.no_grad():
+        ref_sec = sw(seconds)
+    sec_path = os.path.join(args.output_dir, "sa3-seconds_embedder.onnx")
+    torch.onnx.export(
+        sw, (seconds,), sec_path,
+        input_names=["seconds"], output_names=["embed"],
+        opset_version=18, dynamo=False,
+    )
+    print(f"Exported {sec_path} ({os.path.getsize(sec_path)/1e6:.1f} MB)")
+    c2 = parity("seconds", sec_path, {"seconds": seconds.numpy()}, ref_sec)
+
+    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())