wai/tools: wai_jpegai_export_weights.py — reproducible weight export for the JPEG-AI decoder #27
2026-07-18
wai/tools: add wai_jpegai_export_weights.py — reproducible weight export for the decoder
The byte-exact JPEG-AI decoder in wai-rs is driven by weights lifted from the reference
ONNX assets, but the export step lived only in scratch and had to be rewritten each time
it was lost. This commits it.
Two regimes, and choosing wrongly yields plausible-but-wrong output:
ptq float nets (synthesis model_y/model_uv, hyper_decoder, MCM stages) -> 16-bit
per-tensor symmetric PTQ (conv -> int16 + mult=round(scale*2^31), shift=31;
bias/LayerNorm/temperature -> Q16), the layout Weights::from_json expects.
int hyper_scale_decoder -> verbatim. That net is integer end to end, so quantizing
it would corrupt an exactly-representable network; IntNet::from_json reads this.
Dependency-free: only the ONNX initializer table is needed, so it reads the protobuf
directly (no onnx/protobuf/torch import).
Verified to reproduce the weights the merged calibration tests were validated against:
synthesis_y 36 conv + 28 scalar (3,341,769 params, worst round-trip 6.6e-05), MCM stage0
3 conv (76,800 params, 6.9e-05), hyper_scale 9 integer tensors. Re-running
jpegai_hyperscale_calibration against freshly exported weights still reproduces the
reference scale indices 40960/40960 exactly.
Documents the eICCI caveat: post_filters/model_*.onnx export the model's `forward`,
which the decoder never calls, and their initializers are BN-folded and incomplete
(no BNY, scaleU/scaleV, or per-block scale) — take the reference state_dict instead.
Co-Authored-By: Claude Opus 4.8 <[email protected]>