wai/tools: wai_jpegai_export_weights.py — reproducible weight export for the JPEG-AI decoder #27
Loading…
Reference in a new issue
No description provided.
Delete branch "feat/jpegai-weight-tools"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
The byte-exact decoder is driven by weights lifted from the reference ONNX assets, but the export step lived only in scratch and had to be rewritten whenever that was lost. This commits it as a first-class tool.
Two regimes — choosing wrongly yields plausible-but-wrong output:
ptq— float nets (synthesis y/uv, hyper_decoder, MCM) → 16-bit symmetric PTQ (conv → int16 +mult=round(scale*2^31),shift=31; bias/LN/temperature → Q16). WhatWeights::from_jsonexpects.int—hyper_scale_decoder→ verbatim. That net is integer end to end; quantizing it would corrupt an exactly-representable network. WhatIntNet::from_jsonexpects.Dependency-free — only the initializer table is needed, so it reads the ONNX protobuf directly (no onnx/protobuf/torch).
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 (76,800, 6.9e-05), hyper_scale 9 integer tensors. Re-running
jpegai_hyperscale_calibrationon freshly exported weights still gives 40960/40960 exact.Also documents the eICCI caveat:
post_filters/model_*.onnxexport the modelforward, which the decoder never calls, and their initializers are BN-folded and incomplete — take the referencestate_dictinstead.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]>