open-standards/wai/wai-rs/src/int_synth.rs

//! Integer fixed-point synthesis — a learned-image decoder whose reconstruction
//! is **bit-identical on every machine, by construction**.
//!
//! This is the deterministic member of the neural family. The integer rANS
//! recovers latents exactly; the divergence in float decoders like
//! `bmshj2018` / `mbt2018_mean` is the **float ONNX synthesis** — and float
//! round-off is platform-dependent. Here the synthesis is a stack of integer
//! transposed convolutions with integer activations: every operation is `i64`
//! arithmetic + an arithmetic shift, so two sinks honoring the spec produce the
//! same bytes everywhere. No ONNX, no float, nothing implementation-defined.
//!
//! Because it pulls in no float runtime, this decoder lives at the crate root
//! behind the light `neural_int` feature (just `serde_json`), not inside the
//! `ort`-backed `neural` module. That keeps it buildable — and its
//! cross-hardware bit-exactness checkable — in a CI matrix with no ONNX
//! Runtime present. The `neural` module re-exposes it through the
//! `wai.neural.int_synth` capability.
//!
//! Quality needs a quantization-aware-trained model (the weights are integers);
//! producing that model is a tooling step (`tools/wai_int_synth_export.py`). The
//! *decode contract* — the thing that makes cross-hardware bit-exactness real —
//! lives here and is exact regardless.
//!
//! Wire format (little-endian):
//!   u32 H, u32 W, u16 C, u8 S_log2, u8 _reserved   (12 B)
//!   C*S*S bytes   int8 latents, channel-major
//!
//! Sink file (next to the registered model path): `int_synth.json`
//!   { "layers": [ { in_ch, out_ch, k, stride, pad, out_pad,
//!                   weight:[i32; in*out*k*k] (layout ic,oc,ky,kx),
//!                   bias:[i32; out_ch], shift:u32, relu:bool }, … ] }
//! The final layer MUST output 3 channels at H×W; values are clamped to [0,255].

use std::path::Path;

use serde::Deserialize;

const HEADER: usize = 12;

/// Decode error for the integer synthesis path. Deliberately self-contained so
/// the decoder carries no dependency on the `ort`-backed `neural` module; the
/// `neural` dispatch maps this onto its own `DecodeError::InvalidPayload`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum IntSynthError {
    InvalidPayload(String),
}

impl std::fmt::Display for IntSynthError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::InvalidPayload(e) => write!(f, "invalid payload: {e}"),
        }
    }
}
impl std::error::Error for IntSynthError {}

/// Decoded image: RGB row-major u8, H × W × 3 — the integer-exact output.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IntImage {
    pub rgb: Vec<u8>,
    pub width: u32,
    pub height: u32,
}

#[derive(Debug, Clone, Deserialize)]
struct LayerSpec {
    in_ch: usize,
    out_ch: usize,
    k: usize,
    stride: usize,
    pad: usize,
    #[serde(default)]
    out_pad: usize,
    /// Integer weights, layout `[in_ch][out_ch][k][k]` (transposed-conv order).
    weight: Vec<i32>,
    /// Integer bias, `[out_ch]`.
    bias: Vec<i32>,
    /// Arithmetic right-shift applied after accumulate+bias (requantize).
    #[serde(default)]
    shift: u32,
    /// Integer ReLU (clamp at 0) after requantize.
    #[serde(default)]
    relu: bool,
}

#[derive(Debug, Clone, Deserialize)]
struct SynthSpec {
    layers: Vec<LayerSpec>,
}

