Restore JPEG-AI decoders (uv/hyper/hyper_scale/mcm/mans) wiped by the main reset #1
2026-07-07
wai-rs: byte-exact mANS entropy decoder (jpegai_mans) — sgm + fac paths (#55)
The entropy stage of the JPEG-AI decode path: a faithful pure-Rust port of the ISO
reference mANS decoder (cpp_exts/mans/decompressor.cpp). This is the central
non-neural piece — it turns the compressed bitstream into the latent symbols the
synthesis/hyper nets consume.
Mechanism (byte-for-byte with the reference):
• two interleaved ANS states (state1/state2), advanced in a group-of-4
inbound → flush → outbound schedule;
• a rolling 64-bit window read MSB-first (read/flush), consumed backward through
the byte buffer, with a sentinel-bit init (initBackward + initWithStates);
• sgm path (main y-latent): a per-position sigma index selects one of 32
precomputed transition tables; mask=false skips the position (value 0) without
consuming the stream; the outbound/escape path (symbol == -bound) reads a
sign + 3/16-bit magnitude for out-of-range values;
• fac path (z hyper-latent): a per-channel cumulative CDF builds the row table
(transitionsZPreprocess), symbol == MAX_Z(63) escapes to a raw 6-bit value.
State threads across calls, so one MansDecoder decodes a whole stream.
Validated BYTE-EXACT against the ISO reference (the golden was produced by the
reference `ans.so`): decode_sgm 40960/40960 exact on the real SORP stream (160
calls, 703 coded symbols, two states, escapes), decode_factorize 128/128 exact on
an oracle encode→decode roundtrip. Tests feature-gated on the fixtures.
Co-authored-by: Transaction Science <[email protected]>
Co-authored-by: Claude Opus 4.8 <[email protected]>
wai-rs: jpegai_synthesis — add MCM context-model decoders (mcm_stage0/mcm_stage_n) (#54)
Adds the 4-stage Multi-stage Context Model (common_modules/MCM) to the JPEG-AI
integer decode path. Each stage refines the entropy context, conditioned on the
stage's prior `p_i` and the accumulated decoded context `y[<i]`:
• stage0 (no context): p0 -> conv1 -> ReLU -> conv2 -> ReLU -> conv3 -> y[0]
(bias-free 1x1, width 160)
• stage i>=1: cy = conv.1(conv.0(ctx)); Concat(cy, p_i) [->320]
-> fusion.conv1 -> ReLU -> conv2 -> ReLU -> conv3 -> y[i]
(conv.0 1x1 ctx_ch->160, conv.1 3x3 grouped-5 +bias; ctx = concat(y[i-1..0]))
Float nets -> Q16 PTQ (mcm_stage0/mcm_stage_n reuse the shared Weights loader +
conv path). Each stage validated independently against the float ONNX via ort:
SNR 73.86 / 68.94 / 61.98 / 61.18 dB (stages 0..3) — within the 16-bit PTQ.
Structural test on the real per-stage weights + tests/jpegai_mcm_calibration.rs.
Co-authored-by: Transaction Science <[email protected]>
Co-authored-by: Claude Opus 4.8 <[email protected]>
wai-rs: jpegai_synthesis — restore chroma + hyper decoders, add hyper_scale (all byte-exact) (#53)
Consolidated re-land of the JPEG-AI integer decode work that was rolled off main (#51/#52 were merged then force-pushed out by a concurrent main rewrite), plus the new hyper_scale_decoder. On top of the synthesis-Y decoder (#49): • decode_uv — chroma synthesis (model_uv): conv1(256->48) ‖ Slice(last96) -> Concat(144) -> conv2_t; TABs at full res; conv4_t(->8) -> DepthToSpace(CRD,2) -> 2-ch UV. 58.76 dB PSNR vs float ONNX. • decode_hyper — entropy hyper-prior decoder (common_modules): conv chain conv1(1x1)->conv2_t(4x4,up2)->ReLU6->conv3->ReLU6->conv4(->640). 69.55 dB SNR. • decode_hyper_scale — scale-index decoder (common_modules): a FULLY-INTEGER net (int64 symbols -> int32 scale indices). IntNet loader + exact integer conv + per-channel divide-by-2^n (truncation toward zero, not arithmetic shift) + QDQ clip + DepthToSpace(CRD,4) + abs + clip(0,3967). BYTE-EXACT (max|Δ| = 0) vs the onnx.reference evaluator — ORT rejects the model (no int64 Clip kernel), so the reference is the spec-compliant Python evaluator. Pure i64 => byte-identical on every machine. Structural tests on the real weights + per-net calibration tests (feature-gated on assets). Header + module docs updated. Co-authored-by: Transaction Science <[email protected]> Co-authored-by: Claude Opus 4.8 <[email protected]>