#!/usr/bin/env python3 """Export a `.pt` checkpoint to the published ONNX artifacts. Run at **publish time**, from the checkpoint being published, so the artifacts cannot drift from it. See `docs/onnx.md` for why this exists and what each precision costs; the short version is that `onnxruntime` is 27 MB against torch's 336 MB, and the receiving path only ever needs two convolutional passes. Six artifacts per checkpoint -- {encoder, decoder} x {fp32, fp16, int8} -- named after the checkpoint they came from: v1.pt -> v1-encoder-fp32.onnx v1-decoder-fp32.onnx v1-encoder-fp16.onnx v1-decoder-fp16.onnx v1-encoder-int8.onnx v1-decoder-int8.onnx All three precisions are published with every revision because every precision decodes on every other precision's receiver (measured: fp32 ONNX and torch agree to ~1e-6, roughly 112 dB below the channel noise). Publishing them saves third parties from rolling their own export, which is the case that actually risks divergence. **There is one on-air format, and the precisions are not variants of it.** Usage: scripts/export_onnx.py # published checkpoint scripts/export_onnx.py --model out/best.pt scripts/export_onnx.py --push # upload to the Hub Nothing is uploaded without `--push`. Every artifact is verified against the torch model before it is written, and a tolerance breach fails the run rather than publishing a bad codec. """ import argparse import hashlib import json import shutil import sys import tempfile from pathlib import Path import numpy as np import torch sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from sstvae import checkpoint # noqa: E402 from sstvae.codec import load_torch_model # noqa: E402 from sstvae.config import LATENT_CHANNELS, LATENT_H, LATENT_W # noqa: E402 from sstvae.images import IMG_H, IMG_W # noqa: E402 OPSET = 17 PRECISIONS = ("fp32", "fp16", "int8") # The yardstick for quantisation error is not fp32, it is the channel: # at the modem's operating point there is this much RMS noise on # unit-RMS latents. Quantisation noise is just one more small additive # source on a channel that already carries a much larger one. CHANNEL_NOISE_RMS = 0.367 # Gates, not targets. Measured values (docs/onnx.md) sit far below these; # they exist to catch a broken export, not to police the last decimal. # # latent_rms RMS encoder error on unit-RMS latents. Compare against # CHANNEL_NOISE_RMS -- that is the yardstick, not fp32. # quality_db PSNR *lost* against the source image, running the whole # pipeline at this precision versus running it in torch. # # The second one is the metric that means something. PSNR of an ONNX # reconstruction against the torch reconstruction is a difference with no # natural scale: int8 scores ~24 dB on it while costing ~0.15 dB of # actual picture quality, so gating on it would reject a good artifact. # # The int8 gates stay loose enough to survive an untuned export failing # gracefully, but the tuned one is far inside them: latent RMS 7.31e-02, # costing 0.002 dB on photographs and 0.112 dB off-distribution. # `quantize_dynamic` turns every Conv into `ConvInteger`, which supports # only a **per-tensor** weight scale -- `per_channel=True` is silently a # no-op on this graph, verified -- which is why the fix is per-layer # exclusion rather than a quantiser setting. TOLERANCES = { "fp32": {"latent_rms": 1e-4, "quality_db": 0.02}, "fp16": {"latent_rms": 5e-3, "quality_db": 0.05}, "int8": {"latent_rms": 0.25, "quality_db": 0.50}, } def sha256(path: Path) -> str: h = hashlib.sha256() with open(path, "rb") as f: for chunk in iter(lambda: f.read(1 << 20), b""): h.update(chunk) return h.hexdigest() def probe_images(n: int, seed: int = 0) -> torch.Tensor: """Deterministic stand-in images for verification, in [0,1]. Low-frequency content rather than white noise: a convolutional autoencoder trained on photographs behaves quite differently on broadband noise, and a verification input that the model has no idea what to do with tells you little about whether the export is faithful. Pass `--images` for the real thing when it matters. """ g = torch.Generator().manual_seed(seed) yy = torch.linspace(0, 1, IMG_H)[:, None] xx = torch.linspace(0, 1, IMG_W)[None, :] out = [] for _ in range(n): img = torch.zeros(3, IMG_H, IMG_W) for c in range(3): acc = torch.zeros(IMG_H, IMG_W) for _ in range(6): fx, fy = torch.rand(2, generator=g) * 6 ph = torch.rand(1, generator=g) * 6.283 amp = torch.rand(1, generator=g) acc += amp * torch.sin(6.283 * (fx * xx + fy * yy) + ph) img[c] = acc img = (img - img.amin()) / (img.amax() - img.amin() + 1e-8) out.append(img) return torch.stack(out) def load_images(paths: list[Path]) -> torch.Tensor: from PIL import Image from sstvae.images import fit_image, image_to_tensor return torch.stack([image_to_tensor(fit_image(Image.open(p))) for p in paths]) def export_fp32(module: torch.nn.Module, args: tuple, path: Path, input_names: list[str], output_names: list[str]) -> None: """Export one module at fp32. `dynamo=True` is the default and the maintained path; the legacy TorchScript exporter produces a numerically identical graph, so there is no reason to opt out. `external_data=False` is **not** the default and matters: without it the weights land in a `.onnx.data` sidecar, and four artifacts that must arrive together are worse than two. """ torch.onnx.export( module, args, str(path), input_names=input_names, output_names=output_names, opset_version=OPSET, dynamo=True, external_data=False, ) def convert_fp16(src: Path, dst: Path) -> None: import onnx from onnxconverter_common import float16 model = onnx.load(str(src)) # keep_io_types: the caller still hands us fp32 arrays and gets fp32 # back. The precision is an internal storage choice, not part of the # interface -- codec.py should not need to know which file it loaded. onnx.save(float16.convert_float_to_float16(model, keep_io_types=True), str(dst)) # int8 tuning: greedily keep the most quantisation-sensitive conv layers # at fp32, because `ConvInteger` carries one scale per weight *tensor* # and a single outlier-heavy tensor therefore sets a coarse scale for # everything in it. # # **Tune on off-distribution content too, or this measures nothing.** # Quantisation sensitivity barely shows on the photographs the model was # trained on and shows enormously on everything else. Measured on v1's # decoder: fully quantised costs 0.10 dB on COCO but **1.54 dB** on # smooth synthetic probes, and excluding one layer takes the latter to # -0.14 dB. Tuned against COCO alone the search sees 0.10 dB, concludes # there is nothing to fix, and ships the 1.54 dB. That is not an # academic case: operators send test cards, charts and screenshots. INT8_TUNE_SYNTHETIC_PROBES = 2 # Stop once the error is under target -- **absolute, not a # relative-improvement rule**. At the tail a large *fractional* error # reduction is worth 0.00 dB of picture, which is how an earlier # relative-only rule talked itself into three decoder exclusions that # together bought 0.005 dB. `INT8_MIN_GAIN` remains only as a guard # against grinding away at a target that cannot be reached. # # The encoder is measured in latent RMS against the channel noise, since # that is the quantity that goes on the air. The decoder is measured in # dB of picture, since its error reaches nobody else. INT8_TARGET_DB_UNDER_CHANNEL = 13.0 # A tenth of a dB of your own picture. Tighter than this (0.05 was tried) # buys a second decoder exclusion worth 0.02 dB for +1.8 MB, which is the # wrong trade for the precision that exists for constrained devices. INT8_TARGET_DECODER_DB = 0.10 INT8_MIN_GAIN = 0.15 INT8_MAX_EXCLUDED = 3 def convert_int8(src: Path, dst: Path, exclude: list[str] | None = None) -> None: from onnxruntime.quantization import QuantType, quantize_dynamic quantize_dynamic(str(src), str(dst), weight_type=QuantType.QInt8, nodes_to_exclude=list(exclude or [])) def _conv_nodes(path: Path) -> list[str]: import onnx return [n.name for n in onnx.load(str(path)).graph.node if n.op_type in ("Conv", "ConvTranspose")] def tune_int8_exclusions(src: Path, measure, target: float, *, log=print) -> list[str]: """Find which conv layers to leave at fp32, by measuring. Quantisation error is not spread evenly across layers -- it is dominated by a few whose weight distribution has outliers, because `quantize_dynamic` emits `ConvInteger`, which carries **one scale per tensor**. A single bad tensor therefore sets a coarse scale for all of its weights. (This is also why `per_channel=True` does nothing here: ConvInteger has nowhere to put per-channel scales.) Measured on v1: excluding one 0.59 MB layer took the encoder from 1.88e-01 to 7.31e-02 RMS -- 5.8 dB to 14.0 dB under the channel noise -- for 8% more file. That is worth finding automatically rather than hardcoding, because the sensitive layer is a property of the trained weights and will move between revisions. `measure(path) -> float` returns an error, lower being better, and `target` is the value below which the error stops mattering. """ candidates = _conv_nodes(src) excluded: list[str] = [] with tempfile.TemporaryDirectory() as td: def err(excl: list[str]) -> float: p = Path(td) / f"probe{len(excl)}.onnx" convert_int8(src, p, excl) return measure(p) current = err([]) log(f" all quantised: {current:.4e} (target {target:.4e})") while current > target and len(excluded) < INT8_MAX_EXCLUDED: scored = [(err(excluded + [c]), c) for c in candidates if c not in excluded] if not scored: break best_err, best = min(scored) gain = (current - best_err) / current if current else 0.0 if gain < INT8_MIN_GAIN: log(f" stop: best remaining gain {gain:.1%} < " f"{INT8_MIN_GAIN:.0%}") break excluded.append(best) current = best_err log(f" keep {best} at fp32 -> {current:.4e} ({gain:.0%} better)") if current <= target: log(f" at target ({current:.4e} <= {target:.4e})") return excluded def stamp_metadata(path: Path, props: dict) -> None: """Record provenance inside the artifact. Which checkpoint an `.onnx` came from is exactly the question that gets asked when two stations disagree, and an answer that lives in the file cannot be separated from it. """ import onnx model = onnx.load(str(path)) for k, v in props.items(): entry = model.metadata_props.add() entry.key, entry.value = str(k), str(v) onnx.save(model, str(path)) def run_onnx(path: Path, feeds: dict) -> np.ndarray: import onnxruntime as ort opts = ort.SessionOptions() opts.intra_op_num_threads = 4 # measured best; 1 was ~5x worse opts.log_severity_level = 3 sess = ort.InferenceSession(str(path), opts, providers=["CPUExecutionProvider"]) name = sess.get_inputs()[0].name if len(sess.get_inputs()) == 1: return sess.run(None, {name: feeds["primary"]})[0] second = sess.get_inputs()[1].name return sess.run(None, {name: feeds["primary"], second: feeds["secondary"]})[0] def psnr(a: np.ndarray, b: np.ndarray) -> float: mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)) return float("inf") if mse == 0 else 10.0 * float(np.log10(1.0 / mse)) def verify(model, images: torch.Tensor, paths: dict, precision: str) -> dict: """Compare one precision's encoder and decoder against torch. One image at a time: the graphs are exported at a fixed batch of 1, which is what the app actually does and what keeps the export in the easy, fully-static regime. """ lat_errs, vs_torch, torch_q, onnx_q = [], [], [], [] for i in range(images.shape[0]): one = images[i : i + 1] src = one.numpy().astype(np.float32) with torch.no_grad(): ref_latents = model.encoder(one) weights = torch.ones_like(ref_latents) ref_image = model.decoder(ref_latents, weights).numpy() got_latents = run_onnx(paths["encoder"], {"primary": src}) lat_errs.append(got_latents - ref_latents.numpy()) w_np = weights.numpy().astype(np.float32) # Decoder alone, from the *torch* latents: isolates decoder error # instead of compounding the encoder's. vs_torch.append(psnr( run_onnx(paths["decoder"], {"primary": ref_latents.numpy().astype(np.float32), "secondary": w_np}), ref_image, )) # Whole pipeline at this precision -- what a station running these # artifacts actually gets -- measured against the source image. onnx_q.append(psnr( run_onnx(paths["decoder"], {"primary": got_latents.astype(np.float32), "secondary": w_np}), src, )) torch_q.append(psnr(ref_image, src)) lat_err = np.concatenate([e.ravel() for e in lat_errs]) latent_rms = float(np.sqrt(np.mean(lat_err ** 2))) quality_db = float(np.mean(torch_q) - np.mean(onnx_q)) tol = TOLERANCES[precision] return { "precision": precision, "latent_rms": latent_rms, "latent_max": float(np.abs(lat_err).max()), "latent_vs_channel_db": ( float("inf") if latent_rms == 0 else 20.0 * float(np.log10(CHANNEL_NOISE_RMS / latent_rms)) ), "torch_psnr_db": float(np.mean(torch_q)), "onnx_psnr_db": float(np.mean(onnx_q)), "quality_lost_db": quality_db, "decoder_vs_torch_psnr_db": float(np.mean(vs_torch)), "encoder_mb": paths["encoder"].stat().st_size / 1e6, "decoder_mb": paths["decoder"].stat().st_size / 1e6, "ok": latent_rms <= tol["latent_rms"] and quality_db <= tol["quality_db"], } def main() -> int: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--model", default=None, help="checkpoint .pt; defaults to the published one") ap.add_argument("--out", default="onnx", type=Path, help="output directory (default: ./onnx)") ap.add_argument("--stem", default=None, help="artifact name prefix; defaults to the checkpoint stem " "(v1.pt -> v1-encoder-fp16.onnx)") ap.add_argument("--precisions", default=",".join(PRECISIONS), help=f"comma-separated subset of {','.join(PRECISIONS)}") ap.add_argument("--images", type=Path, default=None, help="directory of images to verify against; synthetic " "low-frequency probes are used if omitted") ap.add_argument("--n-probe", type=int, default=4, help="number of verification images (default 4)") ap.add_argument("--no-int8-tuning", action="store_true", help="skip the per-layer int8 sensitivity search (faster, " "but a worse int8 encoder -- see docs/onnx.md)") ap.add_argument("--push", action="store_true", help="upload the artifacts to the Hub after verifying") ap.add_argument("--repo", default=checkpoint.DEFAULT_REPO, help=f"Hub repo to push to (default {checkpoint.DEFAULT_REPO})") args = ap.parse_args() precisions = [p.strip() for p in args.precisions.split(",") if p.strip()] for p in precisions: if p not in PRECISIONS: ap.error(f"unknown precision {p!r}; choose from {', '.join(PRECISIONS)}") ckpt_path = Path(checkpoint.resolve(args.model)) stem = args.stem or ckpt_path.stem ckpt_sha = sha256(ckpt_path) print(f"checkpoint {ckpt_path.name} sha256:{ckpt_sha[:16]}...") model = load_torch_model(args.model) model.eval() if args.images: files = sorted( p for p in args.images.iterdir() if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp"} )[: args.n_probe] if not files: ap.error(f"no images found in {args.images}") images = load_images(files) print(f"verifying against {len(files)} image(s) from {args.images}") else: images = probe_images(args.n_probe) print(f"verifying against {args.n_probe} synthetic probe(s)") args.out.mkdir(parents=True, exist_ok=True) written: list[Path] = [] results = [] with tempfile.TemporaryDirectory() as tmp: tmp = Path(tmp) base = { "encoder": tmp / "encoder-fp32.onnx", "decoder": tmp / "decoder-fp32.onnx", } print("exporting fp32 graphs...") export_fp32(model.encoder, (images[:1],), base["encoder"], ["image"], ["latents"]) z0 = torch.zeros(1, LATENT_CHANNELS, LATENT_H, LATENT_W) export_fp32(model.decoder, (z0, torch.ones_like(z0)), base["decoder"], ["latents", "weights"], ["image"]) exclusions = {"encoder": [], "decoder": []} if "int8" in precisions and not args.no_int8_tuning: # A subset is plenty to *rank* layers -- the full set still # verifies the result afterwards -- but it must include # off-distribution probes. See INT8_TUNE_SYNTHETIC_PROBES. tune_imgs = torch.cat([ images[:min(3, images.shape[0])], probe_images(INT8_TUNE_SYNTHETIC_PROBES, seed=1234), ]) tune_n = tune_imgs.shape[0] src = tune_imgs.numpy() with torch.no_grad(): ref_z = torch.cat([model.encoder(tune_imgs[i:i + 1]) for i in range(tune_n)]) ref_pic = torch.cat([ model.decoder(ref_z[i:i + 1], torch.ones_like(ref_z[i:i + 1])) for i in range(tune_n)]).numpy() ref_z = ref_z.numpy() torch_q = float(np.mean([psnr(ref_pic[i], src[i]) for i in range(tune_n)])) def enc_err(path: Path) -> float: got = np.concatenate([ run_onnx(path, {"primary": src[i:i + 1]}) for i in range(tune_n)]) return float(np.sqrt(np.mean((got - ref_z) ** 2))) def dec_err(path: Path) -> float: """dB of picture lost, decoding the *torch* latents.""" got = np.concatenate([ run_onnx(path, {"primary": ref_z[i:i + 1], "secondary": np.ones_like(ref_z[i:i + 1])}) for i in range(tune_n)]) q = float(np.mean([psnr(got[i], src[i]) for i in range(tune_n)])) return max(torch_q - q, 0.0) plan = { "encoder": (enc_err, CHANNEL_NOISE_RMS / 10 ** (INT8_TARGET_DB_UNDER_CHANNEL / 20)), "decoder": (dec_err, INT8_TARGET_DECODER_DB), } for part, (measure, target) in plan.items(): print(f" tuning int8 {part} " f"({tune_n} probes, {INT8_TUNE_SYNTHETIC_PROBES} synthetic)...") exclusions[part] = tune_int8_exclusions( base[part], measure, target, log=print) for precision in precisions: paths = {} for part in ("encoder", "decoder"): dst = args.out / f"{stem}-{part}-{precision}.onnx" if precision == "fp32": shutil.copyfile(base[part], dst) elif precision == "fp16": convert_fp16(base[part], dst) else: convert_int8(base[part], dst, exclusions[part]) props = { "sstvae.source_checkpoint": ckpt_path.name, "sstvae.source_sha256": ckpt_sha, "sstvae.part": part, "sstvae.precision": precision, "sstvae.opset": OPSET, "sstvae.torch_version": torch.__version__, } if precision == "int8" and exclusions[part]: props["sstvae.int8_fp32_layers"] = ",".join(exclusions[part]) stamp_metadata(dst, props) paths[part] = dst r = verify(model, images, paths, precision) results.append(r) written.extend(paths.values()) flag = "ok" if r["ok"] else "FAIL" print( f" {precision:>4} {r['encoder_mb']:5.1f} + {r['decoder_mb']:4.1f} MB " f"latent RMS {r['latent_rms']:.2e} " f"({r['latent_vs_channel_db']:5.1f} dB under channel) " f"picture {r['onnx_psnr_db']:6.2f} dB " f"({r['quality_lost_db']:+.3f} vs torch) [{flag}]" ) manifest = args.out / f"{stem}-onnx-manifest.json" manifest.write_text(json.dumps({ "source_checkpoint": ckpt_path.name, "source_sha256": ckpt_sha, "opset": OPSET, "torch_version": torch.__version__, "channel_noise_rms": CHANNEL_NOISE_RMS, "int8_fp32_layers": exclusions, "artifacts": { p.name: {"sha256": sha256(p), "bytes": p.stat().st_size} for p in written }, "verification": results, }, indent=2) + "\n") print(f"wrote {len(written)} artifact(s) + manifest to {args.out}/") failed = [r["precision"] for r in results if not r["ok"]] if failed: print(f"\nFAILED verification: {', '.join(failed)} -- nothing pushed", file=sys.stderr) return 1 if args.push: from huggingface_hub import HfApi api = HfApi() api.create_repo(args.repo, exist_ok=True, repo_type="model") for p in [*written, manifest]: print(f" uploading {p.name}") api.upload_file(path_or_fileobj=str(p), path_in_repo=p.name, repo_id=args.repo, commit_message=f"ONNX artifacts for {ckpt_path.name}") print(f"pushed to https://huggingface.co/{args.repo}") else: print("(not pushed; pass --push to upload)") return 0 if __name__ == "__main__": raise SystemExit(main())