/// One integer transposed convolution + bias + requantize + optional ReLU.
/// `input` is `[in_ch, hin, win]` row-major `i64`. Returns `([out_ch, hout,
/// wout] i64, hout, wout)`. Pure integer arithmetic — the whole point.
fn conv_transpose2d(
    input: &[i64],
    cin: usize,
    hin: usize,
    win: usize,
    l: &LayerSpec,
) -> Result<(Vec<i64>, usize, usize), IntSynthError> {
    let (k, s, p, op) = (l.k, l.stride, l.pad, l.out_pad);
    let hout_i = (hin as i64 - 1) * s as i64 + k as i64 + op as i64 - 2 * p as i64;
    let wout_i = (win as i64 - 1) * s as i64 + k as i64 + op as i64 - 2 * p as i64;
    if hout_i <= 0 || wout_i <= 0 {
        return Err(IntSynthError::InvalidPayload(format!(
            "int_synth layer produces non-positive output {hout_i}x{wout_i}")));
    }
    let (hout, wout) = (hout_i as usize, wout_i as usize);
    if l.weight.len() != cin * l.out_ch * k * k {
        return Err(IntSynthError::InvalidPayload(format!(
            "int_synth weight len {} != {}*{}*{}*{}", l.weight.len(), cin, l.out_ch, k, k)));
    }
    if l.bias.len() != l.out_ch {
        return Err(IntSynthError::InvalidPayload(format!(
            "int_synth bias len {} != out_ch {}", l.bias.len(), l.out_ch)));
    }

    let mut out = vec![0i64; l.out_ch * hout * wout];
    for ic in 0..cin {
        for iy in 0..hin {
            for ix in 0..win {
                let v = input[(ic * hin + iy) * win + ix];
                if v == 0 {
                    continue;
                }
                for oc in 0..l.out_ch {
                    let wbase = (ic * l.out_ch + oc) * k * k;
                    let obase = oc * hout * wout;
                    for ky in 0..k {
                        let ty = iy * s + ky;
                        if ty < p {
                            continue;
                        }
                        let oy = ty - p;
                        if oy >= hout {
                            continue;
                        }
                        for kx in 0..k {
                            let tx = ix * s + kx;
                            if tx < p {
                                continue;
                            }
                            let ox = tx - p;
                            if ox >= wout {
                                continue;
                            }
                            out[obase + oy * wout + ox] +=
                                l.weight[wbase + ky * k + kx] as i64 * v;
                        }
                    }
                }
            }
        }
    }

    // Bias, requantize (arithmetic shift — deterministic), integer ReLU.
    for oc in 0..l.out_ch {
        let b = l.bias[oc] as i64;
        let base = oc * hout * wout;
        for v in &mut out[base..base + hout * wout] {
            let mut x = (*v + b) >> l.shift;
            if l.relu && x < 0 {
                x = 0;
            }
            *v = x;
        }
    }
    Ok((out, hout, wout))
}

/// Decode an int_synth payload against the `int_synth.json` sibling of
/// `model_path`. Byte-identical on every machine — the `replicate` contract.
///
/// Filesystem-only convenience: reads the spec and delegates to
/// [`decode_from_spec`], which is the entry point with no `std::fs` (the one a
/// WASM / no-filesystem sink calls with the spec bytes supplied by the host).
pub fn decode(payload: &[u8], model_path: &Path) -> Result<IntImage, IntSynthError> {
    let dir = model_path.parent()
        .ok_or_else(|| IntSynthError::InvalidPayload("model path has no parent".into()))?;
    let spec_path = dir.join("int_synth.json");
    let raw = std::fs::read(&spec_path)
        .map_err(|e| IntSynthError::InvalidPayload(format!("{}: {e}", spec_path.display())))?;
    decode_from_spec(payload, &raw)
}

/// Decode an int_synth payload against the model spec supplied as bytes (the
/// `int_synth.json` contents). No filesystem, no float, no ONNX — the
/// integer-exact decode a browser / WASM sink runs directly, host-supplied
/// weights in, RGB out. Byte-identical on every machine.
pub fn decode_from_spec(payload: &[u8], spec_json: &[u8]) -> Result<IntImage, IntSynthError> {
    if payload.len() < HEADER {
        return Err(IntSynthError::InvalidPayload(format!(
            "int_synth payload {} B < {HEADER} B header", payload.len())));
    }
    let h = u32::from_le_bytes(payload[0..4].try_into().unwrap()) as usize;
    let w = u32::from_le_bytes(payload[4..8].try_into().unwrap()) as usize;
    let c = u16::from_le_bytes(payload[8..10].try_into().unwrap()) as usize;
    let s = 1usize << payload[10];
    let n = c * s * s;
    if payload.len() != HEADER + n {
        return Err(IntSynthError::InvalidPayload(format!(
            "int_synth payload {} != {HEADER} + C*S*S={n}", payload.len())));
    }

    let spec: SynthSpec = serde_json::from_slice(spec_json)
        .map_err(|e| IntSynthError::InvalidPayload(format!("int_synth.json parse: {e}")))?;

    // int8 latents → i64.
    let mut cur: Vec<i64> = payload[HEADER..].iter().map(|&b| b as i8 as i64).collect();
    let (mut cin, mut ch, mut cw) = (c, s, s);
    for (i, l) in spec.layers.iter().enumerate() {
        if l.in_ch != cin {
            return Err(IntSynthError::InvalidPayload(format!(
                "int_synth layer {i} in_ch {} != current {cin}", l.in_ch)));
        }
        let (out, ho, wo) = conv_transpose2d(&cur, cin, ch, cw, l)?;
        cur = out;
        cin = l.out_ch;
        ch = ho;
        cw = wo;
    }
    if cin != 3 || ch != h || cw != w {
        return Err(IntSynthError::InvalidPayload(format!(
            "int_synth output [{cin},{ch},{cw}] != [3,{h},{w}]")));
    }

    let mut rgb = vec![0u8; h * w * 3];
    let hw = h * w;
    for y in 0..h {
        for x in 0..w {
            let pi = y * w + x;
            for ch_i in 0..3 {
                rgb[pi * 3 + ch_i] = cur[ch_i * hw + pi].clamp(0, 255) as u8;
            }
        }
    }
    Ok(IntImage { rgb, width: w as u32, height: h as u32 })
}

