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

//! The perceptual / generative-semantic path — the Rust port of the moved-out
//! `wai-neural` prototype, on-charter this time.
//!
//! Where the bit-exact path ([`crate::neural`]) pins *which model* by hash and
//! reproduces bytes, this path does not reproduce bytes, and says so: it
//! transmits a compact **understanding** of the source and lets the sink
//! **regenerate** something semantically equivalent with whatever conforming
//! generator it has. Conformance is **behavioral** — CLIP cosine similarity
//! against the transmitted understanding ≥ the file's declared threshold
//! ([`ModelRequirement::min_sim`](crate::container::ModelRequirement)) — never
//! byte identity, never a weights hash. Any conforming generator may serve the
//! file; the carrier names *how close is close enough*, not *which model*.
//!
//! The heavy models live at the sink, never in this crate: the generator and
//! the CLIP embedder are [`SemanticGenerator`] / [`SemanticEmbedder`] traits the
//! sink supplies (mirroring how [`crate::neural`] takes decoder ONNX from a
//! registry rather than bundling weights). This module owns the wire format,
//! the conformance metric, and the verify loop — the parts that must be exact.

use std::path::Path;

use crate::container::Wai;

/// RGB image, row-major, 3 bytes/pixel — what a [`SemanticGenerator`] produces
/// and a [`SemanticEmbedder`] consumes.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Frame {
    pub width: u32,
    pub height: u32,
    pub rgb: Vec<u8>,
}

#[derive(Debug)]
pub enum SemanticError {
    /// The semantic-package payload is malformed / truncated.
    Truncated(String),
    /// A frame's `rgb` length does not match `width * height * 3`.
    InvalidImage(String),
    /// The sink generator or embedder failed.
    Generator(String),
    /// ONNX runtime error from the bundled [`OnnxClip`] embedder.
    Ort(String),
}

impl std::fmt::Display for SemanticError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Truncated(e) => write!(f, "truncated semantic package: {e}"),
            Self::InvalidImage(e) => write!(f, "invalid image: {e}"),
            Self::Generator(e) => write!(f, "generator: {e}"),
            Self::Ort(e) => write!(f, "onnxruntime error: {e}"),
        }
    }
}
impl std::error::Error for SemanticError {}

/// The wire payload of a `wai.semantic.*` envelope: a compact, model-agnostic
/// "understanding" of the source. Byte-faithful to the prototype's
/// `semantic_package_v1` (four length-prefixed parts).
#[derive(Debug, Clone, PartialEq)]
pub struct SemanticPackage {
    /// Dequantization scale for the CLIP vector.
    pub scale: f32,
    /// CLIP image embedding, int8-quantized (projection dim, typically 768).
    pub clip_q: Vec<i8>,
    /// Tiny JPEG thumbnail — global colour / composition.
    pub thumb_jpeg: Vec<u8>,
    /// 1-bit Canny edge map, PNG — geometry / layout.
    pub edges_png: Vec<u8>,
    /// Optional short text tag.
    pub tag: String,
}

impl SemanticPackage {
    /// Build a package from a (unit-norm) CLIP embedding plus the side maps.
    /// Quantizes the embedding to int8 exactly as the prototype does.
    pub fn from_embedding(
        embedding: &[f32],
        thumb_jpeg: Vec<u8>,
        edges_png: Vec<u8>,
        tag: impl Into<String>,
    ) -> Self {
        let peak = embedding.iter().fold(0f32, |m, &x| m.max(x.abs()));
        let scale = if peak == 0.0 { 1.0 } else { peak };
        let clip_q = embedding
            .iter()
            .map(|&x| (x / scale * 127.0).round().clamp(-127.0, 127.0) as i8)
            .collect();
        Self { scale, clip_q, thumb_jpeg, edges_png, tag: tag.into() }
    }

    /// Dequantized embedding vector (`q * scale / 127`). Not re-normalized
    /// here; [`cosine_sim`] normalizes both operands.
    pub fn clip(&self) -> Vec<f32> {
        self.clip_q.iter().map(|&q| q as f32 * self.scale / 127.0).collect()
    }

