138 lines
4.4 KiB
Python
138 lines
4.4 KiB
Python
#!/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())
|