#[cfg(test)]
mod tests {
    use super::*;

    fn layer(out_ch: usize, in_ch: usize, k: usize, stride: usize, weight: Vec<i32>) -> LayerSpec {
        LayerSpec {
            in_ch, out_ch, k, stride, pad: 0, out_pad: 0,
            weight, bias: vec![0; out_ch], shift: 0, relu: false,
        }
    }

    #[test]
    fn transposed_conv_matches_hand_computation() {
        // 1 in-channel, 1 out-channel, 2×2 kernel [[1,2],[3,4]], stride 2.
        // input [[5]] (1×1) scatters to a 2×2 output = kernel * 5.
        let l = layer(1, 1, 2, 2, vec![1, 2, 3, 4]);
        let (out, ho, wo) = conv_transpose2d(&[5], 1, 1, 1, &l).unwrap();
        assert_eq!((ho, wo), (2, 2));
        assert_eq!(out, vec![5, 10, 15, 20]);
    }

    #[test]
    fn requantize_shift_and_relu_are_integer_exact() {
        // shift=1 halves (arithmetic), relu clamps negatives.
        let mut l = layer(1, 1, 1, 1, vec![1]);
        l.shift = 1;
        l.relu = true;
        l.bias = vec![0];
        // input -3 → weight*in = -3 → >>1 = -2 → relu → 0.
        let (out, _, _) = conv_transpose2d(&[-3], 1, 1, 1, &l).unwrap();
        assert_eq!(out, vec![0]);
        // input 7 → 7 → >>1 = 3.
        let (out, _, _) = conv_transpose2d(&[7], 1, 1, 1, &l).unwrap();
        assert_eq!(out, vec![3]);
    }

    #[test]
    fn decode_is_deterministic_by_construction() {
        // Two runs of the pure integer op on the same input are byte-identical
        // (the property that makes this a sound `replicate` member).
        let l = layer(2, 1, 3, 2, (0..18).collect());
        let input = [3, -4, 5];
        let a = conv_transpose2d(&input, 1, 1, 3, &l).unwrap();
        let b = conv_transpose2d(&input, 1, 1, 3, &l).unwrap();
        assert_eq!(a, b);
    }

    /// FULL decode path, no ONNX: a constructed 1-layer integer model + a tiny
    /// latent → exact RGB, and byte-identical across runs. This validates the
    /// integer-exact `replicate` member end to end (header parse, latent unpack,
    /// conv, requantize, clamp, RGB layout).
    #[test]
    fn end_to_end_decode_is_exact_and_deterministic() {
        let dir = std::env::temp_dir().join("wai_int_synth_e2e");
        std::fs::create_dir_all(&dir).unwrap();
        // 1 layer, in_ch=1 → out_ch=3 (RGB), 2×2 kernel, stride 2. Each output
        // channel has a constant kernel (1, 2, 3), so a 1×1 latent of 10 scatters
        // to a 2×2 image where every pixel is (10, 20, 30).
        let weight: Vec<i32> = (1..=3).flat_map(|oc| std::iter::repeat(oc).take(4)).collect();
        let spec = serde_json::json!({"layers": [{
            "in_ch": 1, "out_ch": 3, "k": 2, "stride": 2, "pad": 0, "out_pad": 0,
            "weight": weight, "bias": [0, 0, 0], "shift": 0, "relu": false,
        }]});
        std::fs::write(dir.join("int_synth.json"), serde_json::to_vec(&spec).unwrap()).unwrap();
        let model_path = dir.join("model");

        let mut payload = Vec::new();
        payload.extend_from_slice(&2u32.to_le_bytes()); // H
        payload.extend_from_slice(&2u32.to_le_bytes()); // W
        payload.extend_from_slice(&1u16.to_le_bytes()); // C
        payload.push(0); // S_log2 = 0  → S = 1
        payload.push(0); // reserved
        payload.push(10u8); // one int8 latent = 10

        let img = decode(&payload, &model_path).expect("decode");
        assert_eq!((img.width, img.height), (2, 2));
        let expected: Vec<u8> = (0..4).flat_map(|_| [10u8, 20, 30]).collect();
        assert_eq!(img.rgb, expected, "integer synthesis must hit the exact bytes");
        // Byte-identical across runs — the replicate property.
        assert_eq!(decode(&payload, &model_path).unwrap().rgb, img.rgb);
        let _ = std::fs::remove_dir_all(&dir);
    }
}