    /// Serialize to the `semantic_package_v1` wire bytes.
    pub fn pack(&self) -> Vec<u8> {
        let mut clip_part = Vec::with_capacity(4 + self.clip_q.len());
        clip_part.extend_from_slice(&self.scale.to_le_bytes());
        clip_part.extend(self.clip_q.iter().map(|&b| b as u8));

        let parts: [&[u8]; 4] =
            [&clip_part, &self.thumb_jpeg, &self.edges_png, self.tag.as_bytes()];
        let mut out = Vec::new();
        for p in parts {
            out.extend_from_slice(&(p.len() as u32).to_le_bytes());
            out.extend_from_slice(p);
        }
        out
    }

    /// Parse `semantic_package_v1` bytes. Bounds-checked: malformed input
    /// returns `Err`, never panics.
    pub fn unpack(blob: &[u8]) -> Result<Self, SemanticError> {
        fn take<'a>(blob: &'a [u8], o: &mut usize) -> Result<&'a [u8], SemanticError> {
            if *o + 4 > blob.len() {
                return Err(SemanticError::Truncated(format!("length prefix at {o}")));
            }
            let n = u32::from_le_bytes(blob[*o..*o + 4].try_into().unwrap()) as usize;
            *o += 4;
            if *o + n > blob.len() {
                return Err(SemanticError::Truncated(format!("part of {n} bytes at {o}")));
            }
            let s = &blob[*o..*o + n];
            *o += n;
            Ok(s)
        }

        let mut o = 0usize;
        let clip_part = take(blob, &mut o)?;
        let thumb = take(blob, &mut o)?;
        let edges = take(blob, &mut o)?;
        let tag = take(blob, &mut o)?;

        if clip_part.len() < 4 {
            return Err(SemanticError::Truncated("clip part < 4 bytes".into()));
        }
        let scale = f32::from_le_bytes(clip_part[0..4].try_into().unwrap());
        let clip_q = clip_part[4..].iter().map(|&b| b as i8).collect();
        Ok(Self {
            scale,
            clip_q,
            thumb_jpeg: thumb.to_vec(),
            edges_png: edges.to_vec(),
            tag: String::from_utf8_lossy(tag).into_owned(),
        })
    }
}

/// CLIP cosine similarity — the conformance metric. `1.0` identical, `0.0`
/// orthogonal. Both operands are normalized here, so callers may pass raw
/// (unnormalized) embeddings.
pub fn cosine_sim(a: &[f32], b: &[f32]) -> f32 {
    let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
    let na = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let nb = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    if na == 0.0 || nb == 0.0 { 0.0 } else { dot / (na * nb) }
}

/// A sink-supplied CLIP image embedder. The protocol needs embeddings to score
/// conformance; the model that produces them is the sink's, not this crate's.
pub trait SemanticEmbedder {
    /// Unit-norm CLIP image embedding of `frame`.
    fn embed(&self, frame: &Frame) -> Result<Vec<f32>, SemanticError>;
}

/// A sink-supplied generative model that regenerates media from a semantic
/// package. Any conforming generator may serve — SD-XL-Turbo, a tiny on-device
/// model, whatever is real-time on *this* sink's hardware.
pub trait SemanticGenerator {
    fn regenerate(&self, pkg: &SemanticPackage, w: u32, h: u32) -> Result<Frame, SemanticError>;
}

/// Outcome of a verified regeneration.
#[derive(Debug, Clone)]
pub struct SemanticResult {
    pub frame: Frame,
    /// Best similarity (cosine in the manifest's `sim_metric` space) achieved
    /// against the transmitted understanding.
    pub similarity: f32,
    /// Whether `similarity` met the file's declared `min_sim` threshold.
    pub conformed: bool,
    /// The embedding space the score was computed in (the envelope's
    /// `sim_metric`, e.g. `"clip"`, `"dreamsim"`). The sink MUST have scored
    /// with an embedder matching this space for the result to mean anything.
    pub metric: String,
}

