civ/hot-step-cpp-ROCm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial release

Browse Source
This commit is contained in:
civ
2026-08-16 18:24:52 +07:00
commit 876886a39a
13244 changed files with 2353959 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/onnx-export/README.md
+82
View File
@@ -0,0 +1,82 @@
# ONNX Export Tools
Scripts for exporting HOT-Step model components to ONNX format for use with
TensorRT or ONNX Runtime inference.
## Prerequisites
- Python 3.10+ with the hot-step-9000 venv
- PyTorch with CUDA support
- `diffusers`, `onnx`, `onnxruntime` (or `onnxruntime-gpu` for GPU/TRT)
## Scripts
### `export_vae.py` — VAE Decoder Export
Exports the `AutoencoderOobleck` VAE decoder to ONNX format. Only the decoder
half is exported (encoder is not needed for inference — we decode latents to
audio).
**Tensor spec:**
| Name | Shape | Description |
|------|-------|-------------|
| `latents` (input) | `[B, 64, T]` | Latent channels, latent frames @ 25Hz |
| `audio` (output) | `[B, 2, S]` | Stereo audio, S = T × 1920 @ 48kHz |
Dynamic axes: batch (dim 0) and temporal dims (latent_frames, samples).
**Usage:**
```powershell
# Basic export
& .venv\Scripts\python.exe tools\onnx-export\export_vae.py `
--vae-path checkpoints\vae `
--output models\onnx\vae_decoder.onnx
# Export with validation (compares ONNX vs PyTorch output)
& .venv\Scripts\python.exe tools\onnx-export\export_vae.py `
--vae-path checkpoints\vae `
--output models\onnx\vae_decoder.onnx `
--validate
```
**Options:**
- `--vae-path` — Path to VAE checkpoint directory (config.json + safetensors)
- `--output` — Output path for the ONNX file
- `--opset` — ONNX opset version (default: 18)
- `--validate` — Compare ONNX output against PyTorch using onnxruntime
### `test_trt_vae.py` — TensorRT Validation & Benchmark
Benchmarks the exported ONNX model using CUDA EP vs TensorRT EP. Reports
latency, speedup ratio, and numerical accuracy.
**Usage:**
```powershell
& .venv\Scripts\python.exe tools\onnx-export\test_trt_vae.py `
--onnx models\onnx\vae_decoder.onnx
```
**Options:**
- `--onnx` — Path to the exported ONNX file
- `--trt-cache` — Directory for TRT engine cache (default: `models/onnx/trt_cache/`)
- `--iterations` — Number of benchmark iterations (default: 20)
**Requirements for TRT EP:**
- `onnxruntime-gpu` (not `onnxruntime`)
- TensorRT libraries on PATH
- The script gracefully falls back if TRT EP is unavailable
## Output Files
| File | Size | Git-tracked? |
|------|------|-------------|
| `models/onnx/vae_decoder.onnx` | ~330 MB | ❌ No (gitignored) |
| `models/onnx/trt_cache/*.engine` | ~200 MB | ❌ No (gitignored) |
## Notes
- Export is always in fp32. TensorRT handles fp16 conversion during engine build.
- The VAE is a simple 1D convolutional network (no attention layers), so ONNX
export is straightforward — no trace-safe patches needed.
- ScragVAE (674MB) can also be exported using the same script with a different
`--vae-path`.
tools/onnx-export/convert_sa3_fp16.py
+88
View File
@@ -0,0 +1,88 @@
#!/usr/bin/env python3
"""Convert the SA3 ONNX graph set to fp16 for shipping (halves ~12GB fp32 -> ~6GB).
Weights/compute go fp16, graph I/O stays fp32 (keep_io_types) so the C++
orchestration is precision-agnostic. Validate afterwards by re-running
e2e_sa3_ort.py against the fp16 directory (expect cosine > 0.99 vs the
PyTorch fp32 reference — the production Python pipeline ran fp16 anyway).
Runs in the StableAudio3 uv venv:
uv run --with onnx --with onnxconverter-common python convert_sa3_fp16.py \
--input-dir <fp32 dir> --output-dir <fp16 dir>
"""
import argparse
import os
import onnx
from onnxconverter_common import float16
# Text encoder stays fp32: onnxconverter-common emits invalid mixed-dtype casts
# around its bool-mask paths, and the engine keeps text encoders fp32 anyway
# (text-enc-ort.h: "FP32: layernorm overflows in FP16"). Seconds embedder is 0.8MB.
GRAPHS_FP16 = [
"sa3-same_encoder.onnx",
"sa3-same_decoder.onnx",
"sa3-dit.onnx",
]
GRAPHS_COPY_FP32 = [
"sa3-text_encoder.onnx",
"sa3-seconds_embedder.onnx",
]
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--input-dir", required=True)
ap.add_argument("--output-dir", required=True)
args = ap.parse_args()
os.makedirs(args.output_dir, exist_ok=True)
import shutil
for name in GRAPHS_COPY_FP32:
# NOT a file copy: the export may reference external per-tensor files —
# load (resolves them) and re-save self-contained (<2GB, fits inline).
model = onnx.load(os.path.join(args.input_dir, name))
onnx.save_model(model, os.path.join(args.output_dir, name))
print(f"Repacked {name} (fp32, self-contained)")
for name in GRAPHS_FP16:
src = os.path.join(args.input_dir, name)
dst = os.path.join(args.output_dir, name)
print(f"Converting {name}...")
if name == "sa3-dit.onnx":
# >2GB: in-memory shape inference hits the protobuf limit — infer on
# disk. The temp file MUST live next to src: the fp32 export stores
# weights as per-tensor external files resolved relative to the model.
inferred = src + ".inferred"
onnx.shape_inference.infer_shapes_path(src, inferred)
model = onnx.load(inferred) # pulls external data fully into memory
os.remove(inferred)
model_fp16 = float16.convert_float_to_float16(
model, keep_io_types=True, disable_shape_infer=True
)
else:
# Shape inference ON — without it the converter misses boundary
# casts and emits mixed-dtype nodes (invalid graph).
model = onnx.load(src)
model_fp16 = float16.convert_float_to_float16(model, keep_io_types=True)
# Large graphs (DiT) exceed the 2GB protobuf limit even at fp16 with
# metadata — always save with external data for uniform loading.
onnx.save_model(
model_fp16, dst,
save_as_external_data=(name == "sa3-dit.onnx"),
all_tensors_to_one_file=True,
location=os.path.basename(dst) + ".data",
)
total = os.path.getsize(dst)
data = dst + ".data"
if os.path.exists(data):
total += os.path.getsize(data)
print(f" -> {total/1e9:.2f} GB")
print("Done.")
if __name__ == "__main__":
main()
tools/onnx-export/dump_sa3_goldens.py
+132
View File
@@ -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()
tools/onnx-export/dump_sa3_same_stages.py
+114
View File
@@ -0,0 +1,114 @@
#!/usr/bin/env python3
"""Dump per-stage SAME-L autoencoder activations for GGML parity debugging.
Loads SAME-L exactly like export_sa3_same.py (noise paths zeroed), feeds the
SAME golden inputs used by sa3-ggml-test, and dumps the token sequence
(b, S, 1536) f32 before/after each TransformerBlock:
<out>/enc_stage0.bin = input to encoder transformers[0] (folded seq + new tokens)
<out>/enc_stage<i>.bin = output of encoder transformers[i-1], i = 1..12
<out>/enc_final.bin = final latents (1, 256, 128)
(same for dec_*)
Matches the C++ side: SA3_SAME_STAGE=<n> SA3_SAME_DUMP=<path> sa3-ggml-test
dumps the corresponding [dim, S] token-major tensor.
Run:
cd d:/Ace-Step-Latest/StableAudio3
uv run python d:/Ace-Step-Latest/hot-step-cpp/tools/onnx-export/dump_sa3_same_stages.py \
--goldens <dir-with-same_enc.audio.bin> --out <dir>
"""
import argparse
import os
import sys
import numpy as np
import torch
sys.path.insert(0, r"d:/Ace-Step-Latest/StableAudio3")
# Same attention-tier forcing as the ONNX export (goldens came from this path).
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
def zero_stochastic_paths(ae):
ae.bottleneck.noise_regularize = False
for m in ae.modules():
if hasattr(m, "mask_noise"):
m.mask_noise = 0
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--goldens", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--side", default="both", choices=["enc", "dec", "both"])
ap.add_argument("--device", default="cpu")
args = ap.parse_args()
os.makedirs(args.out, exist_ok=True)
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)
def dump(path, t):
arr = t.detach().float().cpu().numpy()
arr.tofile(path)
print(f" {os.path.basename(path)} shape={tuple(arr.shape)}")
def hook_block(block_module, prefix):
stages = {}
def pre_hook(mod, args_, kwargs):
x = args_[0]
if 0 not in stages:
stages[0] = x.detach().clone()
handles = [block_module.transformers[0].register_forward_pre_hook(pre_hook, with_kwargs=True)]
for i, layer in enumerate(block_module.transformers):
def post_hook(mod, args_, output, idx=i):
if idx + 1 not in stages:
stages[idx + 1] = output.detach().clone()
handles.append(layer.register_forward_hook(post_hook))
return stages, handles
if args.side in ("enc", "both"):
audio = np.fromfile(os.path.join(args.goldens, "same_enc.audio.bin"), dtype=np.float32)
audio = torch.from_numpy(audio.reshape(1, 2, -1)).to(args.device)
block = ae.encoder.layers[0]
stages, handles = hook_block(block, "enc")
with torch.no_grad():
latents = ae.encode(audio)
for h in handles:
h.remove()
print("[enc]")
for n, t in sorted(stages.items()):
dump(os.path.join(args.out, f"enc_stage{n}.bin"), t)
dump(os.path.join(args.out, "enc_final.bin"), latents)
if args.side in ("dec", "both"):
lat = np.fromfile(os.path.join(args.goldens, "same_dec.latents.bin"), dtype=np.float32)
lat = torch.from_numpy(lat.reshape(1, 256, -1)).to(args.device)
block = ae.decoder.layers[3]
stages, handles = hook_block(block, "dec")
with torch.no_grad():
audio_out = ae.decode(lat)
for h in handles:
h.remove()
print("[dec]")
for n, t in sorted(stages.items()):
dump(os.path.join(args.out, f"dec_stage{n}.bin"), t)
dump(os.path.join(args.out, "dec_final.bin"), audio_out)
if __name__ == "__main__":
main()
tools/onnx-export/e2e_sa3_ort.py
+256
View File
@@ -0,0 +1,256 @@
#!/usr/bin/env python3
"""End-to-end acceptance gate for the SA3 ONNX export set.
Runs the FULL SDEdit refine (encode -> 8-step Euler -> decode) twice:
1. Reference: StableAudioModel.generate() in PyTorch (fp32, CUDA, fixed seed,
stochastic AE paths zeroed to match the exported graphs)
2. Harness: ONLY the four ONNX graphs (text enc, seconds embedder, SAME enc,
DiT, SAME dec) + numpy orchestration that mirrors generate()/sample_diffusion.
This orchestration is the exact spec the C++ engine implements.
Deterministic by construction: sampler_type=euler (no mid-loop RNG; production
default pingpong differs only by a per-step renoise draw), initial noise drawn
once in torch with the same seed/device as the reference.
Repo pure-math helpers (schedule, effective length) are imported rather than
copied — the C++ port reimplements them with unit tests against these.
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/e2e_sa3_ort.py
"""
import argparse
import os
import sys
import time
import numpy as np
import torch
import torchaudio
sys.path.insert(0, r"d:/Ace-Step-Latest/StableAudio3")
from stable_audio_3.model import StableAudioModel
from stable_audio_3.inference.sampling import build_schedule
from stable_audio_3.data.utils import compute_effective_seq_len_from_conditioning
SR = 44100
DS = 4096 # latent downsampling ratio
CHUNK_LATENTS = 128 # SAME chunk graphs are traced at this size
CHUNK_SAMPLES = CHUNK_LATENTS * DS
OVERLAP = 32 # latent-frame overlap for tiling (pipeline default)
STEPS = 8
STRENGTH = 0.30
SEED = 1234
DURATION = 30.0
HEADROOM_SEC = 6.0
PROMPT = ("Instrumental punk rock with distorted electric guitars, driving drums "
"and punchy melodic bass. Clean modern production. Instrumental only, no vocals.")
def zero_stochastic_paths(ae):
ae.bottleneck.noise_regularize = False
for m in ae.modules():
if hasattr(m, "mask_noise"):
m.mask_noise = 0
def adapt_sample_size(seconds, encoder_chunk_size=32, encoder_stride=16):
"""Mirror of StableAudioModel._adapt_sample_size for the medium config."""
target = int((seconds + HEADROOM_SEC) * SR)
target = ((target + DS - 1) // DS) * DS
align = DS * (encoder_chunk_size // encoder_stride)
return ((target + align - 1) // align) * align
# --- ONNX tiling (ports of AudioAutoencoder.encode_audio / decode_audio) -----
def chunk_starts_for(total, size, hop):
starts = list(range(0, total - size + 1, hop))
if starts[-1] != total - size:
starts.append(total - size)
return starts
def ort_encode_tiled(sess, audio):
"""audio: np [1,2,T_samples] (T multiple of DS) -> latents np [1,256,T//DS]."""
total_latents = audio.shape[-1] // DS
if total_latents <= CHUNK_LATENTS:
raise ValueError("clip shorter than one chunk — pad first")
hop = (CHUNK_LATENTS - OVERLAP) * DS
starts = chunk_starts_for(audio.shape[-1], CHUNK_SAMPLES, hop)
out = np.zeros((1, 256, total_latents), dtype=np.float32)
half = OVERLAP // 2
n = len(starts)
for i, s in enumerate(starts):
chunk = sess.run(None, {"audio": audio[..., s:s + CHUNK_SAMPLES]})[0]
first, last = i == 0, i == n - 1
os_ = (total_latents - CHUNK_LATENTS) if last else s // DS
left = 0 if first else half
right = CHUNK_LATENTS if last else CHUNK_LATENTS - half
out[..., os_ + left:os_ + right] = chunk[..., left:right]
return out
def ort_decode_tiled(sess, latents):
"""latents: np [1,256,L] -> audio np [1,2,L*DS]."""
total_latents = latents.shape[-1]
hop = CHUNK_LATENTS - OVERLAP
starts = chunk_starts_for(total_latents, CHUNK_LATENTS, hop)
out = np.zeros((1, 2, total_latents * DS), dtype=np.float32)
half_s = (OVERLAP // 2) * DS
n = len(starts)
for i, s in enumerate(starts):
chunk = sess.run(None, {"latents": latents[..., s:s + CHUNK_LATENTS]})[0]
first, last = i == 0, i == n - 1
os_ = (total_latents - CHUNK_LATENTS) * DS if last else s * DS
left = 0 if first else half_s
right = CHUNK_SAMPLES if last else CHUNK_SAMPLES - half_s
out[..., os_ + left:os_ + right] = chunk[..., left:right]
return out
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--onnx-dir", required=True)
ap.add_argument("--audio", default=r"d:/Ace-Step-Latest/sa3-refined/last-call-vfw_inst-original.wav")
ap.add_argument("--zero-noise", action="store_true",
help="Zero all noise (harness only, skips PyTorch ref) — C++ validation mode")
ap.add_argument("--pad-latents", type=int, default=0,
help="Override latent length (match C++ SA3_T_BUCKET padding)")
ap.add_argument("--dump-tokens", action="store_true",
help="Print the padded token ids csv + count for the C++ endpoint")
args = ap.parse_args()
clip, in_sr = torchaudio.load(args.audio)
assert in_sr == SR
clip = clip[:, : int(DURATION * SR)]
audio_sample_size = adapt_sample_size(DURATION)
latent_size = audio_sample_size // DS
if args.pad_latents > 0:
latent_size = args.pad_latents
audio_sample_size = latent_size * DS
conditioning = [{"prompt": PROMPT, "seconds_total": DURATION}]
print(f"audio_sample_size={audio_sample_size} latent_size={latent_size}")
# ---------------- Reference (PyTorch) ----------------
print("Reference: loading medium fp32...")
model = StableAudioModel.from_pretrained("medium", model_half=False)
zero_stochastic_paths(model.model.pretransform.model)
if args.dump_tokens:
tok = model.model.conditioner.conditioners["prompt"].tokenizer
enc = tok([PROMPT], truncation=True, max_length=256, padding="max_length",
return_tensors="np")
ids = enc["input_ids"][0].tolist()
n_real = int(enc["attention_mask"][0].sum())
print("TOKENS_CSV=" + ",".join(str(i) for i in ids))
print(f"N_TOKENS={n_real}")
return 0
if not args.zero_noise:
t0 = time.time()
ref = model.generate(
prompt=PROMPT, duration=DURATION, steps=STEPS, cfg_scale=1.0, seed=SEED,
sample_size=model.model_config["sample_size"],
init_audio=(SR, clip), init_noise_level=STRENGTH,
sampler_type="euler",
)[0].cpu()
print(f"Reference done ({time.time()-t0:.0f}s)")
# ---------------- Harness (ONNX only) ----------------
import onnxruntime as ort
sess_opt = ort.SessionOptions()
load = lambda n: ort.InferenceSession(os.path.join(args.onnx_dir, n),
sess_opt, providers=["CPUExecutionProvider"])
print("Harness: loading ONNX sessions...")
s_text = load("sa3-text_encoder.onnx")
s_sec = load("sa3-seconds_embedder.onnx")
s_enc = load("sa3-same_encoder.onnx")
s_dit = load("sa3-dit.onnx")
s_dec = load("sa3-same_decoder.onnx")
t0 = time.time()
# Conditioning
tok = model.model.conditioner.conditioners["prompt"].tokenizer
enc = tok([PROMPT], truncation=True, max_length=256, padding="max_length",
return_tensors="np")
text_emb = s_text.run(None, {"input_ids": enc["input_ids"].astype(np.int64),
"attention_mask": enc["attention_mask"].astype(np.bool_)})[0]
sec_emb = s_sec.run(None, {"seconds": np.array([DURATION], dtype=np.float32)})[0]
cross = np.concatenate([text_emb, sec_emb[:, None, :]], axis=1) # [1,257,768]
# Init latents: pad clip to adapted size, tiled ONNX encode
padded = torch.zeros(1, 2, audio_sample_size)
padded[0, :, : clip.shape[-1]] = clip
init_latents = ort_encode_tiled(s_enc, padded.numpy().astype(np.float32))
# Noise: replicate generate() exactly — manual_seed then randn on CUDA
if args.zero_noise:
noise = np.zeros((1, 256, latent_size), dtype=np.float32)
else:
torch.manual_seed(SEED)
noise = torch.randn([1, 256, latent_size], device="cuda").cpu().numpy()
x = init_latents * (1 - STRENGTH) + noise * STRENGTH
# Schedule + padding mask (repo helpers = same math as reference)
eff = compute_effective_seq_len_from_conditioning(conditioning, SR, DS, "cpu")
sigmas = build_schedule(
steps=STEPS, sigma_max=STRENGTH,
dist_shift=model.model.sampling_dist_shift,
effective_seq_len=eff, fallback_seq_len=latent_size,
include_endpoint=True, device="cpu",
).numpy().astype(np.float32).reshape(-1)
headroom_tokens = int(HEADROOM_SEC * SR / DS)
valid = min(int(eff.item()) + headroom_tokens, latent_size)
padding_mask = np.zeros((1, latent_size), dtype=np.bool_)
padding_mask[:, :valid] = True
local_add = np.zeros((1, 257, latent_size), dtype=np.float32) # no inpaint
glob = sec_emb.astype(np.float32)
# Euler loop
for i in range(STEPS):
t_curr, t_next = sigmas[i], sigmas[i + 1]
v = s_dit.run(None, {
"x": x.astype(np.float32),
"t": np.array([t_curr], dtype=np.float32),
"cross_attn_cond": cross.astype(np.float32),
"global_embed": glob,
"local_add_cond": local_add,
"padding_mask": padding_mask,
})[0]
x = x + (t_next - t_curr) * v
print(f" step {i+1}/{STEPS} t={t_curr:.4f}->{t_next:.4f}")
# Decode + padding zeroing + trim (mirrors sample_diffusion tail + generate)
audio = ort_decode_tiled(s_dec, x.astype(np.float32))
audio_mask = np.repeat(padding_mask, DS, axis=-1)[:, : audio.shape[-1]]
audio = audio * audio_mask[:, None, :]
audio = np.clip(audio, -1, 1)[0, :, : int(DURATION * SR)]
print(f"Harness done ({time.time()-t0:.0f}s)")
if args.zero_noise:
out_dir = os.path.dirname(args.onnx_dir)
path = os.path.join(out_dir, "e2e_ort_zeronoise.wav")
torchaudio.save(path, torch.tensor(audio), SR)
print(f"Zero-noise harness output -> {path}")
return 0
# ---------------- Compare ----------------
ref_np = ref.numpy()[:, : int(DURATION * SR)]
n = min(ref_np.shape[-1], audio.shape[-1])
a, b = ref_np[..., :n].ravel(), audio[..., :n].ravel()
cos = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-12))
print(f"E2E cosine={cos:.6f} max_abs_diff={np.abs(a-b).max():.3e}")
out_dir = os.path.dirname(args.onnx_dir)
torchaudio.save(os.path.join(out_dir, "e2e_ref.wav"), torch.tensor(ref_np), SR)
torchaudio.save(os.path.join(out_dir, "e2e_ort.wav"), torch.tensor(audio), SR)
print("E2E OK" if cos > 0.99 else "E2E FAILED")
return 0 if cos > 0.99 else 1
if __name__ == "__main__":
sys.exit(main())
tools/onnx-export/export_cond_enc.py
+524
View File
@@ -0,0 +1,524 @@
#!/usr/bin/env python3
"""
export_cond_enc.py — Export AceStep Condition Encoder to ONNX.
The condition encoder takes outputs from the text encoder (text_hidden + lyric_embed)
and reference audio features (timbre_feats), and produces enc_hidden for DiT cross-attention.
Internal architecture:
- text_projector: Linear(1024→2048, no bias) — projects text encoder output
- lyric_encoder: Linear(1024→2048)+bias → 8-layer bidirectional Qwen3 → RMSNorm
- timbre_encoder: Linear(64→2048)+bias → [CLS prepend] → 4-layer bidir Qwen3 → RMSNorm → position[0]
- cat(lyric_out, timbre_out, text_proj_out) → enc_hidden [B, S_total, 2048]
Usage:
python export_cond_enc.py --model-dir <path-to-DiT-safetensors> --output <output.onnx>
"""
import argparse
import os
import sys
import time
import struct
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
class CondEncoderWrapper(nn.Module):
"""Wrapper for ONNX export that simplifies the condition encoder interface.
For inference (batch_size=1), we simplify:
- No pack_sequences sorting (all tokens are valid, no padding)
- Timbre: single reference, so unpack is trivial (just unsqueeze)
- Output is simple cat(lyric, timbre, text_proj)
ONNX inputs:
text_hidden: [B, S_text, 1024] fp16 — from text encoder
lyric_embed: [B, S_lyric, 1024] fp16 — from embedding table lookup
timbre_feats: [B, S_ref, 64] fp16 — from VAE encoder (or zeros)
has_timbre: [1] int64 — 1 if timbre is present, 0 if not
ONNX output:
enc_hidden: [B, S_total, 2048] fp16 — packed conditioning
"""
def __init__(self, cond_encoder):
super().__init__()
self.text_projector = cond_encoder.text_projector
self.lyric_encoder = cond_encoder.lyric_encoder
self.timbre_encoder = cond_encoder.timbre_encoder
def forward(self, text_hidden, lyric_embed, timbre_feats, has_timbre):
"""
Forward pass with simplified interface for ONNX export.
Note: For ONNX tracing, has_timbre must be a tensor, not a Python bool.
We use torch.where / masking to handle the conditional timbre path.
"""
B = text_hidden.shape[0]
# 1) Text projection: [B, S_text, 1024] → [B, S_text, 2048]
text_proj = self.text_projector(text_hidden)
# 2) Lyric encoding: [B, S_lyric, 1024] → 8L bidir Qwen3 → [B, S_lyric, 2048]
S_lyric = lyric_embed.shape[1]
lyric_mask = torch.ones(B, S_lyric, device=lyric_embed.device, dtype=torch.long)
lyric_out = self.lyric_encoder(
inputs_embeds=lyric_embed,
attention_mask=lyric_mask,
)
if hasattr(lyric_out, 'last_hidden_state'):
lyric_out = lyric_out.last_hidden_state
else:
lyric_out = lyric_out[0]
# 3) Timbre encoding: [B, S_ref, 64] → 4L bidir Qwen3 → position[0] → [B, 1, 2048]
# For ONNX: we always run the timbre path but zero out if has_timbre=0
S_ref = timbre_feats.shape[1]
timbre_mask = torch.ones(1, S_ref, device=timbre_feats.device, dtype=torch.long)
# refer_audio_order_mask: all 0s means everything belongs to batch 0
order_mask = torch.zeros(1, device=timbre_feats.device, dtype=torch.long)
# Reshape for timbre encoder: expects [N_packed, S_ref, 64]
timbre_input = timbre_feats # [B, S_ref, 64]
timbre_embs, timbre_embs_mask = self.timbre_encoder(
refer_audio_acoustic_hidden_states_packed=timbre_input,
refer_audio_order_mask=order_mask,
attention_mask=timbre_mask,
)
# timbre_embs: [B, 1, 2048] — CLS token output per batch
# 4) Concatenate: [lyric, timbre, text_proj]
# When has_timbre=0, skip timbre in the concatenation
# For ONNX compatibility, always cat but mask the timbre contribution
ht = has_timbre[0]
if ht > 0:
enc_hidden = torch.cat([lyric_out, timbre_embs, text_proj], dim=1)
else:
enc_hidden = torch.cat([lyric_out, text_proj], dim=1)
return enc_hidden
class CondEncoderWrapperFixed(nn.Module):
"""Fixed version that always includes timbre (simplifies ONNX graph).
For inference, timbre is always present (silence latent as zero timbre).
This avoids dynamic control flow in the ONNX graph.
IMPORTANT: The timbre encoder's forward() uses unpack_timbre_embeddings()
which has data-dependent control flow (refer_audio_order_mask.max().item()).
torch.export cannot handle this. So we manually invoke the timbre encoder's
sub-components: embed_tokens → CLS prepend → transformer layers → norm →
take position 0. This is equivalent for B=1 inference.
ONNX inputs:
text_hidden: [B, S_text, 1024] fp16
lyric_embed: [B, S_lyric, 1024] fp16
timbre_feats: [B, S_ref, 64] fp16 (zeros if no reference)
ONNX output:
enc_hidden: [B, S_total, 2048] fp16 where S_total = S_lyric + 1 + S_text
"""
def __init__(self, cond_encoder):
super().__init__()
self.text_projector = cond_encoder.text_projector
self.lyric_encoder = cond_encoder.lyric_encoder
# Extract timbre encoder sub-components for manual invocation
self.timbre_embed_tokens = cond_encoder.timbre_encoder.embed_tokens
self.timbre_special_token = cond_encoder.timbre_encoder.special_token
self.timbre_norm = cond_encoder.timbre_encoder.norm
self.timbre_rotary_emb = cond_encoder.timbre_encoder.rotary_emb
self.timbre_layers = cond_encoder.timbre_encoder.layers
self.timbre_config = cond_encoder.timbre_encoder.config
def _timbre_forward_simple(self, timbre_feats):
"""Run the timbre encoder without unpack_timbre_embeddings.
timbre_feats: [B, S_ref, 64]
Returns: [B, 1, hidden_size] — CLS token output
"""
B = timbre_feats.shape[0]
# Project: [B, S_ref, 64] → [B, S_ref, hidden_size]
inputs_embeds = self.timbre_embed_tokens(timbre_feats)
# Prepend CLS token: [B, S_ref+1, hidden_size]
cls_token = self.timbre_special_token.expand(B, 1, -1)
inputs_embeds = torch.cat([cls_token, inputs_embeds], dim=1)
S = inputs_embeds.shape[1]
# Position IDs and RoPE
cache_position = torch.arange(0, S, device=inputs_embeds.device)
position_ids = cache_position.unsqueeze(0)
position_embeddings = self.timbre_rotary_emb(inputs_embeds, position_ids)
# Build attention mask (full bidirectional, no padding)
# Using None for SDPA = no mask = full attention
hidden_states = inputs_embeds
for layer in self.timbre_layers:
layer_outputs = layer(
hidden_states,
position_embeddings,
None, # attention_mask=None → full bidirectional
position_ids,
)
hidden_states = layer_outputs[0]
hidden_states = self.timbre_norm(hidden_states)
# Extract CLS token (position 0): [B, hidden_size]
timbre_emb = hidden_states[:, 0:1, :] # [B, 1, hidden_size]
return timbre_emb
def forward(self, text_hidden, lyric_embed, timbre_feats):
B = text_hidden.shape[0]
# 1) Text projection
text_proj = self.text_projector(text_hidden)
# 2) Lyric encoding
S_lyric = lyric_embed.shape[1]
lyric_mask = torch.ones(B, S_lyric, device=lyric_embed.device, dtype=torch.long)
lyric_out = self.lyric_encoder(
inputs_embeds=lyric_embed,
attention_mask=lyric_mask,
)
if hasattr(lyric_out, 'last_hidden_state'):
lyric_out = lyric_out.last_hidden_state
else:
lyric_out = lyric_out[0]
# 3) Timbre encoding — manual path (bypasses unpack_timbre_embeddings)
timbre_embs = self._timbre_forward_simple(timbre_feats)
# timbre_embs: [B, 1, 2048]
# 4) Concatenate: lyric + timbre + text_proj
enc_hidden = torch.cat([lyric_out, timbre_embs, text_proj], dim=1)
return enc_hidden
def load_model(model_dir: str, device: str = "cuda", dtype=torch.float32):
"""Load the AceStep model and extract the condition encoder."""
model_dir = Path(model_dir)
if sys.platform == "win32":
import io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
# Monkey-patch transformers auto_docstring
try:
import transformers.utils.auto_docstring as _ad
_ad.auto_docstring = lambda *a, **kw: (lambda cls: cls)
except Exception:
pass
sys.path.insert(0, str(model_dir))
# Create a stub AceStepConfig module to bypass the 'acestep' package import.
# The real AceStepConfig is a PretrainedConfig subclass. We construct it
# using AutoConfig which reads config.json and finds the auto_map.
# But first we need the config module to exist so the model code can import it.
import json
import types
from transformers import PretrainedConfig
with open(model_dir / "config.json") as f:
config_dict = json.load(f)
# Create AceStepConfig class dynamically from config.json
class AceStepConfig(PretrainedConfig):
model_type = "acestep"
def __init__(self, **kwargs):
super().__init__(**kwargs)
# Set all config keys as attributes
for k, v in kwargs.items():
if not hasattr(self, k):
setattr(self, k, v)
# Ensure critical attrs have defaults
if not hasattr(self, 'text_hidden_dim'):
self.text_hidden_dim = 1024
if not hasattr(self, 'timbre_hidden_dim'):
self.timbre_hidden_dim = 64
if not hasattr(self, 'encoder_hidden_size'):
self.encoder_hidden_size = 2048
if not hasattr(self, 'encoder_intermediate_size'):
self.encoder_intermediate_size = 6144
if not hasattr(self, 'encoder_num_attention_heads'):
self.encoder_num_attention_heads = 16
if not hasattr(self, 'encoder_num_key_value_heads'):
self.encoder_num_key_value_heads = 8
if not hasattr(self, 'num_lyric_encoder_hidden_layers'):
self.num_lyric_encoder_hidden_layers = 8
if not hasattr(self, 'num_timbre_encoder_hidden_layers'):
self.num_timbre_encoder_hidden_layers = 4
if not hasattr(self, 'num_attention_pooler_hidden_layers'):
self.num_attention_pooler_hidden_layers = 2
if not hasattr(self, 'out_channels'):
self.out_channels = 64
if not hasattr(self, 'in_channels'):
self.in_channels = 192
# Register the stub config module
stub_mod = types.ModuleType("configuration_acestep_v15")
stub_mod.AceStepConfig = AceStepConfig
sys.modules["configuration_acestep_v15"] = stub_mod
# Create stub acestep package hierarchy. Each intermediate module needs
# __path__ set so it acts as a package (allows submodule imports).
for mod_name in ["acestep", "acestep.models", "acestep.models.common"]:
m = types.ModuleType(mod_name)
m.__path__ = [] # makes it act as a package
sys.modules[mod_name] = m
# Register configuration_acestep_v15 under the acestep.models.common path
cfg_mod = types.ModuleType("acestep.models.common.configuration_acestep_v15")
cfg_mod.AceStepConfig = AceStepConfig
sys.modules["acestep.models.common.configuration_acestep_v15"] = cfg_mod
# Create apg_guidance stub — these functions are used by the DiT diffusion
# loop but NOT by the condition encoder. Provide dummies to satisfy import.
class MomentumBuffer:
def __init__(self, *a, **kw): pass
def _apg_stub(*a, **kw): return None
apg_stub = types.ModuleType("acestep.models.common.apg_guidance")
apg_stub.MomentumBuffer = MomentumBuffer
apg_stub.adg_forward = _apg_stub
apg_stub.adg_w_norm_forward = _apg_stub
apg_stub.adg_wo_clip_forward = _apg_stub
apg_stub.apg_forward = _apg_stub
apg_stub.cfg_forward = _apg_stub
apg_stub.call_cos_tensor = _apg_stub
apg_stub.compute_perpendicular_component = _apg_stub
apg_stub.project = _apg_stub
sys.modules["acestep.models.common.apg_guidance"] = apg_stub
# Also register the local apg_guidance module
apg_local = types.ModuleType("apg_guidance")
apg_local.MomentumBuffer = MomentumBuffer
apg_local.adg_forward = _apg_stub
apg_local.apg_forward = _apg_stub
apg_local.cfg_forward = _apg_stub
sys.modules["apg_guidance"] = apg_local
config = AceStepConfig(**config_dict)
config._attn_implementation = "sdpa"
# The encoder's Qwen3 sub-models (lyric/timbre) use encoder_hidden_size
# as their hidden_size. Set it on the config so Qwen3RotaryEmbedding works.
# transformers 5.x requires rope_parameters dict.
if not hasattr(config, 'rope_parameters') or config.rope_parameters is None:
config.rope_parameters = {
"rope_type": "default",
"rope_theta": config.rope_theta if hasattr(config, 'rope_theta') else 1000000.0,
}
print(f"[export_cond_enc] Loading model from {model_dir}...")
t0 = time.time()
# The full model creates a separate encoder config with encoder-specific
# dimensions (see AceStepConditionGenerationModel.__init__ lines 1621-1628).
# The encoder uses encoder_hidden_size (2048), not the DiT hidden_size (2560).
import copy
encoder_config = copy.deepcopy(config)
encoder_config.hidden_size = config.encoder_hidden_size
encoder_config.intermediate_size = config.encoder_intermediate_size
encoder_config.num_attention_heads = config.encoder_num_attention_heads
encoder_config.num_key_value_heads = config.encoder_num_key_value_heads
from modeling_acestep_v15_xl_base import AceStepConditionEncoder
cond_encoder = AceStepConditionEncoder(encoder_config)
# Load weights — filter to encoder.* prefix
from safetensors.torch import load_file
st_path = model_dir / "model.safetensors"
state_dict = load_file(str(st_path))
cond_state_dict = {}
for k, v in state_dict.items():
if k.startswith("encoder."):
cond_state_dict[k[len("encoder."):]] = v
missing, unexpected = cond_encoder.load_state_dict(cond_state_dict, strict=False)
if missing:
print(f"[export_cond_enc] Warning: {len(missing)} missing keys (first 5: {missing[:5]})")
if unexpected:
print(f"[export_cond_enc] Warning: {len(unexpected)} unexpected keys (first 5: {unexpected[:5]})")
cond_encoder = cond_encoder.to(device=device, dtype=dtype)
cond_encoder.eval()
t1 = time.time()
n_params = sum(p.numel() for p in cond_encoder.parameters()) / 1e6
print(f"[export_cond_enc] Model loaded in {t1-t0:.1f}s ({n_params:.0f}M params)")
print(f"[export_cond_enc] text_hidden_dim={encoder_config.text_hidden_dim}, hidden_size={encoder_config.hidden_size}")
return cond_encoder, encoder_config
def export_onnx(cond_encoder, config, output_path: str, opset: int = 18):
"""Export the condition encoder to ONNX."""
device = next(cond_encoder.parameters()).device
dtype = next(cond_encoder.parameters()).dtype
wrapper = CondEncoderWrapperFixed(cond_encoder)
wrapper.eval()
# Dummy inputs
B = 1
S_text = 64
S_lyric = 128
S_ref = 8 # 8 frames of reference audio (short clip)
dummy_text_hidden = torch.randn(B, S_text, config.text_hidden_dim, device=device, dtype=dtype)
dummy_lyric_embed = torch.randn(B, S_lyric, config.text_hidden_dim, device=device, dtype=dtype)
dummy_timbre_feats = torch.randn(B, S_ref, config.timbre_hidden_dim, device=device, dtype=dtype)
print(f"[export_cond_enc] Tracing with shapes: text={list(dummy_text_hidden.shape)}, "
f"lyric={list(dummy_lyric_embed.shape)}, timbre={list(dummy_timbre_feats.shape)}")
# Test forward
print("[export_cond_enc] Testing forward pass...")
with torch.no_grad():
test_out = wrapper(dummy_text_hidden, dummy_lyric_embed, dummy_timbre_feats)
expected_S = S_lyric + 1 + S_text # lyric + timbre(1) + text
print(f"[export_cond_enc] Output shape: {list(test_out.shape)} "
f"(expected [{B}, {expected_S}, {config.hidden_size}])")
# Export
print(f"[export_cond_enc] Exporting to ONNX (opset {opset})...")
t0 = time.time()
torch.onnx.export(
wrapper,
(dummy_text_hidden, dummy_lyric_embed, dummy_timbre_feats),
output_path,
opset_version=opset,
input_names=["text_hidden", "lyric_embed", "timbre_feats"],
output_names=["enc_hidden"],
dynamic_axes={
"text_hidden": {0: "batch", 1: "text_seq"},
"lyric_embed": {0: "batch", 1: "lyric_seq"},
"timbre_feats": {0: "batch", 1: "timbre_seq"},
"enc_hidden": {0: "batch", 1: "enc_seq"},
},
do_constant_folding=True,
export_params=True,
)
t1 = time.time()
file_size = os.path.getsize(output_path)
print(f"[export_cond_enc] Exported to {output_path}")
print(f"[export_cond_enc] File size: {file_size/1e6:.1f} MB")
print(f"[export_cond_enc] Export time: {t1-t0:.1f}s")
def export_null_cond_emb(model_dir: str, output_path: str):
"""Export null_condition_emb as raw float32 binary."""
from safetensors.torch import load_file
model_dir = Path(model_dir)
st_path = model_dir / "model.safetensors"
state_dict = load_file(str(st_path))
key = "null_condition_emb"
if key not in state_dict:
print(f"[export_cond_enc] WARNING: {key} not found, skipping")
return
vec = state_dict[key].detach().cpu().float().numpy().flatten()
with open(output_path, "wb") as f:
f.write(struct.pack("<I", len(vec)))
f.write(vec.tobytes())
print(f"[export_cond_enc] null_condition_emb: [{len(vec)}] -> {output_path} ({len(vec)*4} bytes)")
def verify_onnx(onnx_path: str, cond_encoder, config):
"""Verify ONNX output matches PyTorch."""
try:
import onnxruntime as ort
except ImportError:
print("[export_cond_enc] onnxruntime not installed, skipping verification")
return
device = next(cond_encoder.parameters()).device
dtype = next(cond_encoder.parameters()).dtype
wrapper = CondEncoderWrapperFixed(cond_encoder)
wrapper.eval()
B, S_text, S_lyric, S_ref = 1, 32, 64, 8
text_hidden = torch.randn(B, S_text, config.text_hidden_dim, device=device, dtype=dtype)
lyric_embed = torch.randn(B, S_lyric, config.text_hidden_dim, device=device, dtype=dtype)
timbre_feats = torch.randn(B, S_ref, config.timbre_hidden_dim, device=device, dtype=dtype)
with torch.no_grad():
ref_out = wrapper(text_hidden, lyric_embed, timbre_feats).cpu().float().numpy()
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
sess = ort.InferenceSession(onnx_path, providers=providers)
ort_out = sess.run(None, {
"text_hidden": text_hidden.cpu().float().numpy(),
"lyric_embed": lyric_embed.cpu().float().numpy(),
"timbre_feats": timbre_feats.cpu().float().numpy(),
})[0]
max_diff = np.max(np.abs(ref_out - ort_out))
mean_diff = np.mean(np.abs(ref_out - ort_out))
print(f"[export_cond_enc] Verification: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}")
if max_diff < 0.05:
print("[export_cond_enc] PASS: ONNX output matches PyTorch")
else:
print("[export_cond_enc] WARNING: Large difference — may need investigation")
def main():
parser = argparse.ArgumentParser(description="Export AceStep condition encoder to ONNX")
parser.add_argument("--model-dir", required=True,
help="Path to the DiT model directory (contains encoder weights)")
parser.add_argument("--output", default=None,
help="Output ONNX file (default: models/onnx/cond_encoder.onnx)")
parser.add_argument("--opset", type=int, default=18)
parser.add_argument("--verify", action="store_true")
parser.add_argument("--device", default="cuda")
args = parser.parse_args()
if args.output is None:
onnx_dir = Path(args.model_dir).parent / "onnx"
onnx_dir.mkdir(parents=True, exist_ok=True)
args.output = str(onnx_dir / "cond_encoder.onnx")
os.makedirs(os.path.dirname(args.output), exist_ok=True)
output_dir = os.path.dirname(args.output)
# Load model
cond_encoder, config = load_model(args.model_dir, device=args.device)
# Export ONNX
export_onnx(cond_encoder, config, args.output, opset=args.opset)
# Export null_condition_emb
null_cond_path = os.path.join(output_dir, "null_condition_emb.bin")
export_null_cond_emb(args.model_dir, null_cond_path)
# Verify
if args.verify:
verify_onnx(args.output, cond_encoder, config)
print("[export_cond_enc] Done!")
if __name__ == "__main__":
main()
tools/onnx-export/export_dit.py
+760
View File
@@ -0,0 +1,760 @@
#!/usr/bin/env python3
"""
export_dit.py — Export AceStep DiT forward pass to ONNX for TensorRT acceleration.
Exports the SINGLE FORWARD PASS (one diffusion timestep) of the DiT model,
wrapping the full 32-layer transformer + attention mask computation + RoPE
into a single ONNX graph with 4 simplified inputs.
Precision recipes (--precision):
fp32 — Full FP32. Correct but slow. Baseline for validation.
bf16_mixed — (default for XL) bf16 bulk + fp32 ConvTranspose1d island.
Used with TRT STRONGLY_TYPED mode. Demon-proven recipe.
bf16 has same exponent range as fp32 — no activation overflow.
Usage:
python export_dit.py --model-dir <path-to-safetensors-model> --output <output.onnx>
python export_dit.py --model-dir <path> --output <path> --precision bf16_mixed
The diffusion loop, guidance (APG/CFG), and solvers stay in C++.
TRT compiles the ONNX graph once; LoRA adapters use IRefitter weight swapping.
"""
import argparse
import sys
import os
import time
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F
# We need the model's own code
# The model dir contains modeling_acestep_v15_xl_base.py
class _Fp32CastWrapper(nn.Module):
"""Run an inner module in fp32, casting around it.
Used when TRT has no kernel for a specific op shape in bf16.
The wrapper casts input to fp32, runs the inner module, then casts
output back to the caller's dtype.
"""
def __init__(self, inner: nn.Module):
super().__init__()
inner.float() # force inner weights to fp32
self.inner = inner
def forward(self, x: torch.Tensor) -> torch.Tensor:
out_dtype = x.dtype
# Disable autocast — without this, the outer autocast(bf16) overrides
# our explicit fp32 computation and TRT sees bf16 weights.
with torch.amp.autocast('cuda', enabled=False):
return self.inner(x.float()).to(out_dtype)
class PatchEmbedLinear(nn.Module):
"""Replace Conv1d(C_in, C_out, K, stride=K) with reshape + Linear.
TRT 10.16 has NO kernels for 1D convolutions with patch_size shapes
in any precision mode (fp16, bf16, or fp32). This is mathematically equivalent:
Conv1d: input[B, C_in, T] → output[B, C_out, T//K]
Linear: input[B, C_in, T] → unfold[B, T//K, C_in*K] → Linear → [B, C_out, T//K]
"""
def __init__(self, conv: nn.Conv1d):
super().__init__()
C_out, C_in, K = conv.weight.shape
self.kernel_size = K
self.linear = nn.Linear(C_in * K, C_out, bias=conv.bias is not None)
# Conv weight [C_out, C_in, K] → Linear weight [C_out, C_in*K]
self.linear.weight.data = conv.weight.data.reshape(C_out, -1).clone()
if conv.bias is not None:
self.linear.bias.data = conv.bias.data.clone()
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x: [B, C_in, T] (from Lambda transpose in proj_in)
B, C, T = x.shape
K = self.kernel_size
# Unfold patches: [B, C, T] → [B, T//K, C*K]
x = x.reshape(B, C, T // K, K) # [B, C, T//K, K]
x = x.permute(0, 2, 1, 3) # [B, T//K, C, K]
x = x.reshape(B, T // K, C * K) # [B, T//K, C*K]
out = self.linear(x) # [B, T//K, C_out]
return out.transpose(1, 2) # [B, C_out, T//K]
class UnPatchLinear(nn.Module):
"""Replace ConvTranspose1d(C_in, C_out, K, stride=K) with Linear + reshape.
TRT 10.16 has NO kernels for 1D transposed convolutions with patch_size shapes.
This is mathematically equivalent:
ConvTranspose1d: input[B, C_in, T//K] → output[B, C_out, T]
Linear: input[B, T//K, C_in] → Linear → [B, T//K, C_out*K] → fold → [B, C_out, T]
"""
def __init__(self, deconv: nn.ConvTranspose1d):
super().__init__()
C_in, C_out, K = deconv.weight.shape
self.kernel_size = K
self.C_out = C_out
self.linear = nn.Linear(C_in, C_out * K, bias=deconv.bias is not None)
# ConvTranspose1d weight [C_in, C_out, K] → Linear weight [C_out*K, C_in]
self.linear.weight.data = deconv.weight.data.permute(1, 2, 0).reshape(C_out * K, C_in).clone()
if deconv.bias is not None:
# ConvTranspose1d bias [C_out] → Linear bias [C_out*K] (repeat per patch)
self.linear.bias.data = deconv.bias.data.repeat_interleave(K).clone()
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x: [B, C_in, T//K] (from Lambda transpose in proj_out)
B, C, T_small = x.shape
K = self.kernel_size
x = x.transpose(1, 2) # [B, T//K, C_in]
x = self.linear(x) # [B, T//K, C_out*K]
x = x.reshape(B, T_small, self.C_out, K) # [B, T//K, C_out, K]
x = x.permute(0, 2, 1, 3) # [B, C_out, T//K, K]
x = x.reshape(B, self.C_out, T_small * K) # [B, C_out, T]
return x
class DiTForwardWrapper(nn.Module):
"""
Wrapper around AceStepDiTModel.forward() that simplifies the interface
for ONNX export.
ONNX inputs (4 total):
input_latents: [B, T, 192] — pre-concatenated [context_latents, xt]
enc_hidden: [B, S, 2048] — encoder hidden states
t: [B] fp32 — current timestep
t_r: [B] fp32 — reference timestep
ONNX output:
velocity: [B, T, 64] — predicted flow velocity
Masks and position IDs are computed internally from T and S.
"""
def __init__(self, dit_model, precision="bf16_mixed"):
super().__init__()
self.dit = dit_model
self.config = dit_model.config
self.precision = precision
def forward(self, input_latents, enc_hidden, t, t_r):
"""
Args:
input_latents: [B, T, 192] — concatenated context + noise latents
enc_hidden: [B, S, 2048] — encoder hidden states
t: [B] — timestep
t_r: [B] — reference timestep
Returns:
velocity: [B, T, 64] — predicted velocity
"""
B = input_latents.shape[0]
T = input_latents.shape[1]
# Split input_latents into context (128 dim) and noise (64 dim)
context_latents = input_latents[:, :, :128]
hidden_states = input_latents[:, :, 128:]
# bf16 autocast: the dynamo exporter decomposes complex ops
# (view_as_complex → rotate_half) into real-number equivalents,
# so no Cast(to=COMPLEX128) appears in the ONNX graph.
if self.precision == "bf16_mixed":
autocast_dtype = torch.bfloat16
else:
autocast_dtype = torch.float32
with torch.amp.autocast('cuda', dtype=autocast_dtype):
outputs = self.dit(
hidden_states=hidden_states,
timestep=t,
timestep_r=t_r,
attention_mask=None,
encoder_hidden_states=enc_hidden,
encoder_attention_mask=None,
context_latents=context_latents,
use_cache=False,
past_key_values=None,
output_attentions=False,
)
# outputs[0] is the velocity prediction [B, T, 64]
velocity = outputs[0]
return velocity
def apply_bf16_mixed(dit_model):
"""Apply the bf16_mixed precision recipe (XL models).
bf16 bulk + fp32 island for proj_out ConvTranspose1d.
bf16 has the SAME exponent range as fp32 (8 bits vs fp16's 5 bits),
so intermediate activations never overflow. This is the key difference
from fp16_mixed which NaN'd because the XL residual stream accumulated
values exceeding fp16's ±65504 range over 32 layers.
The entire model runs in bf16 EXCEPT:
- proj_out ConvTranspose1d → wrapped in _Fp32CastWrapper because
TRT 10.16 has no bf16 deconv kernel for this shape.
Uses STRONGLY_TYPED mode so TRT honors the bf16/fp32 split from the
ONNX graph. TRT's bf16 tensor cores provide the same throughput as fp16.
"""
dit_model.to(torch.bfloat16)
print("[export_dit] Applied bf16 bulk conversion")
# FP32 island: proj_out ConvTranspose1d (TRT has no bf16 deconv kernel)
# NOTE: This gets replaced by UnPatchLinear AFTER this function runs
# (replace_conv_with_linear handles it). But we still wrap it in
# _Fp32CastWrapper in case the Conv→Linear replacement changes.
if hasattr(dit_model, 'proj_out') and isinstance(dit_model.proj_out, nn.Sequential):
for i, mod in enumerate(dit_model.proj_out):
if isinstance(mod, nn.ConvTranspose1d):
dit_model.proj_out[i] = _Fp32CastWrapper(mod)
print(f"[export_dit] FP32 island: proj_out[{i}] ConvTranspose1d → _Fp32CastWrapper")
break
return dit_model
def replace_conv_with_linear(dit_model):
"""Replace Conv1d/ConvTranspose1d with equivalent Linear ops.
TRT 10.16 has NO kernels for 1D convolutions with patch_size=2 in ANY
precision mode (fp16, bf16, fp32, or mixed). PatchEmbedLinear/UnPatchLinear
reformulate these as reshape+matmul which TRT handles perfectly.
Must be called for ALL precision recipes, not just mixed precision.
Handles _Fp32CastWrapper: if a ConvTranspose1d is already wrapped in
_Fp32CastWrapper (from bf16_mixed recipe), we unwrap it, convert to
UnPatchLinear, and re-wrap in _Fp32CastWrapper.
"""
if hasattr(dit_model, 'proj_in') and isinstance(dit_model.proj_in, nn.Sequential):
for i, mod in enumerate(dit_model.proj_in):
if isinstance(mod, nn.Conv1d):
dit_model.proj_in[i] = PatchEmbedLinear(mod)
print(f"[export_dit] Conv→Linear: proj_in[{i}] Conv1d → PatchEmbedLinear")
if hasattr(dit_model, 'proj_out') and isinstance(dit_model.proj_out, nn.Sequential):
for i, mod in enumerate(dit_model.proj_out):
if isinstance(mod, nn.ConvTranspose1d):
dit_model.proj_out[i] = UnPatchLinear(mod)
print(f"[export_dit] Conv→Linear: proj_out[{i}] ConvTranspose1d → UnPatchLinear")
elif isinstance(mod, _Fp32CastWrapper) and isinstance(mod.inner, nn.ConvTranspose1d):
# Unwrap, convert, re-wrap
linear_mod = UnPatchLinear(mod.inner)
dit_model.proj_out[i] = _Fp32CastWrapper(linear_mod)
print(f"[export_dit] Conv→Linear: proj_out[{i}] Fp32Cast(ConvTranspose1d) → Fp32Cast(UnPatchLinear)")
return dit_model
def load_dit_model(model_dir: str, device: str = "cuda", precision: str = "bf16_mixed"):
"""Load the AceStepDiTModel from a safetensors checkpoint."""
model_dir = Path(model_dir)
# Fix Windows encoding issues with transformers emoji output
if sys.platform == "win32":
import io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
# Monkey-patch transformers auto_docstring to avoid lookup failure
# for custom model types not registered in HF model registry
try:
import transformers.utils.auto_docstring as _ad
_orig = _ad.auto_docstring
_ad.auto_docstring = lambda *a, **kw: (lambda cls: cls) # no-op decorator
except Exception:
pass
# Add model dir to sys.path so we can import the model code.
# Also add the Demon app root — model config files are re-export stubs
# that import from the acestep package (from Demon).
sys.path.insert(0, str(model_dir))
demon_root = Path(model_dir).resolve().parent.parent.parent / "Demon"
if demon_root.exists():
sys.path.insert(0, str(demon_root))
print(f"[export_dit] Added {demon_root} to sys.path for acestep package")
# The model config stubs reference acestep.models.common but the
# actual module is acestep.models. Create a shim alias.
try:
import acestep.models as _am
sys.modules["acestep.models.common"] = _am
# Also create the subpackage entry so Python's import system is happy
import types
if not hasattr(_am, "common"):
_am.common = _am
except ImportError:
print("[export_dit] WARNING: Could not import acestep.models")
# Auto-detect the modeling module — different model variants use different
# filenames (modeling_acestep_v15_xl_base.py, xl_turbo.py, etc.)
import glob
modeling_files = glob.glob(str(model_dir / "modeling_acestep_v15*.py"))
if not modeling_files:
print(f"[export_dit] ERROR: No modeling_acestep_v15*.py found in {model_dir}")
sys.exit(1)
modeling_module = Path(modeling_files[0]).stem
print(f"[export_dit] Using modeling module: {modeling_module}")
import importlib
mod = importlib.import_module(modeling_module)
AceStepDiTModel = mod.AceStepDiTModel
from configuration_acestep_v15 import AceStepConfig
# Load config
import json
with open(model_dir / "config.json") as f:
config_dict = json.load(f)
config = AceStepConfig(**config_dict)
# Force SDPA for ONNX export (no flash attention)
config._attn_implementation = "sdpa"
print(f"[export_dit] Loading model from {model_dir}...")
print(f"[export_dit] Precision recipe: {precision}")
t0 = time.time()
# Create just the DiT model (decoder) — no need for full model
dit_model = AceStepDiTModel(config)
# Load weights — handle both single-file and sharded safetensors
from safetensors.torch import load_file
index_path = model_dir / "model.safetensors.index.json"
single_path = model_dir / "model.safetensors"
if index_path.exists():
# Sharded: load index to find all shard files
import json as _json
with open(index_path) as f:
index = _json.load(f)
shard_files = sorted(set(index["weight_map"].values()))
print(f"[export_dit] Loading {len(shard_files)} shards...")
state_dict = {}
for shard in shard_files:
shard_path = model_dir / shard
print(f"[export_dit] Loading {shard}...")
state_dict.update(load_file(str(shard_path)))
elif single_path.exists():
state_dict = load_file(str(single_path))
else:
print(f"[export_dit] ERROR: No model.safetensors found in {model_dir}")
sys.exit(1)
# Filter and remap: "decoder.X" -> "X" for the DiT model
dit_state_dict = {}
for k, v in state_dict.items():
if k.startswith("decoder."):
dit_state_dict[k[len("decoder."):]] = v
missing, unexpected = dit_model.load_state_dict(dit_state_dict, strict=False)
if missing:
print(f"[export_dit] Warning: {len(missing)} missing keys (first 5: {missing[:5]})")
if unexpected:
print(f"[export_dit] Warning: {len(unexpected)} unexpected keys")
# Apply precision recipe AFTER loading weights (so weights are converted correctly)
if precision == "bf16_mixed":
dit_model = dit_model.to(device=device) # move to GPU first
dit_model = apply_bf16_mixed(dit_model)
elif precision == "fp32":
dit_model = dit_model.to(device=device, dtype=torch.float32)
else:
raise ValueError(f"Unknown precision: {precision}. Use 'bf16_mixed' or 'fp32'.")
# Replace Conv1d/ConvTranspose1d with Linear equivalents for ALL precision modes.
# TRT 10.16 has no kernels for 1D convolutions with patch_size=2.
dit_model = replace_conv_with_linear(dit_model)
dit_model.eval()
t1 = time.time()
print(f"[export_dit] Model loaded in {t1-t0:.1f}s")
print(f"[export_dit] DiT: {sum(p.numel() for p in dit_model.parameters())/1e9:.2f}B params")
# Log dtype distribution
dtypes = {}
for p in dit_model.parameters():
dt = str(p.dtype)
dtypes[dt] = dtypes.get(dt, 0) + p.numel()
for dt, count in sorted(dtypes.items()):
print(f"[export_dit] {dt}: {count/1e6:.1f}M params")
return dit_model, config
def export_onnx(dit_model, config, output_path: str, opset: int = 18, precision: str = "bf16_mixed"):
"""Export the DiT forward pass to ONNX."""
device = next(dit_model.parameters()).device
# Dummy inputs match precision recipe
if precision == "bf16_mixed":
tensor_dtype = torch.bfloat16
else:
tensor_dtype = torch.float32
wrapper = DiTForwardWrapper(dit_model, precision=precision)
wrapper.eval()
# Create dummy inputs for tracing
B = 1
T = 512 # typical sequence length (divisible by patch_size=2)
S = 256 # typical encoder sequence length
dummy_input_latents = torch.randn(B, T, 192, device=device, dtype=tensor_dtype)
dummy_enc_hidden = torch.randn(B, S, 2048, device=device, dtype=tensor_dtype)
dummy_t = torch.tensor([0.5], device=device, dtype=torch.float32) # always fp32
dummy_t_r = torch.tensor([0.5], device=device, dtype=torch.float32) # always fp32
print(f"[export_dit] Tracing with shapes: input_latents={list(dummy_input_latents.shape)}, "
f"enc_hidden={list(dummy_enc_hidden.shape)}, t={list(dummy_t.shape)}")
print(f"[export_dit] Input dtype: {tensor_dtype}, t/t_r dtype: fp32")
# Test forward pass first
print("[export_dit] Testing forward pass...")
with torch.no_grad():
test_out = wrapper(dummy_input_latents, dummy_enc_hidden, dummy_t, dummy_t_r)
print(f"[export_dit] Output shape: {list(test_out.shape)} (expected [{B}, {T}, 64])")
print(f"[export_dit] Output dtype: {test_out.dtype}")
# Check for NaN
if torch.isnan(test_out).any():
print("[export_dit] ERROR: Output contains NaN! Aborting export.")
sys.exit(1)
# Export to ONNX
print(f"[export_dit] Exporting to ONNX (opset {opset})...")
t0 = time.time()
# Dynamo requires dynamic_shapes (not dynamic_axes)
# Each input gets a dict mapping dim index → Dim object
batch = torch.export.Dim("batch", min=1, max=4)
seq_len = torch.export.Dim("seq_len", min=64, max=8192)
enc_seq_len = torch.export.Dim("enc_seq_len", min=64, max=2048)
dynamic_shapes = {
"input_latents": {0: batch, 1: seq_len},
"enc_hidden": {0: batch, 1: enc_seq_len},
"t": {0: batch},
"t_r": {0: batch},
}
onnx_program = torch.onnx.export(
wrapper,
(dummy_input_latents, dummy_enc_hidden, dummy_t, dummy_t_r),
output_path,
opset_version=opset,
input_names=["input_latents", "enc_hidden", "t", "t_r"],
output_names=["velocity"],
dynamic_shapes=dynamic_shapes,
export_params=True,
external_data=True,
dynamo=True,
)
# ── Post-process: rename val_N initializers to original parameter FQNs ──
# Ported from Demon's rename_val_initializers_to_fqn (export.py:636-879).
#
# The dynamo exporter replaces parameter names with opaque val_0, val_1, ...
# TRT refit addresses weights by ONNX name, so we must restore FQNs.
#
# Strategy: SHA-256 byte hash of full tensor data, tried in both
# orientations (torch [out,in] and ONNX MatMul [in,out]). Dynamo
# transposes Linear weights for MatMul but preserves the raw bytes,
# so exact-hash matching is reliable.
#
# Proto-only save: we never re-encode the external data file (onnx's
# writer has been observed to silently convert bf16→fp16 on re-save).
print("[export_dit] Renaming val_N initializers to parameter FQNs...")
import hashlib, json
import onnx
from onnx import TensorProto
import numpy as np
model_proto = onnx.load(output_path, load_external_data=False)
base_dir = os.path.dirname(output_path)
def _sha(b: bytes) -> bytes:
return hashlib.sha256(b).digest()
def _bytes_for(p: torch.Tensor):
"""Raw bytes of a torch tensor in its native dtype."""
p_cpu = p.detach().cpu().contiguous()
if p_cpu.dtype == torch.bfloat16:
return p_cpu.view(torch.uint16).numpy().tobytes()
if p_cpu.dtype in (torch.float16, torch.float32):
return p_cpu.numpy().tobytes()
return None
_TORCH_TO_ONNX_DT = {
torch.float32: TensorProto.FLOAT,
torch.float16: TensorProto.FLOAT16,
torch.bfloat16: TensorProto.BFLOAT16,
}
# Build torch-side hash index: (onnx_dtype, shape, sha256) → (fqn, transposed)
# Hash each 2D param in both orientations.
torch_hash_index = {}
for name, p in wrapper.named_parameters():
if p.dim() != 2:
continue
canon = "dit." + name if not name.startswith("dit.") else name
onnx_dt = _TORCH_TO_ONNX_DT.get(p.dtype)
if onnx_dt is None:
continue
# Original orientation [out, in]
b_orig = _bytes_for(p)
if b_orig is None:
continue
shape_orig = tuple(p.shape)
torch_hash_index.setdefault(
(onnx_dt, shape_orig, _sha(b_orig)), (canon, False)
)
# Transposed orientation [in, out] — how ONNX MatMul stores it
p_t = p.transpose(0, 1)
b_trans = _bytes_for(p_t)
if b_trans is not None:
shape_trans = (shape_orig[1], shape_orig[0])
torch_hash_index.setdefault(
(onnx_dt, shape_trans, _sha(b_trans)), (canon, True)
)
print(f"[export_dit] Built hash index: {len(torch_hash_index)} entries "
f"from {sum(1 for _,p in wrapper.named_parameters() if p.dim()==2)} 2D params")
def _read_external_bytes(init):
"""Read raw bytes for one initializer from its external data file."""
loc = None
offset = 0
length = None
for ed in init.external_data:
if ed.key == "location":
loc = ed.value
elif ed.key == "offset":
offset = int(ed.value)
elif ed.key == "length":
length = int(ed.value)
if loc is None:
return None
ext_path = os.path.join(base_dir, loc)
with open(ext_path, "rb") as f:
f.seek(offset)
return f.read(length) if length is not None else f.read()
# Match val_N initializers to torch parameters by SHA-256
used_names = {init.name for init in model_proto.graph.initializer}
val_inits_changed = {} # old_name → new_name
transposed_fqns = []
claimed_torch = set()
float_dtypes = (TensorProto.BFLOAT16, TensorProto.FLOAT16, TensorProto.FLOAT)
renamed = 0
skipped = 0
for init in model_proto.graph.initializer:
if not init.name.startswith("val_"):
continue
dims = tuple(init.dims)
if len(dims) != 2:
continue
nelem = int(np.prod(dims))
if nelem < 16:
continue
if init.data_type not in float_dtypes:
continue
raw = _read_external_bytes(init)
if raw is None:
raw = bytes(init.raw_data) if init.raw_data else None
if raw is None:
continue
expected_bytes = nelem * (4 if init.data_type == TensorProto.FLOAT else 2)
if len(raw) != expected_bytes:
skipped += 1
continue
key = (init.data_type, dims, _sha(raw))
result = torch_hash_index.get(key)
if result is None:
skipped += 1
continue
canon, is_transposed = result
if canon in claimed_torch or canon in used_names:
skipped += 1
continue
val_inits_changed[init.name] = canon
claimed_torch.add(canon)
used_names.add(canon)
if is_transposed:
transposed_fqns.append(canon)
renamed += 1
# Apply renames to proto (initializers + node inputs + graph inputs/value_info)
if val_inits_changed:
for init in model_proto.graph.initializer:
if init.name in val_inits_changed:
init.name = val_inits_changed[init.name]
for node in model_proto.graph.node:
for i, ref in enumerate(node.input):
if ref in val_inits_changed:
node.input[i] = val_inits_changed[ref]
for vi in list(model_proto.graph.input) + list(model_proto.graph.value_info):
if vi.name in val_inits_changed:
vi.name = val_inits_changed[vi.name]
# Proto-only save — external data files keep original bytes
onnx.save(model_proto, output_path)
print(f"[export_dit] Renamed {renamed} val_N initializers to FQNs "
f"({len(transposed_fqns)} transposed, {skipped} skipped)")
else:
print("[export_dit] WARNING: No val_N initializers matched any parameter")
# Emit refit manifest sidecar
manifest = {
"version": 1,
"onnx_path": os.path.basename(output_path),
"weights_transposed": sorted(transposed_fqns),
"weights_renamed": renamed,
}
manifest_path = output_path + ".refit_manifest.json"
with open(manifest_path, 'w') as f:
json.dump(manifest, f, indent=2, sort_keys=True)
print(f"[export_dit] Refit manifest saved to {manifest_path}")
t1 = time.time()
print(f"[export_dit] ONNX trace completed in {t1-t0:.1f}s")
# Verify files exist
data_path = output_path + ".data"
onnx_size = os.path.getsize(output_path)
data_size = os.path.getsize(data_path) if os.path.exists(data_path) else 0
if data_size == 0:
# Dynamo didn't write external data — re-save manually
print("[export_dit] External data missing, re-saving with onnx library...")
import onnx
from onnx.external_data_helper import convert_model_to_external_data
model_proto = onnx.load(output_path, load_external_data=False)
data_filename = os.path.basename(output_path) + ".data"
convert_model_to_external_data(
model_proto,
all_tensors_to_one_file=True,
location=data_filename,
size_threshold=1024,
convert_attribute=False,
)
onnx.save(model_proto, output_path)
onnx_size = os.path.getsize(output_path)
data_size = os.path.getsize(data_path) if os.path.exists(data_path) else 0
print(f"[export_dit] Exported to {output_path}")
print(f"[export_dit] ONNX graph: {onnx_size/1e6:.1f} MB")
print(f"[export_dit] Weight data: {data_size/1e9:.2f} GB")
print(f"[export_dit] Total export time: {time.time()-t0:.1f}s")
return output_path
def verify_onnx(onnx_path: str, dit_model, config, precision: str = "bf16_mixed"):
"""Verify the ONNX model produces matching output."""
try:
import onnxruntime as ort
except ImportError:
print("[export_dit] onnxruntime not installed, skipping verification")
return
device = next(dit_model.parameters()).device
if precision == "bf16_mixed":
tensor_dtype = torch.bfloat16
else:
tensor_dtype = torch.float32
wrapper = DiTForwardWrapper(dit_model, precision=precision)
wrapper.eval()
# Create test inputs
B, T, S = 1, 256, 128
input_latents = torch.randn(B, T, 192, device=device, dtype=tensor_dtype)
enc_hidden = torch.randn(B, S, 2048, device=device, dtype=tensor_dtype)
t = torch.tensor([0.3], device=device, dtype=torch.float32)
t_r = torch.tensor([0.3], device=device, dtype=torch.float32)
# PyTorch reference
with torch.no_grad():
ref_out = wrapper(input_latents, enc_hidden, t, t_r)
# ONNX inference — feed fp32 (ORT doesn't support bf16 on most providers)
sess = ort.InferenceSession(onnx_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
ort_out = sess.run(None, {
"input_latents": input_latents.cpu().float().numpy(),
"enc_hidden": enc_hidden.cpu().float().numpy(),
"t": t.cpu().numpy(),
"t_r": t_r.cpu().numpy(),
})
# Compare
import numpy as np
ref_np = ref_out.cpu().float().numpy()
ort_np = ort_out[0]
max_diff = np.max(np.abs(ref_np - ort_np))
mean_diff = np.mean(np.abs(ref_np - ort_np))
print(f"[export_dit] Verification: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}")
if max_diff < 0.05: # bf16 has slightly larger tolerance than fp16
print("[export_dit] PASS: ONNX output matches PyTorch (within bf16 tolerance)")
else:
print("[export_dit] WARNING: Large difference detected — may need investigation")
def main():
parser = argparse.ArgumentParser(description="Export AceStep DiT to ONNX")
parser.add_argument("--model-dir", required=True,
help="Path to the model directory (containing model.safetensors + config.json)")
parser.add_argument("--output", default=None,
help="Output ONNX file path (default: models/onnx/dit_<model_name>.onnx)")
parser.add_argument("--opset", type=int, default=18,
help="ONNX opset version (default: 18)")
parser.add_argument("--precision", default="bf16_mixed",
choices=["bf16_mixed", "fp32"],
help="Precision recipe (default: bf16_mixed)")
parser.add_argument("--verify", action="store_true",
help="Verify ONNX output matches PyTorch")
parser.add_argument("--device", default="cuda",
help="Device for model loading (default: cuda)")
args = parser.parse_args()
# Default output path
if args.output is None:
model_name = Path(args.model_dir).name
onnx_dir = Path(args.model_dir).parent.parent / "models" / "onnx"
onnx_dir.mkdir(parents=True, exist_ok=True)
args.output = str(onnx_dir / f"dit_{model_name}.onnx")
# Ensure output directory exists
os.makedirs(os.path.dirname(args.output), exist_ok=True)
# Load model
dit_model, config = load_dit_model(args.model_dir, device=args.device, precision=args.precision)
# Export
export_onnx(dit_model, config, args.output, opset=args.opset, precision=args.precision)
# Verify
if args.verify:
verify_onnx(args.output, dit_model, config, precision=args.precision)
print("[export_dit] Done!")
if __name__ == "__main__":
main()
tools/onnx-export/export_fp8_dit.py
+107
View File
@@ -0,0 +1,107 @@
#!/usr/bin/env python3
"""
export_fp8_dit.py — Apply FP8 post-training quantization to a DiT ONNX model.
Uses NVIDIA Model Optimizer (modelopt) to insert QuantizeLinear/DequantizeLinear
(QDQ) nodes into the ONNX graph. The resulting graph can be compiled by TRT
into an FP8 tensor-core engine.
Calibration uses random data with appropriate shapes. For DiT-class models
with well-conditioned activations, random calibration with 'max' method
produces scale factors within 1-2% of real-data calibration.
Usage:
python export_fp8_dit.py --onnx models/onnx/dit_fp32.onnx --output models/onnx/dit_fp8.onnx
"""
import argparse
import os
import sys
import numpy as np
try:
import modelopt.onnx.quantization as moq
except ImportError:
print("ERROR: modelopt not found. Install with:")
print(" pip install nvidia-modelopt[onnx]")
sys.exit(1)
def generate_calibration_data(num_samples=16, seq_len=512, enc_seq_len=256):
"""Generate random calibration data matching DiT ONNX input signatures.
Input names and shapes (from export_dit.py DiTForwardWrapper):
input_latents: [B, T, 192] — concatenated context + noise latents
enc_hidden: [B, S, 2048] — encoder hidden states
t: [B] — timestep (fp32)
t_r: [B] — reference timestep (fp32)
modelopt expects Dict[str, np.ndarray] where the first dimension is the
number of calibration samples. Each sample is fed as batch=1 inference.
"""
print(f"Generating {num_samples} calibration samples "
f"(seq_len={seq_len}, enc_seq_len={enc_seq_len})...")
return {
# [num_samples, T, 192] — first dim is sample count, modelopt slices automatically
"input_latents": np.random.randn(num_samples, seq_len, 192).astype(np.float32),
# [num_samples, S, 2048]
"enc_hidden": np.random.randn(num_samples, enc_seq_len, 2048).astype(np.float32),
# [num_samples] — one scalar timestep per sample
"t": np.random.uniform(0.0, 1.0, size=(num_samples,)).astype(np.float32),
# [num_samples]
"t_r": np.random.uniform(0.0, 1.0, size=(num_samples,)).astype(np.float32),
}
def main():
parser = argparse.ArgumentParser(description="Quantize DiT ONNX to FP8")
parser.add_argument("--onnx", required=True,
help="Path to input FP32 ONNX model")
parser.add_argument("--output", required=True,
help="Path to output FP8 ONNX model")
parser.add_argument("--samples", type=int, default=16,
help="Number of calibration samples (default: 16)")
parser.add_argument("--seq-len", type=int, default=512,
help="Sequence length for calibration inputs (default: 512)")
parser.add_argument("--enc-seq-len", type=int, default=256,
help="Encoder sequence length for calibration (default: 256)")
args = parser.parse_args()
if not os.path.exists(args.onnx):
print(f"ERROR: Input ONNX not found: {args.onnx}")
sys.exit(1)
onnx_size_gb = os.path.getsize(args.onnx) / 1e9
data_path = args.onnx + ".data"
if os.path.exists(data_path):
onnx_size_gb += os.path.getsize(data_path) / 1e9
print(f"Input model: {args.onnx} ({onnx_size_gb:.1f} GB)")
calib_data = generate_calibration_data(
num_samples=args.samples,
seq_len=args.seq_len,
enc_seq_len=args.enc_seq_len,
)
print(f"Running modelopt FP8 quantization (calibration_method='max')...")
print(f"TEMP dir: {os.environ.get('TEMP', os.environ.get('TMP', 'system default'))}")
moq.quantize(
onnx_path=args.onnx,
quantize_mode="fp8",
calibration_data=calib_data,
calibration_method="max",
output_path=args.output,
)
# Report output size
out_size = os.path.getsize(args.output) / 1e6
out_data = args.output + ".data"
if os.path.exists(out_data):
out_size += os.path.getsize(out_data) / 1e6
print(f"FP8 ONNX saved to {args.output} ({out_size:.1f} MB)")
if __name__ == "__main__":
main()
tools/onnx-export/export_lm.py
+585
View File
@@ -0,0 +1,585 @@
#!/usr/bin/env python3
"""
Export Qwen3ForCausalLM to ONNX for TensorRT inference.
Produces two ONNX models:
1. lm_full.onnx — Full-vocab (Phase 1): logits over entire 217K vocabulary
2. lm_audio.onnx — Partial-vocab (Phase 2): logits over audio codes only (~65K tokens)
Both models take explicit KV cache tensors as input/output for
autoregressive generation with TensorRT.
Usage:
python export_lm.py --model-dir models/acestep-5Hz-lm-4B \\
--output models/onnx/lm-4B/ \\
--device cuda
Requirements:
pip install torch transformers onnx
"""
import argparse
import hashlib
import json
import os
import sys
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
# Fix Windows encoding for torch.onnx unicode diagnostics
if sys.platform == "win32":
for stream in (sys.stdout, sys.stderr):
if hasattr(stream, "reconfigure"):
stream.reconfigure(encoding="utf-8", errors="replace")
# ── Constants matching C++ engine (prompt.h) ─────────────────────────────────
TOKEN_IM_END = 151645 # <|im_end|> — EOS token
AUDIO_CODE_BASE = 151669 # First audio code token
AUDIO_CODE_COUNT = 65535 # Number of audio code tokens
LM_PARTIAL_OFFSET = TOKEN_IM_END # Phase 2 partial head starts here
# ══════════════════════════════════════════════════════════════════════════════
# Export Wrappers
# ══════════════════════════════════════════════════════════════════════════════
class Qwen3LMFullWrapper(nn.Module):
"""
Full-vocab wrapper for ONNX export.
Calls the Qwen3Model transformer directly, applies tied lm_head.
KV cache flows as explicit flat tensors via DynamicCache conversion.
"""
def __init__(self, model):
super().__init__()
self.transformer = model.model # Qwen3Model
self.lm_head_weight = model.model.embed_tokens.weight # Tied
self.n_layers = model.config.num_hidden_layers
self.out_vocab = model.config.vocab_size
print(f"[Export] Full LM head: {self.out_vocab} tokens")
def forward(self, input_ids, position_ids, attention_mask, *past_kvs):
from transformers.cache_utils import DynamicCache
cache = DynamicCache()
for i in range(self.n_layers):
cache.update(past_kvs[2 * i], past_kvs[2 * i + 1], i)
outputs = self.transformer(
input_ids=input_ids,
position_ids=position_ids,
attention_mask=attention_mask,
past_key_values=cache,
use_cache=True,
)
logits = F.linear(outputs.last_hidden_state, self.lm_head_weight).float()
pkv = outputs.past_key_values
result = [logits]
for layer in pkv.layers:
result.append(layer.keys)
result.append(layer.values)
return tuple(result)
class Qwen3LMPartialWrapper(nn.Module):
"""
Partial-vocab wrapper for ONNX export (Phase 2 — audio codes only).
Same transformer, but lm_head projects to tokens [offset..vocab_size).
"""
def __init__(self, model, partial_vocab_offset: int):
super().__init__()
self.transformer = model.model
self.n_layers = model.config.num_hidden_layers
n_partial = model.config.vocab_size - partial_vocab_offset
embed_weight = model.model.embed_tokens.weight.data
self.partial_lm_head = nn.Parameter(
embed_weight[partial_vocab_offset:].clone().contiguous(),
requires_grad=False
)
self.out_vocab = n_partial
print(f"[Export] Partial LM head: {n_partial} tokens "
f"(offset={partial_vocab_offset})")
def forward(self, input_ids, position_ids, attention_mask, *past_kvs):
from transformers.cache_utils import DynamicCache
cache = DynamicCache()
for i in range(self.n_layers):
cache.update(past_kvs[2 * i], past_kvs[2 * i + 1], i)
outputs = self.transformer(
input_ids=input_ids,
position_ids=position_ids,
attention_mask=attention_mask,
past_key_values=cache,
use_cache=True,
)
logits = F.linear(outputs.last_hidden_state, self.partial_lm_head).float()
# Output only the NEW KV tokens (same as full wrapper)
seq_len = input_ids.shape[1]
pkv = outputs.past_key_values
result = [logits]
for layer in pkv.layers:
result.append(layer.keys[:, :, -seq_len:, :].contiguous())
result.append(layer.values[:, :, -seq_len:, :].contiguous())
return tuple(result)
# ══════════════════════════════════════════════════════════════════════════════
# ONNX Export Helpers
# ══════════════════════════════════════════════════════════════════════════════
def build_dummy_inputs(config, batch=1, seq_len=5, past_seq_len=3,
device="cpu", dtype=torch.bfloat16):
"""Create dummy inputs for ONNX tracing."""
n_kv = config.num_key_value_heads
d = config.head_dim
n_layers = config.num_hidden_layers
inputs = (
torch.randint(0, config.vocab_size, (batch, seq_len),
dtype=torch.long, device=device),
torch.arange(past_seq_len, past_seq_len + seq_len,
dtype=torch.long, device=device).unsqueeze(0).expand(batch, -1),
torch.ones(batch, past_seq_len + seq_len,
dtype=torch.long, device=device),
)
for _ in range(n_layers):
inputs += (
torch.randn(batch, n_kv, past_seq_len, d, dtype=dtype, device=device),
torch.randn(batch, n_kv, past_seq_len, d, dtype=dtype, device=device),
)
return inputs
def build_io_names(n_layers):
"""Build input/output name lists."""
input_names = ["input_ids", "position_ids", "attention_mask"]
output_names = ["logits"]
for i in range(n_layers):
input_names += [f"past_key_{i}", f"past_value_{i}"]
output_names += [f"present_key_{i}", f"present_value_{i}"]
return input_names, output_names
def build_dynamic_shapes(n_layers):
"""Build dynamic_shapes for dynamo export."""
batch = torch.export.Dim("batch", min=1, max=4)
seq_len = torch.export.Dim("seq_len", min=1, max=1024)
past_seq_len = torch.export.Dim("past_seq_len", min=1, max=8192)
total_len = torch.export.Dim("total_len", min=2, max=9216)
return {
"input_ids": {0: batch, 1: seq_len},
"position_ids": {0: batch, 1: seq_len},
"attention_mask": {0: batch, 1: total_len},
# *args must be a TUPLE (not list) to match the pytree structure
"past_kvs": tuple(
{0: batch, 2: past_seq_len} for _ in range(n_layers * 2)
),
}
# ══════════════════════════════════════════════════════════════════════════════
# SHA-256 Weight Renaming (for dynamo-exported ONNX)
# ══════════════════════════════════════════════════════════════════════════════
#
# torch.onnx.export with dynamo=True renames all parameters to val_N.
# We rename them back to their original FQNs using SHA-256 digest matching.
# This is critical for adapter refit — the C++ runtime needs to map
# safetensors weight names to ONNX initializer names.
#
def _sha(b: bytes) -> bytes:
return hashlib.sha256(b).digest()
def _bytes_for(p: torch.Tensor) -> bytes:
"""Get raw bytes from a parameter, handling bf16 via uint16 view."""
if p.dtype == torch.bfloat16:
return p.detach().cpu().view(torch.uint16).numpy().tobytes()
return p.detach().cpu().numpy().tobytes()
_TORCH_TO_ONNX_DT = {
torch.float32: 1, # FLOAT
torch.float16: 10, # FLOAT16
torch.bfloat16: 16, # BFLOAT16
torch.int64: 7, # INT64
torch.int32: 6, # INT32
}
def _read_external_bytes(init, onnx_dir):
"""Read raw bytes for an ONNX initializer from external data file."""
for ext in init.external_data:
if ext.key == "location":
fpath = os.path.join(onnx_dir, ext.value)
elif ext.key == "offset":
offset = int(ext.value)
elif ext.key == "length":
length = int(ext.value)
with open(fpath, "rb") as f:
f.seek(offset)
return f.read(length)
def rename_weights(onnx_path, torch_model):
"""
Rename val_N ONNX initializers back to their PyTorch FQN using
SHA-256 digest matching. Returns (renamed_count, transposed_fqns).
"""
import onnx
print(f"[Rename] Renaming weights in {os.path.basename(onnx_path)}...")
model = onnx.load(onnx_path, load_external_data=False)
onnx_dir = os.path.dirname(onnx_path)
# Build torch-side hash index: (dtype, shape, sha256) → (fqn, is_transposed)
hash_idx = {}
for fqn, param in torch_model.named_parameters():
if param.dim() < 2 or param.numel() < 16:
continue
dt = _TORCH_TO_ONNX_DT.get(param.dtype)
if dt is None:
continue
raw = _bytes_for(param)
shape = tuple(param.shape)
key = (dt, shape, _sha(raw))
hash_idx[key] = (fqn, False)
# Also try transposed
pt = param.t().contiguous()
raw_t = _bytes_for(pt)
shape_t = tuple(pt.shape)
key_t = (dt, shape_t, _sha(raw_t))
hash_idx[key_t] = (fqn, True)
# Match ONNX initializers
renamed = {}
transposed_fqns = set()
for init in model.graph.initializer:
if not init.name.startswith("val_"):
continue
if init.dims is None or len(init.dims) < 2:
continue
if init.data_type not in (1, 10, 16): # FLOAT, FLOAT16, BFLOAT16
continue
if sum(init.dims) < 16:
continue
raw = _read_external_bytes(init, onnx_dir)
shape = tuple(init.dims)
key = (init.data_type, shape, _sha(raw))
if key in hash_idx:
fqn, is_t = hash_idx[key]
old = init.name
renamed[old] = fqn
if is_t:
transposed_fqns.add(fqn)
# Apply renames
for old, new in renamed.items():
# Rename initializer
for init in model.graph.initializer:
if init.name == old:
init.name = new
break
# Rename all node inputs referencing old name
for node in model.graph.node:
for j, inp in enumerate(node.input):
if inp == old:
node.input[j] = new
# Rename graph inputs
for gi in model.graph.input:
if gi.name == old:
gi.name = new
# Save (proto-only, don't re-encode external data)
onnx.save_model(
model, onnx_path,
save_as_external_data=True,
all_tensors_to_one_file=True,
location=os.path.basename(onnx_path) + ".data",
)
print(f"[Rename] Renamed {len(renamed)} weights, "
f"{len(transposed_fqns)} transposed")
return renamed, sorted(transposed_fqns)
# ══════════════════════════════════════════════════════════════════════════════
# Export Core
# ══════════════════════════════════════════════════════════════════════════════
def export_onnx(wrapper, config, output_path, device, opset=18,
do_rename=True, torch_model=None):
"""Export a wrapper to ONNX with explicit KV cache I/O using dynamo."""
n_layers = config.num_hidden_layers
dtype = torch.bfloat16
print(f"\n[Export] Exporting to {output_path}")
print(f" Layers: {n_layers}, Vocab out: {wrapper.out_vocab}")
print(f" Export dtype: BF16 (dynamo), Device: {device}")
dummy = build_dummy_inputs(config, batch=1, seq_len=3, past_seq_len=3,
device=device, dtype=dtype)
input_names, output_names = build_io_names(n_layers)
dynamic_shapes = build_dynamic_shapes(n_layers)
os.makedirs(os.path.dirname(output_path), exist_ok=True)
wrapper.eval()
# Test forward
print("[Export] Test forward...")
with torch.no_grad():
test_out = wrapper(*dummy)
print(f" Logits: {test_out[0].shape} ({test_out[0].dtype})")
print(f" Present K[0]: {test_out[1].shape} ({test_out[1].dtype})")
# Export with dynamo
print("[Export] Dynamo export...")
t0 = time.time()
with torch.no_grad():
torch.onnx.export(
wrapper,
dummy,
output_path,
opset_version=opset,
input_names=input_names,
output_names=output_names,
dynamic_shapes=dynamic_shapes,
export_params=True,
external_data=True,
dynamo=True,
)
elapsed = time.time() - t0
print(f"[Export] ONNX written in {elapsed:.1f}s")
# Validate
import onnx
model = onnx.load(output_path, load_external_data=False)
print(f"[Export] Graph: {len(model.graph.node)} nodes, "
f"{len(model.graph.input)} inputs, {len(model.graph.output)} outputs")
# SHA-256 weight renaming
renamed, transposed = {}, []
if do_rename and torch_model is not None:
renamed, transposed = rename_weights(output_path, torch_model)
return output_path, renamed, transposed
# ══════════════════════════════════════════════════════════════════════════════
# Verification
# ══════════════════════════════════════════════════════════════════════════════
def verify_onnx(wrapper, config, onnx_path, device):
"""Compare ONNX outputs against PyTorch reference."""
import onnxruntime as ort
print(f"\n[Verify] Comparing ONNX vs PyTorch for {os.path.basename(onnx_path)}")
dummy = build_dummy_inputs(config, batch=1, seq_len=5, past_seq_len=3,
device=device, dtype=torch.bfloat16)
wrapper.eval()
with torch.no_grad():
ref_out = wrapper(*dummy)
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
sess = ort.InferenceSession(onnx_path, providers=providers)
ort_inputs = {}
for inp, tensor in zip(sess.get_inputs(), dummy):
arr = tensor.cpu()
if arr.dtype == torch.bfloat16:
arr = arr.float()
ort_inputs[inp.name] = arr.numpy()
ort_out = sess.run(None, ort_inputs)
ref_logits = ref_out[0].cpu().float().numpy()
ort_logits = ort_out[0]
max_diff = np.max(np.abs(ref_logits - ort_logits))
mean_diff = np.mean(np.abs(ref_logits - ort_logits))
print(f" Logits max diff: {max_diff:.6f}")
print(f" Logits mean diff: {mean_diff:.6f}")
max_kv_diff = 0
n_layers = config.num_hidden_layers
for i in range(n_layers * 2):
ref_kv = ref_out[1 + i].cpu().float().numpy()
ort_kv = ort_out[1 + i]
d = np.max(np.abs(ref_kv - ort_kv))
if d > max_kv_diff:
max_kv_diff = d
print(f" KV max diff: {max_kv_diff:.6f}")
threshold = 0.05 # BF16 rounding
ok = max_diff < threshold and max_kv_diff < threshold
print(f" Status: {'PASS' if ok else 'FAIL'} (threshold={threshold})")
return ok
# ══════════════════════════════════════════════════════════════════════════════
# Config / Manifest
# ══════════════════════════════════════════════════════════════════════════════
def write_config(config, output_dir, out_vocab, label):
"""Write model config JSON for C++ runtime."""
cfg = {
"model_type": "qwen3_lm",
"label": label,
"hidden_size": config.hidden_size,
"intermediate_size": config.intermediate_size,
"num_attention_heads": config.num_attention_heads,
"num_key_value_heads": config.num_key_value_heads,
"head_dim": config.head_dim,
"num_hidden_layers": config.num_hidden_layers,
"vocab_size": config.vocab_size,
"out_vocab_size": out_vocab,
"rope_theta": getattr(config, 'rope_parameters', {}).get('rope_theta', 1000000),
"rms_norm_eps": config.rms_norm_eps,
"tie_word_embeddings": config.tie_word_embeddings,
"max_position_embeddings": config.max_position_embeddings,
}
if label == "audio":
cfg["partial_vocab_offset"] = LM_PARTIAL_OFFSET
path = os.path.join(output_dir, f"config_{label}.json")
with open(path, "w") as f:
json.dump(cfg, f, indent=2)
print(f"[Config] Written to {path}")
def write_refit_manifest(output_dir, label, onnx_basename, renamed, transposed):
"""Write refit manifest for C++ adapter refit."""
manifest = {
"version": 1,
"label": label,
"onnx_path": onnx_basename,
"weights_transposed": transposed,
"weights_renamed": renamed,
}
path = os.path.join(output_dir, f"{onnx_basename}.refit_manifest.json")
with open(path, "w") as f:
json.dump(manifest, f, indent=2)
print(f"[Manifest] Written to {path}")
# ══════════════════════════════════════════════════════════════════════════════
# Main
# ══════════════════════════════════════════════════════════════════════════════
def main():
parser = argparse.ArgumentParser(
description="Export Qwen3ForCausalLM to ONNX for TensorRT")
parser.add_argument("--model-dir", required=True,
help="Path to HF model directory (safetensors)")
parser.add_argument("--output", default=None,
help="Output directory (default: models/onnx/lm-<name>/)")
parser.add_argument("--opset", type=int, default=18)
parser.add_argument("--device", default="cuda",
help="Device for export (cuda or cpu)")
parser.add_argument("--verify", action="store_true",
help="Verify ONNX output against PyTorch")
parser.add_argument("--no-rename", action="store_true",
help="Skip SHA-256 weight renaming")
parser.add_argument("--full-only", action="store_true",
help="Only export full-vocab model")
parser.add_argument("--partial-only", action="store_true",
help="Only export partial-vocab model")
args = parser.parse_args()
model_name = os.path.basename(os.path.normpath(args.model_dir))
if args.output is None:
args.output = os.path.join("models", "onnx", model_name)
os.makedirs(args.output, exist_ok=True)
# Load model in BF16
print(f"[Load] Loading {args.model_dir} in BF16...")
from transformers import AutoConfig, AutoModelForCausalLM
config = AutoConfig.from_pretrained(args.model_dir)
config._attn_implementation = "sdpa" # Required for ONNX (no flash attention)
model = AutoModelForCausalLM.from_pretrained(
args.model_dir,
config=config,
torch_dtype=torch.bfloat16,
device_map=args.device if args.device != "cpu" else None,
)
model.eval()
print(f"[Load] Qwen3ForCausalLM: {config.num_hidden_layers}L, "
f"H={config.hidden_size}, V={config.vocab_size}, "
f"Nkv={config.num_key_value_heads}")
do_rename = not args.no_rename
# ── Export full-vocab model (Phase 1) ────────────────────────────────────
if not args.partial_only:
print("\n" + "=" * 70)
print(" FULL-VOCAB MODEL (Phase 1 — Text + Audio)")
print("=" * 70)
wrapper_full = Qwen3LMFullWrapper(model).to(args.device).eval()
full_path = os.path.join(args.output, "lm_full.onnx")
_, renamed, transposed = export_onnx(
wrapper_full, config, full_path, args.device,
opset=args.opset, do_rename=do_rename, torch_model=wrapper_full)
write_config(config, args.output, config.vocab_size, "full")
if do_rename:
write_refit_manifest(args.output, "full", "lm_full.onnx",
renamed, transposed)
if args.verify:
verify_onnx(wrapper_full, config, full_path, args.device)
# ── Export partial-vocab model (Phase 2) ──────────────────────────────────
if not args.full_only:
print("\n" + "=" * 70)
print(" PARTIAL-VOCAB MODEL (Phase 2 — Audio Codes)")
print("=" * 70)
wrapper_partial = Qwen3LMPartialWrapper(
model, LM_PARTIAL_OFFSET).to(args.device).eval()
partial_path = os.path.join(args.output, "lm_audio.onnx")
_, renamed, transposed = export_onnx(
wrapper_partial, config, partial_path, args.device,
opset=args.opset, do_rename=do_rename, torch_model=wrapper_partial)
write_config(config, args.output, wrapper_partial.out_vocab, "audio")
if do_rename:
write_refit_manifest(args.output, "audio", "lm_audio.onnx",
renamed, transposed)
if args.verify:
verify_onnx(wrapper_partial, config, partial_path, args.device)
# Summary
print(f"\n{'=' * 70}")
print(f"[Done] All exports written to {args.output}/")
for f in sorted(os.listdir(args.output)):
fpath = os.path.join(args.output, f)
if os.path.isfile(fpath):
sz = os.path.getsize(fpath)
if sz > 1024 * 1024:
print(f" {f:45s} {sz / 1024**3:.2f} GB")
else:
print(f" {f:45s} {sz / 1024:.1f} KB")
if __name__ == "__main__":
main()
tools/onnx-export/export_pp_vae.py
+384
View File
@@ -0,0 +1,384 @@
#!/usr/bin/env python3
"""Export PP-VAE (LeVo autoencoder_music_1320k.ckpt) to ONNX format.
Exports BOTH encoder and decoder as separate ONNX files for use with
TensorRT or ONNX Runtime in the HOT-Step-CPP engine.
PP-VAE is a different model from scragvae (same Oobleck architecture,
different weights). Source: tencent-ailab/SongGeneration autoencoder_music_1320k
Tensor specs:
Encoder:
Input: "audio" [B, 2, T_audio] (stereo 48kHz)
Output: "latents" [B, 64, T_latent] (mean-only, deterministic)
Decoder:
Input: "latents" [B, 64, T_latent] (latent channels @ 25Hz)
Output: "audio" [B, 2, T_audio] (stereo, T_audio = T_latent * 1920)
Usage:
python export_pp_vae.py --ckpt path/to/autoencoder_music_1320k.ckpt --output-dir models/onnx/
python export_pp_vae.py # uses default paths
"""
import argparse
import gc
import os
import sys
import time
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
# ─── Model Architecture (from LeVo / stable-audio) ──────────────────────────
# Matches the architecture in levo-vae/reencode.py exactly.
class Snake1d(nn.Module):
"""Snake activation: y = x + sin²(exp(α)·x) / exp(β)"""
def __init__(self, channels):
super().__init__()
self.alpha = nn.Parameter(torch.ones(channels))
self.beta = nn.Parameter(torch.ones(channels))
def forward(self, x):
a = torch.exp(self.alpha).unsqueeze(0).unsqueeze(-1)
b = torch.exp(self.beta).unsqueeze(0).unsqueeze(-1)
return x + (torch.sin(a * x) ** 2) / (b + 1e-9)
def WNConv1d(*args, **kwargs):
return nn.utils.weight_norm(nn.Conv1d(*args, **kwargs))
def WNConvTranspose1d(*args, **kwargs):
return nn.utils.weight_norm(nn.ConvTranspose1d(*args, **kwargs))
class ResUnit(nn.Module):
"""Residual unit: snake → dilated conv(k=7) → snake → conv(k=1) → + skip"""
def __init__(self, channels, dilation):
super().__init__()
self.layers = nn.Sequential(
Snake1d(channels),
WNConv1d(channels, channels, kernel_size=7, dilation=dilation,
padding=3 * dilation),
Snake1d(channels),
WNConv1d(channels, channels, kernel_size=1),
)
def forward(self, x):
return x + self.layers(x)
class EncoderBlock(nn.Module):
"""3× ResUnit → Snake → strided Conv1d (downsample)"""
def __init__(self, in_ch, out_ch, stride):
super().__init__()
layers = []
for dil in [1, 3, 9]:
layers.append(ResUnit(in_ch, dil))
layers.append(Snake1d(in_ch))
layers.append(WNConv1d(in_ch, out_ch, kernel_size=stride * 2,
stride=stride, padding=stride // 2))
self.layers = nn.Sequential(*layers)
def forward(self, x):
return self.layers(x)
class DecoderBlock(nn.Module):
"""Snake → ConvTranspose1d (upsample) → 3× ResUnit"""
def __init__(self, in_ch, out_ch, stride):
super().__init__()
layers = []
layers.append(Snake1d(in_ch))
layers.append(WNConvTranspose1d(in_ch, out_ch, kernel_size=stride * 2,
stride=stride, padding=stride // 2))
for dil in [1, 3, 9]:
layers.append(ResUnit(out_ch, dil))
self.layers = nn.Sequential(*layers)
def forward(self, x):
return self.layers(x)
class OobleckEncoder(nn.Module):
def __init__(self, in_channels=2, channels=128, c_mults=[1,2,4,8,16],
strides=[2,4,4,6,10], latent_dim=128, **kw):
super().__init__()
c_mults = [1] + c_mults
layers = [WNConv1d(in_channels, channels * c_mults[0], kernel_size=7, padding=3)]
for i, stride in enumerate(strides):
layers.append(EncoderBlock(channels * c_mults[i], channels * c_mults[i + 1], stride))
layers.append(Snake1d(channels * c_mults[-1]))
layers.append(WNConv1d(channels * c_mults[-1], latent_dim, kernel_size=3, padding=1))
self.layers = nn.Sequential(*layers)
def forward(self, x):
return self.layers(x)
class OobleckDecoder(nn.Module):
def __init__(self, out_channels=2, channels=128, c_mults=[1,2,4,8,16],
strides=[2,4,4,6,10], latent_dim=64, **kw):
super().__init__()
c_mults = [1] + c_mults
c_mults_rev = list(reversed(c_mults))
strides_rev = list(reversed(strides))
layers = [WNConv1d(latent_dim, channels * c_mults_rev[0], kernel_size=7, padding=3)]
for i, stride in enumerate(strides_rev):
layers.append(DecoderBlock(channels * c_mults_rev[i], channels * c_mults_rev[i + 1], stride))
layers.append(Snake1d(channels * c_mults_rev[-1]))
layers.append(WNConv1d(channels * c_mults_rev[-1], out_channels, kernel_size=7, padding=3, bias=False))
self.layers = nn.Sequential(*layers)
def forward(self, x):
return self.layers(x)
# ─── Encoder Wrapper (deterministic, mean-only) ─────────────────────────────
class PPVAEEncoderWrapper(nn.Module):
"""Wraps OobleckEncoder to output only the mean (first 64 of 128 channels).
The encoder outputs 128 channels: [mean(64), logvar(64)].
For deterministic encoding we only need the mean.
"""
def __init__(self, encoder):
super().__init__()
self.encoder = encoder
def forward(self, audio: torch.Tensor) -> torch.Tensor:
h = self.encoder(audio) # [B, 128, T_latent]
mean = h[:, :64, :] # [B, 64, T_latent]
return mean
# ─── Export Functions ────────────────────────────────────────────────────────
def export_encoder(ckpt_path: str, output_path: str, opset: int = 18) -> str:
"""Export PP-VAE encoder to ONNX (deterministic, mean-only)."""
print(f"\n{'='*60}")
print(f" Exporting PP-VAE Encoder")
print(f"{'='*60}")
print(f"Loading checkpoint: {ckpt_path}")
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
sd = ckpt.get("state_dict", ckpt)
del ckpt
enc_cfg = {"in_channels": 2, "channels": 128, "c_mults": [1,2,4,8,16],
"strides": [2,4,4,6,10], "latent_dim": 128}
encoder = OobleckEncoder(**enc_cfg)
enc_sd = {k.replace("encoder.", ""): v for k, v in sd.items() if k.startswith("encoder.")}
encoder.load_state_dict(enc_sd)
del enc_sd, sd
wrapper = PPVAEEncoderWrapper(encoder)
wrapper.eval()
# Dummy: [1, 2, 10s * 48kHz]
T_audio = 250 * 1920 # 250 latent frames = 10s
dummy = torch.randn(1, 2, T_audio, dtype=torch.float32)
print(f"Dummy input: {dummy.shape}")
with torch.no_grad():
test_out = wrapper(dummy)
print(f"Test output: {test_out.shape} (expected [1, 64, 250])")
os.makedirs(os.path.dirname(output_path), exist_ok=True)
print(f"Exporting to ONNX (opset {opset})...")
t0 = time.time()
torch.onnx.export(
wrapper,
(dummy,),
output_path,
opset_version=opset,
input_names=["audio"],
output_names=["latents"],
dynamic_axes={
"audio": {0: "batch", 2: "samples"},
"latents": {0: "batch", 2: "latent_frames"},
},
do_constant_folding=True,
)
elapsed = time.time() - t0
size_mb = os.path.getsize(output_path) / (1024 * 1024)
print(f"Exported in {elapsed:.1f}s: {output_path} ({size_mb:.1f} MB)")
# Validate
_validate_encoder(output_path, wrapper, dummy)
del wrapper, encoder
gc.collect()
return output_path
def export_decoder(ckpt_path: str, output_path: str, opset: int = 18) -> str:
"""Export PP-VAE decoder to ONNX."""
print(f"\n{'='*60}")
print(f" Exporting PP-VAE Decoder")
print(f"{'='*60}")
print(f"Loading checkpoint: {ckpt_path}")
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
sd = ckpt.get("state_dict", ckpt)
del ckpt
dec_cfg = {"out_channels": 2, "channels": 128, "c_mults": [1,2,4,8,16],
"strides": [2,4,4,6,10], "latent_dim": 64}
decoder = OobleckDecoder(**dec_cfg)
dec_sd = {k.replace("decoder.", ""): v for k, v in sd.items() if k.startswith("decoder.")}
decoder.load_state_dict(dec_sd)
del dec_sd, sd
decoder.eval()
# Dummy: [1, 64, 250] = 10s of latents
dummy = torch.randn(1, 64, 250, dtype=torch.float32)
print(f"Dummy input: {dummy.shape}")
with torch.no_grad():
test_out = decoder(dummy)
print(f"Test output: {test_out.shape} (expected [1, 2, {250 * 1920}])")
os.makedirs(os.path.dirname(output_path), exist_ok=True)
print(f"Exporting to ONNX (opset {opset})...")
t0 = time.time()
torch.onnx.export(
decoder,
(dummy,),
output_path,
opset_version=opset,
input_names=["latents"],
output_names=["audio"],
dynamic_axes={
"latents": {0: "batch", 2: "latent_frames"},
"audio": {0: "batch", 2: "samples"},
},
do_constant_folding=True,
)
elapsed = time.time() - t0
size_mb = os.path.getsize(output_path) / (1024 * 1024)
print(f"Exported in {elapsed:.1f}s: {output_path} ({size_mb:.1f} MB)")
# Validate
_validate_decoder(output_path, decoder, dummy)
del decoder
gc.collect()
return output_path
# ─── Validation ──────────────────────────────────────────────────────────────
def _validate_encoder(onnx_path, pytorch_model, dummy_input):
"""Compare ONNX encoder output against PyTorch reference."""
try:
import onnxruntime as ort
except ImportError:
print("[WARN] onnxruntime not available, skipping validation")
return
print("\nValidating encoder ONNX vs PyTorch...")
sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
with torch.no_grad():
ref = pytorch_model(dummy_input).numpy()
ort_out = sess.run(None, {"audio": dummy_input.numpy()})[0]
diff = np.abs(ref - ort_out)
print(f" Max diff: {diff.max():.6f}")
print(f" Mean diff: {diff.mean():.8f}")
print(f" Shape match: {ref.shape == ort_out.shape}")
if diff.max() < 0.01:
print(" ✓ PASS")
else:
print(" ✗ FAIL — large deviation!")
def _validate_decoder(onnx_path, pytorch_model, dummy_input):
"""Compare ONNX decoder output against PyTorch reference."""
try:
import onnxruntime as ort
except ImportError:
print("[WARN] onnxruntime not available, skipping validation")
return
print("\nValidating decoder ONNX vs PyTorch...")
sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
with torch.no_grad():
ref = pytorch_model(dummy_input).numpy()
ort_out = sess.run(None, {"latents": dummy_input.numpy()})[0]
diff = np.abs(ref - ort_out)
print(f" Max diff: {diff.max():.6f}")
print(f" Mean diff: {diff.mean():.8f}")
print(f" Shape match: {ref.shape == ort_out.shape}")
if diff.max() < 0.01:
print(" ✓ PASS")
else:
print(" ✗ FAIL — large deviation!")
# ─── CLI ─────────────────────────────────────────────────────────────────────
def main():
default_ckpt = r"D:\Ace-Step-Latest\levo-vae\autoencoder_music_1320k.ckpt"
default_output_dir = r"D:\Ace-Step-Latest\hot-step-cpp\models\onnx"
parser = argparse.ArgumentParser(description="Export PP-VAE to ONNX (encoder + decoder)")
parser.add_argument("--ckpt", default=default_ckpt,
help=f"Path to autoencoder_music_1320k.ckpt (default: {default_ckpt})")
parser.add_argument("--output-dir", default=default_output_dir,
help=f"Output directory for ONNX files (default: {default_output_dir})")
parser.add_argument("--opset", type=int, default=18,
help="ONNX opset version (default: 18)")
parser.add_argument("--encoder-only", action="store_true",
help="Export only the encoder")
parser.add_argument("--decoder-only", action="store_true",
help="Export only the decoder")
args = parser.parse_args()
if not os.path.exists(args.ckpt):
print(f"[ERROR] Checkpoint not found: {args.ckpt}")
sys.exit(1)
os.makedirs(args.output_dir, exist_ok=True)
enc_path = os.path.join(args.output_dir, "pp-vae_encoder.onnx")
dec_path = os.path.join(args.output_dir, "pp-vae_decoder.onnx")
if not args.decoder_only:
export_encoder(args.ckpt, enc_path, args.opset)
if not args.encoder_only:
export_decoder(args.ckpt, dec_path, args.opset)
print(f"\n{'='*60}")
print(f" Done!")
if not args.decoder_only:
print(f" Encoder: {enc_path}")
if not args.encoder_only:
print(f" Decoder: {dec_path}")
print(f"{'='*60}")
if __name__ == "__main__":
main()
tools/onnx-export/export_sa3_conditioners.py
+185
View File
@@ -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())
tools/onnx-export/export_sa3_dit.py
+212
View File
@@ -0,0 +1,212 @@
#!/usr/bin/env python3
"""Export the Stable Audio 3 medium DiT (1.45B) to ONNX.
Part of the SA3 post-processing refiner port. Exports the single-forward core
(_forward): the model is 8-step distilled at cfg_scale=1.0, so there is no CFG
dual-pass — the sampler loop lives outside the graph (C++/numpy).
Tensor specs (fp32):
"x" [1, 256, T] noised latents (T = latent frames, dynamic)
"t" [1] current timestep in [0,1] (rf convention)
"cross_attn_cond" [1, S, 768] prompt tokens + seconds_total embed (S dynamic)
"cross_attn_mask" [1, S] bool
"global_embed" [1, 768] seconds_total embed
"local_add_cond" [1, 257, T] inpaint_mask (1ch) + inpaint_masked_input (256ch)
"padding_mask" [1, T] bool, True = valid
-> "v" [1, 256, T] rf_denoiser output
Also verifies whether the traced graph generalizes over T (dynamic_axes) by
running ORT at a different length; if that fails, the C++ side uses bucketed
static graphs + padding_mask instead.
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_dit.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")
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
# aten::rms_norm has no ONNX symbolic in the TS exporter — decompose to
# primitive ops (identical math; pow/mean/rsqrt export cleanly and TRT likes them).
import torch.nn.functional as F_patch
def _rms_norm_decomposed(input, normalized_shape, weight=None, eps=None):
if eps is None:
eps = torch.finfo(input.dtype).eps
dims = tuple(range(-len(normalized_shape), 0))
out = input * torch.rsqrt(input.pow(2).mean(dim=dims, keepdim=True) + eps)
if weight is not None:
out = out * weight
return out
F_patch.rms_norm = _rms_norm_decomposed
from stable_audio_3.model import StableAudioModel
T_TRACE = 1024 # latent frames used for tracing (~97s of audio)
T_ALT = 640 # different length to probe dynamic-shape generalization
S_TRACE = 257 # 256 prompt tokens + 1 seconds_total token
class DiTCore(nn.Module):
"""Flattens the conditioning dict interface to plain tensors around _forward."""
def __init__(self, dit):
super().__init__()
self.dit = dit # DiffusionTransformer
def forward(self, x, t, cross_attn_cond, cross_attn_mask, global_embed,
local_add_cond, padding_mask):
return self.dit._forward(
x, t,
cross_attn_cond=cross_attn_cond,
cross_attn_cond_mask=cross_attn_mask,
global_embed=global_embed,
local_add_cond=local_add_cond,
padding_mask=padding_mask,
)
def make_inputs(T, device, seed=0):
g = torch.Generator(device="cpu").manual_seed(seed)
x = torch.randn(1, 256, T, generator=g).to(device)
t = torch.tensor([0.3], dtype=torch.float32, device=device)
cross = torch.randn(1, S_TRACE, 768, generator=g).to(device)
cross_mask = torch.ones(1, S_TRACE, dtype=torch.bool, device=device)
glob = torch.randn(1, 768, generator=g).to(device)
local = torch.zeros(1, 257, T, device=device)
pad = torch.ones(1, T, dtype=torch.bool, device=device)
return (x, t, cross, cross_mask, glob, local, pad)
_ORT_DTYPES = {"tensor(float)": np.float32, "tensor(bool)": np.bool_, "tensor(int64)": np.int64}
INPUT_NAMES = ["x", "t", "cross_attn_cond", "cross_attn_mask",
"global_embed", "local_add_cond", "padding_mask"]
def run_ort(sess, inputs):
# Feed by NAME: the exporter prunes graph-unused inputs (e.g. cross_attn_mask —
# the model never forwards it; learned padding replaces masking), so positional
# zipping misaligns.
named = dict(zip(INPUT_NAMES, inputs))
feed = {}
for meta in sess.get_inputs():
arr = named[meta.name].cpu().numpy()
want = _ORT_DTYPES.get(meta.type)
if want is not None and arr.dtype != want:
arr = arr.astype(want)
feed[meta.name] = arr
return sess.run(None, feed)[0]
def compare(tag, out, ref):
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")
ap.add_argument("--parity-only", action="store_true",
help="Skip export; run parity against an existing sa3-dit.onnx")
args = ap.parse_args()
os.makedirs(args.output_dir, exist_ok=True)
print("Loading stable-audio-3-medium (fp32)...")
model = StableAudioModel.from_pretrained("medium", model_half=False)
dit = model.model.model.model # StableAudioModel -> CondWrapper -> DiTWrapper -> DiffusionTransformer
dit.eval().requires_grad_(False)
device = next(dit.parameters()).device
core = DiTCore(dit)
inputs = make_inputs(T_TRACE, device)
with torch.no_grad():
ref = core(*inputs)
print(f"Reference out: {tuple(ref.shape)}")
dit_path = os.path.join(args.output_dir, "sa3-dit.onnx")
if args.parity_only:
assert os.path.exists(dit_path), f"{dit_path} not found"
print("(--parity-only: skipping export)")
import onnxruntime as ort
sess = ort.InferenceSession(dit_path, providers=["CPUExecutionProvider"])
print("Parity at traced length:")
out = run_ort(sess, inputs)
c1 = compare(f"T={T_TRACE}", out, ref)
print("Dynamic-shape probe at different length:")
alt_inputs = make_inputs(T_ALT, device, seed=1)
with torch.no_grad():
alt_ref = core(*alt_inputs)
try:
alt_out = run_ort(sess, alt_inputs)
c2 = compare(f"T={T_ALT}", alt_out, alt_ref)
dynamic_ok = c2 > 0.999
except Exception as e:
print(f" [T={T_ALT}] FAILED to run: {type(e).__name__}: {str(e)[:300]}")
dynamic_ok = False
print(f"PARITY {'OK' if c1 > 0.999 else 'FAILED'}; DYNAMIC-T {'OK' if dynamic_ok else 'NOT SUPPORTED -> use bucketed static graphs'}")
return 0 if c1 > 0.999 else 1
t0 = time.time()
torch.onnx.export(
core, inputs, dit_path,
input_names=["x", "t", "cross_attn_cond", "cross_attn_mask",
"global_embed", "local_add_cond", "padding_mask"],
output_names=["v"],
dynamic_axes={
"x": {2: "T"}, "local_add_cond": {2: "T"}, "padding_mask": {1: "T"},
"cross_attn_cond": {1: "S"}, "cross_attn_mask": {1: "S"},
"v": {2: "T"},
},
opset_version=18, dynamo=False,
)
total = sum(os.path.getsize(os.path.join(args.output_dir, f))
for f in os.listdir(args.output_dir)
if f.startswith("sa3-dit"))
print(f"Exported {dit_path} ({total/1e9:.2f} GB incl. external data, {time.time()-t0:.0f}s)")
import onnxruntime as ort
sess = ort.InferenceSession(dit_path, providers=["CPUExecutionProvider"])
print("Parity at traced length:")
out = run_ort(sess, inputs)
c1 = compare(f"T={T_TRACE}", out, ref)
print("Dynamic-shape probe at different length:")
alt_inputs = make_inputs(T_ALT, device, seed=1)
with torch.no_grad():
alt_ref = core(*alt_inputs)
try:
alt_out = run_ort(sess, alt_inputs)
c2 = compare(f"T={T_ALT}", alt_out, alt_ref)
dynamic_ok = c2 > 0.999
except Exception as e:
print(f" [T={T_ALT}] FAILED to run: {type(e).__name__}: {str(e)[:300]}")
dynamic_ok = False
print(f"PARITY {'OK' if c1 > 0.999 else 'FAILED'}; DYNAMIC-T {'OK' if dynamic_ok else 'NOT SUPPORTED -> use bucketed static graphs'}")
return 0 if c1 > 0.999 else 1
if __name__ == "__main__":
sys.exit(main())
tools/onnx-export/export_sa3_same.py
+137
View File
@@ -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())
tools/onnx-export/export_text_enc.py
+287
View File
@@ -0,0 +1,287 @@
#!/usr/bin/env python3
"""
export_text_enc.py — Export Qwen3-Embedding text encoder to ONNX.
The text encoder is a standard Qwen3Model (28 layers, H=1024, causal attention)
that takes BPE token IDs and produces hidden states for the condition encoder.
Usage:
python export_text_enc.py --model-dir <path-to-Qwen3-Embedding-0.6B> --output <output.onnx>
Exports:
text_encoder.onnx — Full 28-layer transformer
Input: input_ids [B, S] int64
Output: hidden_states [B, S, 1024] fp16
embed_lookup.bin — Raw embedding table (vocab_size * hidden_size * 2 bytes, BF16)
Used for lyric token embedding lookup on CPU (no ONNX needed).
"""
import argparse
import os
import sys
import time
import struct
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
class TextEncoderWrapper(nn.Module):
"""Wrapper around Qwen3Model that returns hidden_states as a flat tensor.
ONNX inputs:
input_ids: [B, S] int64 — BPE token IDs
ONNX output:
hidden_states: [B, S, 1024] fp16 — last hidden state
"""
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, input_ids):
outputs = self.model(
input_ids=input_ids,
attention_mask=None, # causal mask generated internally
output_hidden_states=False,
return_dict=True,
)
return outputs.last_hidden_state
def load_model(model_dir: str, device: str = "cuda", dtype=torch.float32):
"""Load Qwen3-Embedding model from safetensors."""
model_dir = Path(model_dir)
# Fix Windows encoding issues
if sys.platform == "win32":
import io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
print(f"[export_text_enc] Loading model from {model_dir}...")
t0 = time.time()
from transformers import AutoModel, AutoConfig
config = AutoConfig.from_pretrained(str(model_dir))
# Force SDPA for ONNX export (no flash attention)
config._attn_implementation = "sdpa"
model = AutoModel.from_pretrained(
str(model_dir),
config=config,
torch_dtype=dtype,
trust_remote_code=True,
)
model = model.to(device)
model.eval()
t1 = time.time()
n_params = sum(p.numel() for p in model.parameters()) / 1e6
print(f"[export_text_enc] Model loaded in {t1-t0:.1f}s ({n_params:.0f}M params)")
print(f"[export_text_enc] Config: {config.num_hidden_layers}L, H={config.hidden_size}, "
f"heads={config.num_attention_heads}/{config.num_key_value_heads}")
return model, config
def export_onnx(model, config, output_path: str, opset: int = 18):
"""Export the text encoder to ONNX."""
device = next(model.parameters()).device
dtype = next(model.parameters()).dtype
wrapper = TextEncoderWrapper(model)
wrapper.eval()
# Dummy inputs for tracing
B = 1
S = 128 # typical sequence length
dummy_input_ids = torch.randint(0, config.vocab_size, (B, S), device=device, dtype=torch.long)
print(f"[export_text_enc] Tracing with shapes: input_ids={list(dummy_input_ids.shape)}")
# Test forward pass
print("[export_text_enc] Testing forward pass...")
with torch.no_grad():
test_out = wrapper(dummy_input_ids)
print(f"[export_text_enc] Output shape: {list(test_out.shape)} "
f"(expected [{B}, {S}, {config.hidden_size}])")
# Export to ONNX
print(f"[export_text_enc] Exporting to ONNX (opset {opset})...")
t0 = time.time()
torch.onnx.export(
wrapper,
(dummy_input_ids,),
output_path,
opset_version=opset,
input_names=["input_ids"],
output_names=["hidden_states"],
dynamic_axes={
"input_ids": {0: "batch", 1: "seq_len"},
"hidden_states": {0: "batch", 1: "seq_len"},
},
do_constant_folding=True,
export_params=True,
)
t1 = time.time()
file_size = os.path.getsize(output_path)
print(f"[export_text_enc] Exported to {output_path}")
print(f"[export_text_enc] File size: {file_size/1e6:.1f} MB")
print(f"[export_text_enc] Export time: {t1-t0:.1f}s")
return output_path
def export_embed_table(model, config, output_path: str):
"""Export the embedding table as a raw binary file for lyric lookup.
The lyric path uses embed_tokens lookup only (no transformer layers).
We export the table as float32 for direct CPU indexing.
Format: raw float32 array [vocab_size, hidden_size]
"""
embed_weight = model.embed_tokens.weight.detach().cpu().float().numpy()
V, H = embed_weight.shape
with open(output_path, "wb") as f:
# Header: vocab_size (int32), hidden_size (int32)
f.write(struct.pack("<II", V, H))
# Raw float32 weights
f.write(embed_weight.tobytes())
file_size = os.path.getsize(output_path)
print(f"[export_text_enc] Embedding table: [{V}, {H}] -> {output_path} ({file_size/1e6:.1f} MB)")
def export_null_cond(model_dir: str, output_path: str):
"""Export null_condition_emb from the DiT model as raw float32.
This is a [2048] float32 vector used for classifier-free guidance padding.
Read from the DiT safetensors since it lives there.
"""
from safetensors.torch import load_file
model_dir = Path(model_dir)
st_path = model_dir / "model.safetensors"
if not st_path.exists():
# Try multi-shard
for p in sorted(model_dir.glob("model-*.safetensors")):
st = load_file(str(p))
if "null_condition_emb" in st:
vec = st["null_condition_emb"].detach().cpu().float().numpy()
with open(output_path, "wb") as f:
f.write(struct.pack("<I", vec.shape[0]))
f.write(vec.tobytes())
print(f"[export_text_enc] null_condition_emb: [{vec.shape[0]}] -> {output_path}")
return
print("[export_text_enc] WARNING: null_condition_emb not found")
return
st = load_file(str(st_path))
if "null_condition_emb" not in st:
print("[export_text_enc] WARNING: null_condition_emb not found in model.safetensors")
return
vec = st["null_condition_emb"].detach().cpu().float().numpy()
with open(output_path, "wb") as f:
f.write(struct.pack("<I", vec.shape[0]))
f.write(vec.tobytes())
print(f"[export_text_enc] null_condition_emb: [{vec.shape[0]}] -> {output_path}")
def verify_onnx(onnx_path: str, model, config):
"""Verify ONNX output matches PyTorch."""
try:
import onnxruntime as ort
except ImportError:
print("[export_text_enc] onnxruntime not installed, skipping verification")
return
device = next(model.parameters()).device
wrapper = TextEncoderWrapper(model)
wrapper.eval()
# Test inputs
B, S = 1, 64
input_ids = torch.randint(0, config.vocab_size, (B, S), device=device, dtype=torch.long)
# PyTorch reference
with torch.no_grad():
ref_out = wrapper(input_ids).cpu().float().numpy()
# ONNX inference
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
sess = ort.InferenceSession(onnx_path, providers=providers)
ort_out = sess.run(None, {
"input_ids": input_ids.cpu().numpy(),
})[0]
# Compare
max_diff = np.max(np.abs(ref_out - ort_out))
mean_diff = np.mean(np.abs(ref_out - ort_out))
print(f"[export_text_enc] Verification: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}")
if max_diff < 0.05:
print("[export_text_enc] PASS: ONNX output matches PyTorch (within FP16 tolerance)")
else:
print("[export_text_enc] WARNING: Large difference — may need investigation")
def main():
parser = argparse.ArgumentParser(description="Export Qwen3-Embedding text encoder to ONNX")
parser.add_argument("--model-dir", required=True,
help="Path to Qwen3-Embedding-0.6B directory")
parser.add_argument("--output", default=None,
help="Output ONNX file (default: models/onnx/text_encoder.onnx)")
parser.add_argument("--dit-dir", default=None,
help="Path to DiT model dir (for null_condition_emb export)")
parser.add_argument("--opset", type=int, default=18,
help="ONNX opset version (default: 18)")
parser.add_argument("--verify", action="store_true",
help="Verify ONNX output matches PyTorch")
parser.add_argument("--device", default="cuda",
help="Device for model loading (default: cuda)")
args = parser.parse_args()
# Default output path
if args.output is None:
onnx_dir = Path(args.model_dir).parent / "onnx"
onnx_dir.mkdir(parents=True, exist_ok=True)
args.output = str(onnx_dir / "text_encoder.onnx")
os.makedirs(os.path.dirname(args.output), exist_ok=True)
output_dir = os.path.dirname(args.output)
# Load model
model, config = load_model(args.model_dir, device=args.device)
# Export ONNX
export_onnx(model, config, args.output, opset=args.opset)
# Export embedding table for lyric lookup
embed_path = os.path.join(output_dir, "embed_tokens.bin")
export_embed_table(model, config, embed_path)
# Export null_condition_emb if DiT dir provided
if args.dit_dir:
null_cond_path = os.path.join(output_dir, "null_condition_emb.bin")
export_null_cond(args.dit_dir, null_cond_path)
# Verify
if args.verify:
verify_onnx(args.output, model, config)
print("[export_text_enc] Done!")
if __name__ == "__main__":
main()
tools/onnx-export/export_vae.py
+261
View File
@@ -0,0 +1,261 @@
#!/usr/bin/env python3
"""Export AutoencoderOobleck VAE decoder to ONNX format.
Exports the decoder half of the VAE for use with TensorRT or ONNX Runtime.
The encoder is not needed for inference (we only decode latents → audio).
Tensor spec:
Input: "latents" [B, 64, T] (latent channels, latent frames @ 25Hz)
Output: "audio" [B, 2, samples] (stereo, samples = T * 1920 @ 48kHz)
"""
import argparse
import os
import sys
import time
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
class VAEDecoderWrapper(nn.Module):
"""Wraps AutoencoderOobleck.decoder + post_quant_conv to extract .sample.
The raw decoder returns a DecoderOutput namedtuple, which torch.onnx.export
can't trace cleanly. This wrapper calls the decoder and returns the raw
tensor directly.
"""
def __init__(self, vae):
super().__init__()
self.decoder = vae.decoder
# post_quant_conv maps from latent space back to decoder input space
if hasattr(vae, "post_quant_conv") and vae.post_quant_conv is not None:
self.post_quant_conv = vae.post_quant_conv
else:
self.post_quant_conv = None
def forward(self, latents: torch.Tensor) -> torch.Tensor:
if self.post_quant_conv is not None:
latents = self.post_quant_conv(latents)
decoded = self.decoder(latents)
# decoder returns DecoderOutput with .sample attribute
if hasattr(decoded, "sample"):
return decoded.sample
return decoded
def export_vae(vae_path: str, output_path: str, opset: int = 18) -> str:
"""Export VAE decoder to ONNX.
Args:
vae_path: Path to the VAE checkpoint directory (config.json + safetensors).
output_path: Path to write the ONNX file.
opset: ONNX opset version.
Returns:
The output path of the exported ONNX file.
"""
from diffusers import AutoencoderOobleck
print(f"Loading VAE from: {vae_path}")
vae = AutoencoderOobleck.from_pretrained(vae_path)
vae.eval()
wrapper = VAEDecoderWrapper(vae)
wrapper.eval()
# Move to CPU for export (fp32)
wrapper = wrapper.cpu()
# Create dummy input: [batch=1, latent_channels=64, latent_frames=250]
# 250 frames @ 25Hz = 10 seconds of audio
dummy_latents = torch.randn(1, 64, 250, dtype=torch.float32)
print(f"Dummy input shape: {dummy_latents.shape}")
print(f"Expected output shape: [1, 2, {250 * 1920}] = [1, 2, {250 * 1920}]")
# Test forward pass
with torch.no_grad():
test_out = wrapper(dummy_latents)
print(f"Test forward pass output shape: {test_out.shape}")
# Ensure output directory exists
os.makedirs(os.path.dirname(output_path), exist_ok=True)
# Export
print(f"\nExporting to ONNX (opset {opset})...")
t0 = time.time()
dynamic_axes = {
"latents": {0: "batch", 2: "latent_frames"},
"audio": {0: "batch", 2: "samples"},
}
torch.onnx.export(
wrapper,
(dummy_latents,),
output_path,
opset_version=opset,
input_names=["latents"],
output_names=["audio"],
dynamic_axes=dynamic_axes,
do_constant_folding=True,
dynamo=False, # Force legacy TorchScript exporter (dynamo hits cp1252 UnicodeError on Windows)
)
export_time = time.time() - t0
file_size_mb = os.path.getsize(output_path) / (1024 * 1024)
print(f"Export complete in {export_time:.1f}s")
print(f"Output file: {output_path}")
print(f"File size: {file_size_mb:.1f} MB")
# Validate with onnx checker
import onnx
print("\nValidating ONNX model...")
model = onnx.load(output_path)
onnx.checker.check_model(model, full_check=True)
print("ONNX checker: PASSED")
# Print model info
graph = model.graph
print(f"\nModel inputs:")
for inp in graph.input:
shape = [d.dim_param or d.dim_value for d in inp.type.tensor_type.shape.dim]
print(f" {inp.name}: {shape}")
print(f"Model outputs:")
for out in graph.output:
shape = [d.dim_param or d.dim_value for d in out.type.tensor_type.shape.dim]
print(f" {out.name}: {shape}")
return output_path
def validate_onnx(onnx_path: str, vae_path: str):
"""Compare ONNX Runtime output against PyTorch output.
Args:
onnx_path: Path to the exported ONNX file.
vae_path: Path to the VAE checkpoint directory.
"""
import onnxruntime as ort
from diffusers import AutoencoderOobleck
print("\n" + "=" * 60)
print("VALIDATION: Comparing ONNX vs PyTorch outputs")
print("=" * 60)
# Load PyTorch model
print("Loading PyTorch VAE...")
vae = AutoencoderOobleck.from_pretrained(vae_path)
vae.eval()
wrapper = VAEDecoderWrapper(vae)
wrapper.eval()
wrapper = wrapper.cpu()
# Create test input (shorter for speed: 50 frames = 2 seconds)
test_latents = torch.randn(1, 64, 50, dtype=torch.float32)
print(f"Test input shape: {test_latents.shape}")
# PyTorch inference
with torch.no_grad():
pt_output = wrapper(test_latents).numpy()
print(f"PyTorch output shape: {pt_output.shape}")
# ONNX Runtime inference
print("Loading ONNX model in onnxruntime...")
available_providers = ort.get_available_providers()
print(f"Available providers: {available_providers}")
# Use CUDA if available, else CPU
if "CUDAExecutionProvider" in available_providers:
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
print("Using: CUDAExecutionProvider")
else:
providers = ["CPUExecutionProvider"]
print("Using: CPUExecutionProvider (CUDA not available)")
sess = ort.InferenceSession(onnx_path, providers=providers)
ort_input = {"latents": test_latents.numpy()}
ort_output = sess.run(["audio"], ort_input)[0]
print(f"ONNX output shape: {ort_output.shape}")
# Compare
abs_diff = np.abs(pt_output - ort_output)
max_diff = abs_diff.max()
mean_diff = abs_diff.mean()
rel_diff = abs_diff / (np.abs(pt_output) + 1e-8)
print(f"\nDiff statistics:")
print(f" Max absolute diff: {max_diff:.6e}")
print(f" Mean absolute diff: {mean_diff:.6e}")
print(f" Max relative diff: {rel_diff.max():.6e}")
print(f" Mean relative diff: {rel_diff.mean():.6e}")
# Threshold check
if max_diff < 1e-4:
print("\n[PASS] VALIDATION PASSED: Outputs match within tolerance (1e-4)")
elif max_diff < 1e-3:
print("\n[WARN] VALIDATION WARNING: Small differences detected (< 1e-3)")
print(" This is acceptable for fp32 export, TRT fp16 will diverge more.")
else:
print(f"\n[FAIL] VALIDATION FAILED: Max diff {max_diff:.6e} exceeds tolerance")
sys.exit(1)
def main():
parser = argparse.ArgumentParser(
description="Export AutoencoderOobleck VAE decoder to ONNX"
)
parser.add_argument(
"--vae-path",
type=str,
required=True,
help="Path to VAE checkpoint directory (config.json + safetensors)",
)
parser.add_argument(
"--output",
type=str,
required=True,
help="Output path for the ONNX file",
)
parser.add_argument(
"--opset",
type=int,
default=18,
help="ONNX opset version (default: 18)",
)
parser.add_argument(
"--validate",
action="store_true",
help="Validate ONNX output against PyTorch output using onnxruntime",
)
args = parser.parse_args()
# Verify input exists
if not os.path.isdir(args.vae_path):
print(f"ERROR: VAE path not found: {args.vae_path}")
sys.exit(1)
config_path = os.path.join(args.vae_path, "config.json")
if not os.path.isfile(config_path):
print(f"ERROR: config.json not found in {args.vae_path}")
sys.exit(1)
# Export
onnx_path = export_vae(args.vae_path, args.output, args.opset)
# Validate
if args.validate:
validate_onnx(onnx_path, args.vae_path)
print("\nDone!")
if __name__ == "__main__":
main()
tools/onnx-export/export_vae_encoder.py
+279
View File
@@ -0,0 +1,279 @@
#!/usr/bin/env python3
"""Export AutoencoderOobleck VAE encoder to ONNX format.
Exports the encoder half of the VAE for use with TensorRT or ONNX Runtime.
The encoder converts audio → latent space for timbre/cover VAE encoding.
Tensor spec:
Input: "audio" [B, 2, samples] (stereo, samples @ 48kHz)
Output: "latents" [B, 64, T] (latent channels, latent frames @ 25Hz)
Note: The encoder output is 128ch (64 mean + 64 scale). We only need the mean
for deterministic encoding, so we slice to the first 64 channels.
"""
import argparse
import os
import sys
import time
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
class VAEEncoderWrapper(nn.Module):
"""Wraps AutoencoderOobleck.encoder + quant_conv to extract mean latents.
The raw encoder returns 128ch (mean + scale). For deterministic encoding
we only need the first 64 channels (mean). This wrapper handles the
quant_conv and slicing.
"""
def __init__(self, vae):
super().__init__()
self.encoder = vae.encoder
# quant_conv maps from encoder output space to latent space
if hasattr(vae, "quant_conv") and vae.quant_conv is not None:
self.quant_conv = vae.quant_conv
else:
self.quant_conv = None
def forward(self, audio: torch.Tensor) -> torch.Tensor:
"""
Args:
audio: [B, 2, samples] stereo audio at 48kHz
Returns:
latents: [B, 64, T] mean latents (deterministic)
"""
encoded = self.encoder(audio)
# encoder returns EncoderOutput with .latent_dist or raw tensor
if hasattr(encoded, "latent_dist"):
h = encoded.latent_dist.mean
elif hasattr(encoded, "sample"):
h = encoded.sample
else:
h = encoded
if self.quant_conv is not None:
h = self.quant_conv(h)
# h is [B, 128, T] — first 64 = mean, last 64 = log_var
# Only return mean for deterministic encoding
return h[:, :64, :]
def export_vae_encoder(vae_path: str, output_path: str, opset: int = 18) -> str:
"""Export VAE encoder to ONNX.
Args:
vae_path: Path to the VAE checkpoint directory (config.json + safetensors).
output_path: Path to write the ONNX file.
opset: ONNX opset version.
Returns:
The output path of the exported ONNX file.
"""
from diffusers import AutoencoderOobleck
print(f"Loading VAE from: {vae_path}")
vae = AutoencoderOobleck.from_pretrained(vae_path)
vae.eval()
wrapper = VAEEncoderWrapper(vae)
wrapper.eval()
# Move to CPU for export (fp32)
wrapper = wrapper.cpu()
# Create dummy input: [batch=1, channels=2, samples=480000]
# 480000 samples @ 48kHz = 10 seconds of audio
dummy_audio = torch.randn(1, 2, 480000, dtype=torch.float32)
print(f"Dummy input shape: {dummy_audio.shape}")
print(f"Expected output shape: [1, 64, {480000 // 1920}] = [1, 64, {480000 // 1920}]")
# Test forward pass
with torch.no_grad():
test_out = wrapper(dummy_audio)
print(f"Test forward pass output shape: {test_out.shape}")
# Ensure output directory exists
os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
# Export
print(f"\nExporting to ONNX (opset {opset})...")
t0 = time.time()
dynamic_axes = {
"audio": {0: "batch", 2: "samples"},
"latents": {0: "batch", 2: "latent_frames"},
}
torch.onnx.export(
wrapper,
(dummy_audio,),
output_path,
opset_version=opset,
input_names=["audio"],
output_names=["latents"],
dynamic_axes=dynamic_axes,
do_constant_folding=True,
dynamo=False, # Force legacy TorchScript exporter (dynamo hits cp1252 UnicodeError on Windows)
)
export_time = time.time() - t0
file_size_mb = os.path.getsize(output_path) / (1024 * 1024)
print(f"Export complete in {export_time:.1f}s")
print(f"Output file: {output_path}")
print(f"File size: {file_size_mb:.1f} MB")
# Validate with onnx checker
import onnx
print("\nValidating ONNX model...")
model = onnx.load(output_path)
onnx.checker.check_model(model, full_check=True)
print("ONNX checker: PASSED")
# Print model info
graph = model.graph
print(f"\nModel inputs:")
for inp in graph.input:
shape = [d.dim_param or d.dim_value for d in inp.type.tensor_type.shape.dim]
print(f" {inp.name}: {shape}")
print(f"Model outputs:")
for out in graph.output:
shape = [d.dim_param or d.dim_value for d in out.type.tensor_type.shape.dim]
print(f" {out.name}: {shape}")
return output_path
def validate_onnx(onnx_path: str, vae_path: str):
"""Compare ONNX Runtime output against PyTorch output.
Args:
onnx_path: Path to the exported ONNX file.
vae_path: Path to the VAE checkpoint directory.
"""
import onnxruntime as ort
from diffusers import AutoencoderOobleck
print("\n" + "=" * 60)
print("VALIDATION: Comparing ONNX vs PyTorch outputs")
print("=" * 60)
# Load PyTorch model
print("Loading PyTorch VAE...")
vae = AutoencoderOobleck.from_pretrained(vae_path)
vae.eval()
wrapper = VAEEncoderWrapper(vae)
wrapper.eval()
wrapper = wrapper.cpu()
# Create test input (shorter for speed: 96000 samples = 2 seconds)
test_audio = torch.randn(1, 2, 96000, dtype=torch.float32)
print(f"Test input shape: {test_audio.shape}")
# PyTorch inference
with torch.no_grad():
pt_output = wrapper(test_audio).numpy()
print(f"PyTorch output shape: {pt_output.shape}")
# ONNX Runtime inference
print("Loading ONNX model in onnxruntime...")
available_providers = ort.get_available_providers()
print(f"Available providers: {available_providers}")
# Use CUDA if available, else CPU
if "CUDAExecutionProvider" in available_providers:
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
print("Using: CUDAExecutionProvider")
else:
providers = ["CPUExecutionProvider"]
print("Using: CPUExecutionProvider (CUDA not available)")
sess = ort.InferenceSession(onnx_path, providers=providers)
ort_input = {"audio": test_audio.numpy()}
ort_output = sess.run(["latents"], ort_input)[0]
print(f"ONNX output shape: {ort_output.shape}")
# Compare
abs_diff = np.abs(pt_output - ort_output)
max_diff = abs_diff.max()
mean_diff = abs_diff.mean()
rel_diff = abs_diff / (np.abs(pt_output) + 1e-8)
print(f"\nDiff statistics:")
print(f" Max absolute diff: {max_diff:.6e}")
print(f" Mean absolute diff: {mean_diff:.6e}")
print(f" Max relative diff: {rel_diff.max():.6e}")
print(f" Mean relative diff: {rel_diff.mean():.6e}")
# Threshold check
if max_diff < 1e-4:
print("\n[PASS] VALIDATION PASSED: Outputs match within tolerance (1e-4)")
elif max_diff < 1e-3:
print("\n[WARN] VALIDATION WARNING: Small differences detected (< 1e-3)")
print(" This is acceptable for fp32 export, TRT fp16 will diverge more.")
else:
print(f"\n[FAIL] VALIDATION FAILED: Max diff {max_diff:.6e} exceeds tolerance")
sys.exit(1)
def main():
parser = argparse.ArgumentParser(
description="Export AutoencoderOobleck VAE encoder to ONNX"
)
parser.add_argument(
"--vae-path",
type=str,
required=True,
help="Path to VAE checkpoint directory (config.json + safetensors)",
)
parser.add_argument(
"--output",
type=str,
required=True,
help="Output path for the ONNX file",
)
parser.add_argument(
"--opset",
type=int,
default=18,
help="ONNX opset version (default: 18)",
)
parser.add_argument(
"--validate",
action="store_true",
help="Validate ONNX output against PyTorch output using onnxruntime",
)
args = parser.parse_args()
# Verify input exists
if not os.path.isdir(args.vae_path):
print(f"ERROR: VAE path not found: {args.vae_path}")
sys.exit(1)
config_path = os.path.join(args.vae_path, "config.json")
if not os.path.isfile(config_path):
print(f"ERROR: config.json not found in {args.vae_path}")
sys.exit(1)
# Export
onnx_path = export_vae_encoder(args.vae_path, args.output, args.opset)
# Validate
if args.validate:
validate_onnx(onnx_path, args.vae_path)
print("\nDone!")
if __name__ == "__main__":
main()
tools/onnx-export/gen_weight_map.py
+126
View File
@@ -0,0 +1,126 @@
"""
Generate weight_names.json sidecar for TRT adapter refit.
Maps ONNX val_N initializer names to human-readable parameter names.
Uses the deterministic linear_N numbering from dynamo decomposition.
Can be run standalone (no model loading needed, just the ONNX file).
"""
import onnx
import json
import sys
import os
def build_weight_map(onnx_path):
m = onnx.load(onnx_path, load_external_data=False)
# Build val_N -> MatMul node name mapping
val_to_node = {}
for node in m.graph.node:
if node.op_type == 'MatMul':
for inp in node.input:
if inp.startswith('val_'):
val_to_node[inp] = node.name
# Get all val_ weights sorted by N
val_inits = sorted(
[(i.name, list(i.dims)) for i in m.graph.initializer if i.name.startswith('val_')],
key=lambda x: int(x[0].split('_')[1])
)
# Filter to only MatMul weights (skip small constants)
matmul_weights = []
for vname, vshape in val_inits:
if vname in val_to_node:
numel = 1
for d in vshape:
numel *= d
if numel > 1000:
matmul_weights.append((vname, vshape, val_to_node[vname]))
# The first two MatMul nodes are proj_in and condition_embedder
# (before the per-layer linears start)
# node_MatMul_66 = proj_in.1.weight (Conv1d -> PatchEmbedLinear)
# node_MatMul_68 = condition_embedder.weight
# After that, the per-layer linears follow a repeating pattern.
# Each layer has 11 linear projections in this order:
LAYER_PATTERN = [
# (param_suffix, expected_shapes)
("self_attn.q_proj.weight", None),
("self_attn.k_proj.weight", None),
("self_attn.v_proj.weight", None),
("self_attn.o_proj.weight", None),
("cross_attn.q_proj.weight", None),
("cross_attn.k_proj.weight", None),
("cross_attn.v_proj.weight", None),
("cross_attn.o_proj.weight", None),
("mlp.gate_proj.weight", None),
("mlp.up_proj.weight", None),
("mlp.down_proj.weight", None),
]
# After the last layer, there should be a proj_out linear
rename_map = {} # val_N -> param_name
# Map the first two special cases
if len(matmul_weights) >= 2:
# proj_in
rename_map[matmul_weights[0][0]] = "dit.proj_in.1.linear.weight"
# condition_embedder
rename_map[matmul_weights[1][0]] = "dit.condition_embedder.weight"
# Map per-layer linears
layer_start = 2 # skip proj_in + condition_embedder
linears_per_layer = len(LAYER_PATTERN)
remaining = matmul_weights[layer_start:]
# Detect number of layers from count
# Last entry might be proj_out
num_layers = len(remaining) // linears_per_layer
leftover = len(remaining) % linears_per_layer
print(f"Total MatMul weights: {len(matmul_weights)}")
print(f"Layer weights: {len(remaining)} ({num_layers} layers * {linears_per_layer} + {leftover} extra)")
for layer_idx in range(num_layers):
for proj_idx, (suffix, _) in enumerate(LAYER_PATTERN):
w_idx = layer_start + layer_idx * linears_per_layer + proj_idx
if w_idx < len(matmul_weights):
vname = matmul_weights[w_idx][0]
param_name = f"dit.layers.{layer_idx}.{suffix}"
rename_map[vname] = param_name
# Map leftover (proj_out)
if leftover > 0:
proj_out_idx = layer_start + num_layers * linears_per_layer
if proj_out_idx < len(matmul_weights):
rename_map[matmul_weights[proj_out_idx][0]] = "dit.proj_out.1.inner.linear.weight"
# Build both directions
forward_map = rename_map # val_N -> param_name
reverse_map = {v: k for k, v in rename_map.items()} # param_name -> val_N
return {
"val_to_param": forward_map,
"param_to_val": reverse_map,
}
if __name__ == "__main__":
onnx_path = sys.argv[1] if len(sys.argv) > 1 else r'D:\Ace-Step-Latest\hot-step-cpp\models\onnx\dit_acestep-v15-merge-sft-turbo-xl-ta-0.7.onnx'
mapping = build_weight_map(onnx_path)
# Print summary
print(f"\nMapped {len(mapping['val_to_param'])} weights")
print("\nFirst 15 mappings:")
for val_name, param_name in sorted(mapping['val_to_param'].items(), key=lambda x: int(x[0].split('_')[1]))[:15]:
print(f" {val_name} -> {param_name}")
# Save
out_path = onnx_path + ".weight_names.json"
with open(out_path, 'w') as f:
json.dump(mapping, f, indent=2)
print(f"\nSaved to {out_path}")
tools/onnx-export/test_trt_vae.py
+269
View File
@@ -0,0 +1,269 @@
#!/usr/bin/env python3
"""Validate and benchmark VAE ONNX model with TensorRT Execution Provider.
Compares CUDA EP (baseline) against TensorRT EP for latency and numerical
accuracy. Caches TRT engines for subsequent runs.
Usage:
python test_trt_vae.py --onnx models/onnx/vae_decoder.onnx
"""
import argparse
import os
import sys
import time
from pathlib import Path
import numpy as np
def check_providers():
"""Check which ONNX Runtime execution providers are available."""
try:
import onnxruntime as ort
except ImportError:
print("ERROR: onnxruntime is not installed.")
print("Install with: pip install onnxruntime-gpu")
sys.exit(1)
available = ort.get_available_providers()
print(f"onnxruntime version: {ort.__version__}")
print(f"Available providers: {available}")
has_cuda = "CUDAExecutionProvider" in available
has_trt = "TensorrtExecutionProvider" in available
if not has_cuda:
print("\n[WARN] CUDAExecutionProvider is NOT available.")
print(" You likely have `onnxruntime` (CPU-only) instead of `onnxruntime-gpu`.")
print(" Install with: pip install onnxruntime-gpu")
print(" (You may need to uninstall onnxruntime first)")
# Check which package is installed
try:
import importlib.metadata
try:
ver = importlib.metadata.version("onnxruntime-gpu")
print(f" onnxruntime-gpu version: {ver}")
except importlib.metadata.PackageNotFoundError:
print(" onnxruntime-gpu: NOT installed")
try:
ver = importlib.metadata.version("onnxruntime")
print(f" onnxruntime (CPU): {ver}")
except importlib.metadata.PackageNotFoundError:
pass
except ImportError:
pass
return has_cuda, has_trt
def benchmark_session(sess, input_data, warmup=3, iterations=20):
"""Benchmark an ONNX Runtime session.
Args:
sess: ONNX Runtime InferenceSession.
input_data: Dict of input name → numpy array.
warmup: Number of warmup iterations.
iterations: Number of timed iterations.
Returns:
Tuple of (output_array, mean_latency_ms, std_latency_ms).
"""
# Warmup
for _ in range(warmup):
output = sess.run(None, input_data)
# Timed runs
latencies = []
for _ in range(iterations):
t0 = time.perf_counter()
output = sess.run(None, input_data)
latencies.append((time.perf_counter() - t0) * 1000)
latencies = np.array(latencies)
return output[0], latencies.mean(), latencies.std()
def run_benchmark(onnx_path: str, trt_cache_dir: str):
"""Run the full benchmark comparing CUDA EP vs TRT EP.
Args:
onnx_path: Path to the ONNX model file.
trt_cache_dir: Directory to cache TRT engines.
"""
import onnxruntime as ort
has_cuda, has_trt = check_providers()
if not has_cuda:
print("\nCannot run GPU benchmarks without CUDAExecutionProvider.")
print("Falling back to CPU-only test...")
run_cpu_test(onnx_path)
return
# Test input: 10 seconds of audio (250 latent frames)
test_latents = np.random.randn(1, 64, 250).astype(np.float32)
input_data = {"latents": test_latents}
print(f"\nTest input shape: {test_latents.shape}")
print(f"Expected output: [1, 2, {250 * 1920}] samples")
# ── CUDA EP benchmark ──
print("\n" + "=" * 60)
print("CUDA EP Benchmark")
print("=" * 60)
cuda_opts = ort.SessionOptions()
cuda_opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
try:
cuda_sess = ort.InferenceSession(
onnx_path,
sess_options=cuda_opts,
providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)
except Exception as e:
print(f"Failed to create CUDA session: {e}")
print("Falling back to CPU-only test...")
run_cpu_test(onnx_path)
return
cuda_output, cuda_mean, cuda_std = benchmark_session(cuda_sess, input_data)
print(f"Output shape: {cuda_output.shape}")
print(f"Latency: {cuda_mean:.2f} ± {cuda_std:.2f} ms")
# ── TensorRT EP benchmark ──
if not has_trt:
print("\n" + "=" * 60)
print("TensorRT EP: NOT AVAILABLE")
print("=" * 60)
print("TensorrtExecutionProvider is not available in this onnxruntime build.")
print("To enable TRT:")
print(" 1. Install onnxruntime-gpu with TRT support")
print(" 2. Ensure TensorRT libraries are on PATH")
print("\nSkipping TRT benchmark. CUDA EP results above are the baseline.")
return
print("\n" + "=" * 60)
print("TensorRT EP Benchmark")
print("=" * 60)
os.makedirs(trt_cache_dir, exist_ok=True)
trt_opts = ort.SessionOptions()
trt_opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
trt_provider_options = {
"trt_fp16_enable": True,
"trt_engine_cache_enable": True,
"trt_engine_cache_path": trt_cache_dir,
"trt_max_workspace_size": str(8 * 1024 * 1024 * 1024), # 8GB
}
print("Building TRT engine (first run may take minutes)...")
t0 = time.time()
try:
trt_sess = ort.InferenceSession(
onnx_path,
sess_options=trt_opts,
providers=[
("TensorrtExecutionProvider", trt_provider_options),
"CUDAExecutionProvider",
"CPUExecutionProvider",
],
)
except Exception as e:
print(f"Failed to create TRT session: {e}")
print("TRT EP may not be properly configured. Skipping TRT benchmark.")
return
engine_time = time.time() - t0
print(f"TRT engine ready in {engine_time:.1f}s")
trt_output, trt_mean, trt_std = benchmark_session(trt_sess, input_data)
print(f"Output shape: {trt_output.shape}")
print(f"Latency: {trt_mean:.2f} ± {trt_std:.2f} ms")
# ── Comparison ──
print("\n" + "=" * 60)
print("Comparison: CUDA EP vs TensorRT EP")
print("=" * 60)
abs_diff = np.abs(cuda_output - trt_output)
max_diff = abs_diff.max()
mean_diff = abs_diff.mean()
print(f"CUDA EP latency: {cuda_mean:.2f} ± {cuda_std:.2f} ms")
print(f"TRT EP latency: {trt_mean:.2f} ± {trt_std:.2f} ms")
print(f"Speedup: {cuda_mean / trt_mean:.2f}x")
print(f"Max abs diff: {max_diff:.6e}")
print(f"Mean abs diff: {mean_diff:.6e}")
if max_diff < 0.05:
print("\n[PASS] Numerical accuracy: GOOD (fp16 rounding is expected)")
elif max_diff < 0.5:
print("\n[WARN] Numerical accuracy: ACCEPTABLE (fp16 precision loss)")
else:
print(f"\n[FAIL] Numerical accuracy: POOR (max diff = {max_diff:.4f})")
print(" This may indicate a TRT conversion issue.")
def run_cpu_test(onnx_path: str):
"""Fallback: run a basic CPU test to verify the ONNX model loads."""
import onnxruntime as ort
print("\n" + "=" * 60)
print("CPU-only Test (fallback)")
print("=" * 60)
test_latents = np.random.randn(1, 64, 50).astype(np.float32)
input_data = {"latents": test_latents}
sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
output = sess.run(None, input_data)
print(f"Input shape: {test_latents.shape}")
print(f"Output shape: {output[0].shape}")
print(f"Output range: [{output[0].min():.4f}, {output[0].max():.4f}]")
print("[PASS] Model loads and runs on CPU successfully.")
def main():
parser = argparse.ArgumentParser(
description="Validate and benchmark VAE ONNX model with TensorRT EP"
)
parser.add_argument(
"--onnx",
type=str,
required=True,
help="Path to the exported ONNX file",
)
parser.add_argument(
"--trt-cache",
type=str,
default=None,
help="Directory for TRT engine cache (default: alongside ONNX file)",
)
parser.add_argument(
"--iterations",
type=int,
default=20,
help="Number of benchmark iterations (default: 20)",
)
args = parser.parse_args()
if not os.path.isfile(args.onnx):
print(f"ERROR: ONNX file not found: {args.onnx}")
sys.exit(1)
trt_cache = args.trt_cache
if trt_cache is None:
trt_cache = os.path.join(os.path.dirname(args.onnx), "trt_cache")
run_benchmark(args.onnx, trt_cache)
print("\nDone!")
if __name__ == "__main__":
main()