From 3b0fbe3abe2f6e33ebb07eca19a61d35528492dc Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Wed, 15 Jul 2026 03:38:56 -0400 Subject: [PATCH] =?UTF-8?q?wai-rs:=20jpegai=5Ficci=20=E2=80=94=20the=20eIC?= =?UTF-8?q?CI=20network=20body=20(conv,=20BN-fold,=20residual=20block,=20p?= =?UTF-8?q?lane=20path)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second piece of eICCI. The reference network is float, so this runs in f32 rather than quantising: IEEE-754 add/mul with a fixed accumulation order is deterministic across machines, so a faithful float port is byte-identical AND carries no PTQ error — no quantise/calibrate loop is needed for this filter. Conv stride-1 zero-padded conv, [cout,cin,kh,kw]; supports the 3x3 stem and the separable [1,3] / [3,1] block kernels Conv::fold_bn folds BatchNorm into the preceding conv (exact at inference): BN(conv(x)) = (w·γ/√(σ²+ε))·x + ((b−μ)·γ/√(σ²+ε) + β) ResBlock x + scale · BN2(conv2_3x1(relu(BN1(conv1_1x3(x))))) [BNs folded] IcciPath relu(BN(conv_first(yuv48))) -> N residual blocks -> conv_last -> out·scale + subbands16 -> idwt2 -> plane Unit tests: identity conv + zero-padding at the edge, fold_bn matching an explicit BatchNorm, and the residual block (scale=0 is the identity; scale=1 with identity convs adds relu(x) back). Remaining for eICCI: load the 20 models' float weights from the ONNX, and the per-tile model selection. Co-Authored-By: Claude Opus 4.8 --- wai/wai-rs/src/jpegai_icci.rs | 188 ++++++++++++++++++++++++++++++++++ 1 file changed, 188 insertions(+) diff --git a/wai/wai-rs/src/jpegai_icci.rs b/wai/wai-rs/src/jpegai_icci.rs index f00fe2d53..08972f7e5 100644 --- a/wai/wai-rs/src/jpegai_icci.rs +++ b/wai/wai-rs/src/jpegai_icci.rs @@ -112,6 +112,136 @@ pub fn icci_dwt_planes(y: &[f32], u: &[f32], v: &[f32], h: usize, w: usize) (y16, u16, v16, yuv48, h2, w2) } +// ---- network body ---- +// +// The eICCI CNN is float in the reference, so we run it in f32 rather than quantising: +// IEEE-754 add/mul with a fixed accumulation order is deterministic across machines, so +// a faithful float port is byte-identical *and* avoids any PTQ error. BatchNorm is +// folded into the preceding convolution, which is exact at inference: +// +// BN(conv(x)) = γ·(w·x + b − μ)/√(σ²+ε) + β +// = (w · γ/√(σ²+ε))·x + ((b − μ)·γ/√(σ²+ε) + β) + +/// A convolution with its BatchNorm already folded in. Weights are `[cout, cin, kh, kw]` +/// row-major; `bias` is per output channel. Padding is `(ph, pw)`, stride 1. +#[derive(Clone)] +pub struct Conv { + pub w: Vec, + pub b: Vec, + pub cin: usize, + pub cout: usize, + pub kh: usize, + pub kw: usize, + pub ph: usize, + pub pw: usize, +} + +impl Conv { + /// Fold a BatchNorm into a convolution's weights/bias (exact at inference). + /// `gamma`/`beta`/`mean`/`var` are per output channel. + pub fn fold_bn(mut self, gamma: &[f32], beta: &[f32], mean: &[f32], var: &[f32], eps: f32) -> Self { + assert_eq!(gamma.len(), self.cout, "BN params must be per output channel"); + let tap = self.cin * self.kh * self.kw; + for o in 0..self.cout { + let s = gamma[o] / (var[o] + eps).sqrt(); + for t in 0..tap { + self.w[o * tap + t] *= s; + } + self.b[o] = (self.b[o] - mean[o]) * s + beta[o]; + } + self + } + + /// Stride-1 zero-padded convolution: `[cin, h, w]` → `[cout, h, w]`. + pub fn forward(&self, x: &[f32], h: usize, w: usize) -> Vec { + assert_eq!(x.len(), self.cin * h * w, "conv input must be [cin, h, w]"); + let mut out = vec![0f32; self.cout * h * w]; + for o in 0..self.cout { + for oy in 0..h { + for ox in 0..w { + let mut acc = self.b[o]; + for i in 0..self.cin { + for ky in 0..self.kh { + let iy = oy + ky; + if iy < self.ph || iy - self.ph >= h { continue; } + let sy = iy - self.ph; + for kx in 0..self.kw { + let ix = ox + kx; + if ix < self.pw || ix - self.pw >= w { continue; } + let sx = ix - self.pw; + acc += self.w[((o * self.cin + i) * self.kh + ky) * self.kw + kx] + * x[(i * h + sy) * w + sx]; + } + } + } + out[(o * h + oy) * w + ox] = acc; + } + } + } + out + } +} + +/// One `ResidualBlock_BN_RectKernel`: separable rectangular kernels with a scaled residual. +/// +/// ```text +/// out = x + scale · BN2( conv2_{3×1}( relu( BN1( conv1_{1×3}(x) ) ) ) ) +/// ``` +/// +/// Both BatchNorms are folded into their convolutions, so this is two convs + a ReLU. +#[derive(Clone)] +pub struct ResBlock { + pub conv1: Conv, // [1,3] kernel, pad [0,1] + pub conv2: Conv, // [3,1] kernel, pad [1,0] + pub scale: f32, +} + +impl ResBlock { + pub fn forward(&self, x: &[f32], h: usize, w: usize) -> Vec { + let mut t = self.conv1.forward(x, h, w); + for v in t.iter_mut() { if *v < 0.0 { *v = 0.0; } } + let t = self.conv2.forward(&t, h, w); + x.iter().zip(&t).map(|(&a, &b)| a + b * self.scale).collect() + } +} + +/// One plane-path of the eICCI network: `conv_first` + BN + ReLU → N residual blocks → +/// `conv_last`, then a scaled residual add onto the plane's own subbands. +/// +/// ```text +/// fea = relu( BN( conv_first(yuv48) ) ) // BN folded into conv_first +/// fea = resblocks(fea) +/// out16 = conv_last(fea) · scale + subbands16 +/// plane = idwt2(out16) +/// ``` +#[derive(Clone)] +pub struct IcciPath { + pub conv_first: Conv, // 48 → nf, 3×3, pad 1 (BN folded) + pub blocks: Vec, + pub conv_last: Conv, // nf → 16, 3×3, pad 1 + pub scale: f32, +} + +impl IcciPath { + /// Run the path and return the refined 16 subbands (before the inverse transform). + pub fn refine_subbands(&self, yuv48: &[f32], subbands16: &[f32], h: usize, w: usize) -> Vec { + let mut fea = self.conv_first.forward(yuv48, h, w); + for v in fea.iter_mut() { if *v < 0.0 { *v = 0.0; } } + for blk in &self.blocks { + fea = blk.forward(&fea, h, w); + } + let res = self.conv_last.forward(&fea, h, w); + assert_eq!(res.len(), subbands16.len(), "conv_last must emit the 16 subbands"); + res.iter().zip(subbands16).map(|(&r, &s)| r * self.scale + s).collect() + } + + /// Full plane: refine the subbands then invert the two-level transform. + pub fn forward_plane(&self, yuv48: &[f32], subbands16: &[f32], h: usize, w: usize) -> Vec { + let refined = self.refine_subbands(yuv48, subbands16, h, w); + haar_idwt2(&refined, h, w).0 + } +} + #[cfg(test)] mod tests { use super::*; @@ -168,4 +298,62 @@ mod tests { // the Y block is the leading 16 channels assert_eq!(&yuv48[..y16.len()], &y16[..]); } + + fn ident_conv(c: usize, kh: usize, kw: usize, ph: usize, pw: usize) -> Conv { + // identity: centre tap 1 on the matching in/out channel, everything else 0 + let mut w = vec![0f32; c * c * kh * kw]; + for o in 0..c { + w[((o * c + o) * kh + kh / 2) * kw + kw / 2] = 1.0; + } + Conv { w, b: vec![0f32; c], cin: c, cout: c, kh, kw, ph, pw } + } + + #[test] + fn conv_identity_and_zero_padding() { + let x = vec![1.0f32, 2.0, 3.0, 4.0]; + // 1×3 identity with pad [0,1] reproduces the input + let o = ident_conv(1, 1, 3, 0, 1).forward(&x, 2, 2); + assert_eq!(o, x, "identity conv must be a no-op"); + + // a pure left-shift tap reads the zero pad at the left edge + let mut c = Conv { w: vec![0.0, 0.0, 1.0], b: vec![0.0], cin: 1, cout: 1, kh: 1, kw: 3, ph: 0, pw: 1 }; + c.b[0] = 0.0; + let o = c.forward(&[5.0, 7.0], 1, 2); + // tap at kx=2 reads x[ox+2-1] = x[ox+1]: out[0]=x[1]=7, out[1]=pad=0 + assert_eq!(o, vec![7.0, 0.0], "zero padding at the right edge"); + } + + #[test] + fn fold_bn_matches_explicit_batchnorm() { + // conv = identity, BN with known params → folded conv must equal BN(x) directly + let (gamma, beta, mean, var, eps) = ([2.0f32], [0.5f32], [1.0f32], [4.0f32], 1e-5f32); + let folded = ident_conv(1, 1, 1, 0, 0).fold_bn(&gamma, &beta, &mean, &var, eps); + let x = [3.0f32, -1.0]; + let got = folded.forward(&x, 1, 2); + for i in 0..x.len() { + let want = gamma[0] * (x[i] - mean[0]) / (var[0] + eps).sqrt() + beta[0]; + assert!((got[i] - want).abs() < 1e-5, "fold != BN: {} vs {want}", got[i]); + } + } + + #[test] + fn resblock_is_identity_when_scale_is_zero() { + let c = 2usize; + let (h, w) = (3usize, 3usize); + let x = ramp(c * h * w); + let blk = ResBlock { + conv1: ident_conv(c, 1, 3, 0, 1), + conv2: ident_conv(c, 3, 1, 1, 0), + scale: 0.0, + }; + assert_eq!(blk.forward(&x, h, w), x, "scale=0 must leave the identity path alone"); + + // scale=1 with identity convs adds relu(x) back onto x + let blk1 = ResBlock { scale: 1.0, ..blk }; + let got = blk1.forward(&x, h, w); + for i in 0..x.len() { + let want = x[i] + if x[i] > 0.0 { x[i] } else { 0.0 }; + assert!((got[i] - want).abs() < 1e-5, "at {i}: {} vs {want}", got[i]); + } + } } -- 2.54.0 (Apple Git-157)