From a065082807d176c5292a4a38e6cbeb5785a171eb Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Mon, 24 Aug 2026 23:18:11 -0400 Subject: [PATCH] =?UTF-8?q?wai-web:=20full=20JPEG-AI=20codestream=E2=86=92?= =?UTF-8?q?RGB=20decode=20in=20wasm=20(decodeJpegAiFull)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The whole byte-exact decoder (wai::jpegai_decode::decode) as a browser entry: 13 byte inputs (codestream + the model weight bundle from wai_jpegai_build_e2e_inputs.py + norm tables), returns the decoded RGB. Same integer decoder a native sink runs, compiled to wasm32. Verified in node: decodes test256 to 256x256 in ~8s and the RGB is byte-identical to the native decode (sha256 dbc494a5…, the same digest the reconstruction-binding test produced natively) — decode-equivalence across host and browser. The JS side passes raw bytes; the wasm parses, so there is no JS-side double-parse. Drives the full-image path of the /jpegai demo. Co-Authored-By: Claude Opus 4.8 --- wai/wai-web/src/neural_int.rs | 54 +++++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/wai/wai-web/src/neural_int.rs b/wai/wai-web/src/neural_int.rs index a3ce40e18..3c703f949 100644 --- a/wai/wai-web/src/neural_int.rs +++ b/wai/wai-web/src/neural_int.rs @@ -172,3 +172,57 @@ pub fn decode_jpegai_mans( dec.decode_sgm(sigma, &mut out, &masks_b); out } + +// ---- JPEG-AI full decode → RGB, in the browser (PTQ value/synthesis path) ---- + +/// Decode a full ISO/IEC 6048 JPEG-AI codestream to RGB in the browser +/// (`wai.image.jpegai`, decode-equivalence). The same decoder a native sink runs +/// (`wai::jpegai_decode::decode`), compiled to wasm32: integer entropy front end + +/// deterministic value/synthesis/post-filter path. One pinned reconstruction, +/// byte-identical on every machine — so the decoded RGB has a portable hash. The 13 +/// byte inputs are the model weight bundle (`tools/wai_jpegai_build_e2e_inputs.py`): +/// per-plane hyper-scale / hyper-decoder / synthesis, the four luma MCM stages, the +/// eICCI net, and the normative tables (`norm`). +#[wasm_bindgen(js_name = decodeJpegAiFull)] +#[allow(clippy::too_many_arguments)] +pub fn decode_jpegai_full( + codestream: &[u8], + hs_y: &[u8], hyp_y: &[u8], synth_y: &[u8], + hs_uv: &[u8], hyp_uv: &[u8], synth_uv: &[u8], + mcm0: &[u8], mcm1: &[u8], mcm2: &[u8], mcm3: &[u8], + icci: &[u8], norm: &[u8], +) -> Result { + use wai::jpegai_decode::{decode, DecoderWeights, NormTables, PlaneWeights}; + use wai::jpegai_icci::IcciNet; + use wai::jpegai_synthesis::{IntNet, Weights}; + + let parse = |b: &[u8]| -> Result { + serde_json::from_slice(b).map_err(|e| JsError::new(&e.to_string())) + }; + let w = |b: &[u8]| -> Result { Ok(Weights::from_json(&parse(b)?)) }; + let net = |b: &[u8]| -> Result { Ok(IntNet::from_json(&parse(b)?)) }; + + let weights = DecoderWeights { + luma: PlaneWeights { hyper_scale: net(hs_y)?, hyper_decoder: w(hyp_y)?, synthesis: w(synth_y)? }, + chroma: PlaneWeights { hyper_scale: net(hs_uv)?, hyper_decoder: w(hyp_uv)?, synthesis: w(synth_uv)? }, + mcm_y: [w(mcm0)?, w(mcm1)?, w(mcm2)?, w(mcm3)?], + icci: IcciNet::from_json(&parse(icci)?), + }; + + let nv = parse(norm)?; + let i64v = |k: &str| nv[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect::>(); + let u8v = |k: &str| nv[k].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect::>(); + let i32v = |k: &str| nv[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i32).collect::>(); + let i3 = |k: &str| { let a = nv[k].as_array().unwrap(); [a[0].as_i64().unwrap(), a[1].as_i64().unwrap(), a[2].as_i64().unwrap()] }; + let f3 = |k: &str| { let a = nv[k].as_array().unwrap(); [a[0].as_f64().unwrap() as f32, a[1].as_f64().unwrap() as f32, a[2].as_f64().unwrap() as f32] }; + let norm_t = NormTables { + pmf: nv["pmf"].as_array().ok_or_else(|| JsError::new("norm.pmf"))?.iter() + .map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()).collect(), + bounds: u8v("bounds"), per_channel: nv["per_channel"].as_u64().unwrap() as usize, + cdfs_y: u8v("cdfs_y"), cdfs_uv: u8v("cdfs_uv"), gain_y: i64v("gain_y"), gain_uv: i64v("gain_uv"), + lsbs_t0: i32v("lsbs_t0"), lsbs_t1: i32v("lsbs_t1"), lef_thr: i3("lef_thr0"), lef_mag: f3("lef_mag0"), + }; + + let img = decode(codestream, &weights, &norm_t).ok_or_else(|| JsError::new("unsupported geometry"))?; + Ok(DecodedImage { rgb: img.data, width: img.w as u32, height: img.h as u32 }) +} -- 2.54.0 (Apple Git-157)