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

//! WAI — Web AI media transport + execution.
//!
//! What this crate IS:
//!   - the **envelope** (length-prefixed magic + JSON manifest + payload)
//!   - the **capability dispatch** model (capability is named, not shipped)
//!   - thin, well-tested **wrappers around SOTA standard codecs** for
//!     the zeroth condition (AVIF / JPEG-XL / PNG / Opus / FLAC / AV1 /
//!     zstd / XZ — all royalty-free, all production-grade, all the same
//!     libraries every modern media stack uses)
//!   - the **FFI surface** so this Rust crate is the SDK other languages
//!     consume (C ABI via `extern "C"`, cdylib/staticlib outputs)
//!
//! What this crate is NOT:
//!   - a from-scratch codec implementation. The zeroth-condition codecs
//!     are wrappers around the field's SOTA libraries. The standard's
//!     value is the envelope + capability model, not "rewrite JPEG".
//!
//! Neural capabilities (`wai.neural.*`) are declared here so the
//! dispatch table knows them; the actual neural model runs in whatever
//! ML runtime the sink has installed (PyTorch / Core ML / ONNX / etc.).

pub mod codecs;
pub mod container;
#[cfg(feature = "codecs")]
pub mod ffi;

// The integer-exact synthesis decoder lives at the crate root behind the
// light `neural_int` feature: pure i64 arithmetic, no ONNX runtime, so it
// builds (and its cross-hardware bit-exactness is checkable) without `ort`.
// The `neural` module re-exposes it via the `wai.neural.int_synth` capability.
#[cfg(feature = "neural_int")]
pub mod int_synth;

// Integer Gaussian-conditional entropy — the reusable scale→bucket primitive
// that makes a conditional codec's entropy decode hardware-deterministic (the
// integer half of the integer-entropy contract). Pure i64, no `ort`, like
// int_synth; the integer rANS + int_synth are the arithmetic-coder + synthesis
// halves a full `NeuralIntegerExact` conditional codec also needs.
#[cfg(feature = "neural_int")]
pub mod int_entropy;

// The integer rANS arithmetic coder (CompressAI rans64, precision 16) — a pure
// i64 entropy coder, no `ort`, so it sits with the other `neural_int` primitives
// rather than behind the ONNX `neural` feature that also consumes it. The
// encode+decode round-trip is the arithmetic-coder half of a `NeuralIntegerExact`
// conditional codec (with `int_entropy` for bucketing, `int_hyper_synth` for the
// transforms).
#[cfg(feature = "neural_int")]
pub mod rans;

// Integer-exact transform primitives (softmax/layernorm/matmul) — the missing
// piece for a NeuralIntegerExact *attention-based* decode (the SOTA / JPEG-AI
// MetaFormer class). Pure i64, no `ort`, like int_synth/int_entropy.
#[cfg(feature = "neural_int")]
pub mod int_transform;

// Integer-exact synthesis of a SOTA-class learned codec (per-channel int12 convs
// + integer IGDN) — the decode side of a NeuralIntegerExact conditional codec at
// SOTA quality. Pure i64, no `ort`. Runs the model the export tool emits.
#[cfg(feature = "neural_int")]
pub mod int_hyper_synth;

// Byte-exact integer TRANSFORMER kernels (softmax / attention / LayerNorm /
// L2-norm / ELU / Sigmoid / MDTA) — the pure-i64 arithmetic a learned ATTENTION
// image codec (JPEG-AI's Restormer synthesis) needs beyond the conv/IGDN stack
// above. NOTE: distinct from `int_transform` (a signal transform). Pure i64.
#[cfg(feature = "neural_int")]
pub mod int_transformer;

// Whole-net byte-exact integer JPEG-AI synthesis (Restormer, model_y) — threads the
// real traced topology through int_transformer + int_hyper_synth ops on sink-supplied
// 16-bit-PTQ weights → Y plane. Pure i64. (Structural; PSNR-vs-float pending.)
#[cfg(feature = "neural_int")]
pub mod jpegai_synthesis;

// Byte-exact mANS entropy decoder for JPEG-AI (port of the ISO reference
// decompressor.cpp) — two interleaved ANS states, sgm (y-latent) + fac (z-latent)
// paths. Pure integer ⇒ byte-identical everywhere.
#[cfg(feature = "neural_int")]
pub mod jpegai_mans;

