wai-rs: jpegai_pipeline::reconstruct_yhat — the SGMM autoregressive spine #7
Loading…
Reference in a new issue
No description provided.
Delete branch "feat/jpegai-reconstruct"
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?
Rebuilds the synthesis-input latent y_hat from the dequantized residual + hyper params — the last large structural piece of the decode chain. Ports the reference context.pred loop: chunk psi into 4 per-stage params, down_shuffle the residual into a 2x2 checkerboard (order (0,0),(1,1),(0,1),(1,0) — diagonal pair first), run each MCM stage predicting a mean from the hyper params + already-reconstructed stages (context ASCENDING, y_0 first — pinned from the reference source, not the ONLX input label y[1-0]), add residual, up_shuffle back.
Reuses the byte-exact MCM stages. Validated on real test256.bin data: 76.83 dB SNR (max|Δ|=0.005) vs the reference yhat_y, carrying the MCM nets 16-bit PTQ floor (the float reference loop matches to ~139 dB). Unit test pins the shuffle round-trip + checkerboard order.
Decode chain now structurally complete: entropy → dequant → sigma → reconstruct → synthesis, each validated on real data.
Rebuilds the synthesis-input latent y_hat from the dequantized residual + the hyper params — the last large structural piece of the JPEG-AI decode chain. Ports the reference context.pred loop: p = chunk(psi, 4) # per-stage hyper params, ascending r = down_shuffle(resi) -> (k00,k11,k01,k10) # 2x2 checkerboard sub-lattices for i in 0..4: mu_i = MCM_i(p_i, cat(y_0 .. y_{i-1})) # ctx grows 0/c/2c/3c channels y_i = r_i + mu_i y_hat = up_shuffle(y_0 .. y_3) Reuses the byte-exact MCM stages (mcm_stage0 / mcm_stage_n). Two subtleties pinned from the reference source (not the ONNX): the checkerboard order is (0,0),(1,1), (0,1),(1,0) — the diagonal pair first — and the context concatenation is ASCENDING (y_0 first), despite the ONNX input being labelled y[1-0]; stage 0 has no context. Validated on real test256.bin data: reconstruct_yhat(resi_dq_y, psi_y) vs the reference yhat_y = 76.83 dB SNR (max|Δ| = 0.005), carrying the MCM nets' 16-bit PTQ floor (the float reference loop matches yhat_y to ~139 dB). Unit test pins the shuffle round-trip + checkerboard ordering. Co-Authored-By: Claude Opus 4.8 <[email protected]>dcharlot referenced this pull request2026-11-06 13:54:17 -05:00