Restore JPEG-AI decoders (uv/hyper/hyper_scale/mcm/mans) wiped by the main reset #1

Merged
dcharlot merged 3 commits from recovery/jpegai-decoders into main 2026-07-08 04:47:44 -04:00
Owner

These 5 byte-exact decoder PRs (#53/#54/#55 on the former GitHub canonical) were dropped from main during the migration to this Gitea. This branch restores them as a clean 3-commit fast-forward on top of #49. All validated: synthesis 58 dB, hyper 69 dB, hyper_scale + mANS entropy byte-exact.

These 5 byte-exact decoder PRs (#53/#54/#55 on the former GitHub canonical) were dropped from main during the migration to this Gitea. This branch restores them as a clean 3-commit fast-forward on top of #49. All validated: synthesis 58 dB, hyper 69 dB, hyper_scale + mANS entropy byte-exact.
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]>
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]>
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]>
Sign in to join this conversation.
No reviewers
No labels
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set

Reference
Transaction-Science/open-standards!1
No description provided.