// JPEG-AI decode-pipeline glue (dequant + reconstruction wiring).
#[cfg(feature = "neural_int")]
pub mod jpegai_pipeline;

// JPEG-AI eICCI post-filter: the Haar wavelet domain its CNN operates in.
#[cfg(feature = "neural_int")]
pub mod jpegai_icci;

// JPEG-AI bitstream container (marker segments) + the MSB-first header bit reader —
// the decoder's entry point.
#[cfg(feature = "neural_int")]
pub mod jpegai_bitstream;

// Deterministic integer generative synthesis (`wai.generate.*`) — the
// `intent=create` analog of the integer-exact codecs: a generator whose entire
// pipeline is integer/fixed-point, pinned to (model, prompt, seed, steps), so the
// generated image is byte-identical on every machine (its hash is portable, can
// go in a receipt). Pure i64, no float/ONNX; compiles to wasm32.
#[cfg(feature = "generate")]
pub mod generate;

// Deterministic integer inter-frame video (`wai.video.int_motion`) — byte-exact
// frame-by-frame decode: integer block motion compensation + integer-quantized
// residual over the rANS coder. Frame-state-equivalence (per-frame portable hash),
// the temporal extension of the integer contract. Pure i64, no float/ONNX.
#[cfg(feature = "video_int")]
pub mod int_video;

// Integer-exact LEARNED audio codec (`wai.audio.int_codec`) — a learned 1-D
// transform codec whose synthesis is a stack of integer transposed convolutions +
// integer IGDN, decoding to int16 PCM byte-identically on every machine
// (NeuralIntegerExact; *sample-equivalence*, a portable PCM hash). The audio analog
// of int_hyper. Pure i64, no float/ONNX; compiles to wasm32.
#[cfg(feature = "audio_int")]
pub mod int_audio;

// Integer-exact INTERACTIVE WORLD MODEL (`int_world`) — the byte-exact renderer for
// wai.world.model: an action-conditioned next-frame model (transition + synthesis are
// both int_hyper IntModels), so given (model, z0, action trace) every generated frame
// is byte-identical on every machine. Pure i64; compiles to wasm32.
#[cfg(feature = "world_int")]
pub mod int_world;

// Integer-exact LEARNED Gaussian-splat ATTRIBUTE codec (`int_splat`) — "the int_hyper
// of 3D": a learned transform over the [14,S,S] attribute image, integer synthesis →
// byte-exact attribute tensor (splat-equivalence). Pure i64; compiles to wasm32.
#[cfg(feature = "splat_int")]
pub mod int_splat;

#[cfg(feature = "neural")]
pub mod neural;
// The neural C FFI wraps the codec FFI's shared buffer/error types, so it is
// available only when the codec chain is too. The Rust neural decoders (above)
// need neither — they build and are tested without the C codec chain.
#[cfg(all(feature = "neural", feature = "codecs"))]
pub mod ffi_neural;

// Shared deterministic ONNX session loader, used by both ort-backed paths.
#[cfg(any(feature = "neural", feature = "semantic"))]
pub(crate) mod ort_session;

// The perceptual / generative-semantic path: transmit a compact
// "understanding" (CLIP embedding + thumb + edges), regenerate at the sink
// with any conforming generator, conform by CLIP similarity (not byte
// identity). The Rust port of the moved-out `wai-neural` prototype — the
// generator and embedder are sink-supplied traits, never bundled.
#[cfg(feature = "semantic")]
pub mod semantic;

// Interactive worlds (extensions/interactive-worlds): media that consumes
// input. A world is a deterministic content-addressed artifact; every event
// is a signed CausalOp; conformance is replay-equivalence — identical
// state hash at every checkpoint, no tolerance parameter. Pure Rust,
// closed-form fixed-point; no ML dep chain.
#[cfg(feature = "world")]
pub mod world;

// Object-based spatial audio (extensions/spatial-audio): the deterministic
// Fx mixdown engine. Conformance is mixdown-equivalence — same scene +
// declared layout → identical canonical PCM. Reuses the world `Fx` floor.
#[cfg(feature = "audio_scene")]
pub mod audio_scene;

// Haptics (extensions/haptics): touch as instructions — a keyframed
// intensity/sharpness envelope reconstructed in the Fx floor.
#[cfg(feature = "haptic")]
pub mod haptic;

