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