/// Regenerate semantically-equivalent media from a `wai.semantic.*` envelope,
/// best-of-`k`, and gate on **behavioral conformance**: the regenerated
/// frame's CLIP embedding must reach the file's declared
/// [`min_sim`](crate::container::ModelRequirement::min_sim) against the
/// transmitted understanding. Byte identity is never required; semantic
/// equivalence is. Returns the best candidate with its score and whether it
/// conformed — the caller decides what to do with a non-conforming result.
pub fn regenerate_verified<G: SemanticGenerator, E: SemanticEmbedder>(
    env: &Wai,
    generator: &G,
    embedder: &E,
    w: u32,
    h: u32,
    k: usize,
) -> Result<SemanticResult, SemanticError> {
    let pkg = SemanticPackage::unpack(&env.payload)?;
    let target = pkg.clip();
    let threshold = env.manifest.model_requirement.min_sim.unwrap_or(0.0);

    let mut best: Option<(Frame, f32)> = None;
    for _ in 0..k.max(1) {
        let frame = generator.regenerate(&pkg, w, h)?;
        let sim = cosine_sim(&embedder.embed(&frame)?, &target);
        let better = match &best {
            None => true,
            Some((_, b)) => sim > *b,
        };
        if better {
            best = Some((frame, sim));
        }
    }
    let (frame, sim) = best.expect("k.max(1) >= 1 guarantees a candidate");
    Ok(SemanticResult {
        frame,
        similarity: sim,
        conformed: sim >= threshold,
        metric: env.sim_metric().to_string(),
    })
}

// ---- ORT-backed CLIP embedder (sink reference implementation) -------------

/// CLIP image embedder backed by a sink-installed ONNX model, run through the
/// shared deterministic runtime. This is a *reference* `SemanticEmbedder` — the
/// model file is the sink's; this crate bundles no weights.
pub struct OnnxClip {
    session: std::sync::Mutex<ort::session::Session>,
}

impl OnnxClip {
    /// Load a CLIP image-encoder ONNX (output = image embedding).
    pub fn load(model: &Path) -> Result<Self, SemanticError> {
        let s = crate::ort_session::load_deterministic(model).map_err(SemanticError::Ort)?;
        Ok(Self { session: std::sync::Mutex::new(s) })
    }
}

impl SemanticEmbedder for OnnxClip {
    fn embed(&self, frame: &Frame) -> Result<Vec<f32>, SemanticError> {
        const SIZE: u32 = 224;
        // CLIP ViT-L/14 preprocessing.
        const MEAN: [f32; 3] = [0.481_454_66, 0.457_827_5, 0.408_210_73];
        const STD: [f32; 3] = [0.268_629_54, 0.261_302_58, 0.275_777_11];

        if frame.rgb.len() != (frame.width * frame.height * 3) as usize {
            return Err(SemanticError::InvalidImage("rgb len != w*h*3".into()));
        }
        let img = image::RgbImage::from_raw(frame.width, frame.height, frame.rgb.clone())
            .ok_or_else(|| SemanticError::InvalidImage("from_raw failed".into()))?;
        let resized =
            image::imageops::resize(&img, SIZE, SIZE, image::imageops::FilterType::Triangle);

        let mut input = ndarray::Array4::<f32>::zeros((1, 3, SIZE as usize, SIZE as usize));
        for y in 0..SIZE {
            for x in 0..SIZE {
                let px = resized.get_pixel(x, y);
                for c in 0..3 {
                    input[[0, c, y as usize, x as usize]] =
                        (px[c] as f32 / 255.0 - MEAN[c]) / STD[c];
                }
            }
        }

        let mut sess = self.session.lock().unwrap();
        let input_name = sess.inputs().first().map(|i| i.name().to_string())
            .unwrap_or_else(|| "pixel_values".into());
        let output_name = sess.outputs().first().map(|o| o.name().to_string())
            .unwrap_or_else(|| "image_embeds".into());
        let value = ort::value::Value::from_array(input)
            .map_err(|e| SemanticError::Ort(format!("input tensor: {e}")))?;
        let outputs = sess.run(ort::inputs![input_name => value])
            .map_err(|e| SemanticError::Ort(format!("clip run: {e}")))?;
        let out = outputs.get(output_name.as_str())
            .ok_or_else(|| SemanticError::Ort(format!("output {output_name} missing")))?;
        let view = out.try_extract_array::<f32>()
            .map_err(|e| SemanticError::Ort(format!("extract embedding: {e}")))?;
        let emb: Vec<f32> = view.view().iter().copied().collect();
        let norm = emb.iter().map(|x| x * x).sum::<f32>().sqrt();
        Ok(if norm > 0.0 { emb.iter().map(|x| x / norm).collect() } else { emb })
    }
}