// Volumetric / 4D temporal Gaussian splats (extensions/volumetric): each
// Gaussian carries keyframed tracks; the sink interpolates the cloud at
// frame N in the Fx floor (frame-state-equivalence). Reuses the world Fx
// floor + blake3; the EWA render stays the sink's.
#[cfg(feature = "splat4d")]
pub mod splat4d;

// Music as instructions (extensions/score): a score of note events
// reconstructed to a canonical PCM mixdown by a deterministic trig-free
// synth in the Fx floor (mixdown-equivalence). The soundfont/neural
// instrument render stays the sink's. Reuses the world Fx floor + blake3.
#[cfg(feature = "score")]
pub mod score;

// Asset provenance (extensions/provenance): the "verify whatever it
// carries" half — a signed, energy-accounted lineage for opaque captured
// content WAI does not reconstruct. Reuses the JWP Merkle + Ed25519.
#[cfg(feature = "provenance")]
pub mod provenance;

// wai.world.model — generative-world transport: a signed, replay-equivalent action
// trace (the world CausalOp log) + per-frame provenance. The wire carries the
// action trace + a declared model_id, not the pixels; the sink renders. The neural
// sibling of wai.world.replay.
#[cfg(feature = "world_model")]
pub mod world_model;

// Linear film (extensions/linear-film): "a film is a sealed, camera-on-rails
// world." The world replay (replay-equivalence) + a keyframed camera track in
// the Fx floor + a frame schedule; verified frame-by-frame (film-equivalence).
// Reuses the world engine + JWP Merkle + Ed25519; the render stays the sink's.
#[cfg(feature = "film")]
pub mod film;

pub use codecs::{
    sink_capabilities,
    CAP_AUDIO_FLAC, CAP_AUDIO_OPUS,
    CAP_IMAGE_AVIF, CAP_IMAGE_JPEG, CAP_IMAGE_JXL, CAP_IMAGE_PNG,
    CAP_NEURAL_BMSHJ2018, CAP_NEURAL_DAC, CAP_NEURAL_DCVC_RT, CAP_NEURAL_ENCODEC32,
    CAP_NEURAL_GLC, CAP_NEURAL_MIMI, CAP_NEURAL_VIDEO_BMSHJ2018, CAP_NEURAL_WAVTOKENIZER,
    CAP_TEXT_XZ, CAP_TEXT_ZSTD,
    CAP_VIDEO_AV1, CAP_VIDEO_AV1_LOSSLESS,
};
pub use container::{Conditioning, Manifest, ModelRequirement, Wai};

#[cfg(all(test, feature = "codecs"))]
mod integration_tests {
    use super::*;

    #[test]
    fn end_to_end_image_jxl_lossless() {
        // build an image, encode as JXL lossless, wrap in a WAI envelope,
        // round-trip the envelope, decode, byte-compare.
        let h = 64u32;
        let w = 64u32;
        let mut rgb = vec![0u8; (h * w * 3) as usize];
        for i in 0..h as usize {
            for j in 0..w as usize {
                let p = (i * w as usize + j) * 3;
                rgb[p] = (i * 4) as u8;
                rgb[p + 1] = (j * 4) as u8;
                rgb[p + 2] = ((i + j) * 2) as u8;
            }
        }
        let payload = codecs::image::jxl_encode(&rgb, h, w, None).unwrap();
        let manifest = Manifest {
            wai: "1.0".into(),
            media: "image".into(),
            intent: "replicate".into(),
            model_requirement: ModelRequirement {
                capability: CAP_IMAGE_JXL.into(),
                fallback: Some(CAP_IMAGE_PNG.into()),
                min_sim: None,
                sim_metric: None,
            },
            conditioning: Conditioning { kind: "jxl".into(), prior: None },
            target: serde_json::json!({"h": h, "w": w}),
        };
        let wai = Wai::new(manifest, payload);
        let bytes = wai.to_bytes().unwrap();
        let recovered = Wai::from_bytes(&bytes).unwrap();
        assert_eq!(recovered.capability(), CAP_IMAGE_JXL);
        let (rec_rgb, rh, rw) = codecs::image::jxl_decode(&recovered.payload).unwrap();
        assert_eq!((rh, rw), (h, w));
        assert_eq!(rec_rgb, rgb, "JXL lossless via WAI envelope must round-trip exactly");
    }
}