open-standards/wai/wai-rs/tests/neural_audio_vectors.rs

//! The neural audio route over the committed vectors of `tests/vectors/neural_audio/`
//! (SPEC §5, *Payloads of the neural audio codecs*): one envelope each of
//! `wai.neural.encodec24`, `wai.neural.encodec48`, `wai.neural.dac16` and
//! `wai.neural.dac24`, pinned to a tiny decoder with its model's geometry, and the decode
//! `tools/wai_neural_audio_vectors.py` computes for it on its own (the second
//! implementation). No model download: every input is committed.
//!
//! Each vector decodes through the dispatch to its expected samples, at its model's
//! sample rate and channel count; a payload outside the model is refused before the
//! decoder runs; a decoder the pin does not name is refused; the candidate
//! `wai.neural.soundstream` is never dispatched; and the C ABI, which has no form for
//! two channels, refuses the stereo capability rather than return it as mono.

#![cfg(feature = "neural")]

use std::path::{Path, PathBuf};

use serde_json::Value;
use wai::codecs::{is_candidate, is_dispatchable, CAP_NEURAL_SOUNDSTREAM};
use wai::container::Wai;
use wai::neural::{decode_envelope, decode_payload, decodes_to_several_channels, DecodeError, Decoded, ModelRegistry};

const VECTORS: [&str; 4] = ["dac16", "dac24", "encodec24", "encodec48"];

fn dir(name: &str) -> PathBuf {
    Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/neural_audio").join(name)
}

fn meta(name: &str) -> Value {
    serde_json::from_slice(&std::fs::read(dir(name).join("meta.json")).unwrap()).unwrap()
}

fn envelope(name: &str) -> Wai {
    Wai::from_bytes(&std::fs::read(dir(name).join("vector.wai")).unwrap()).unwrap()
}

fn registry(cap: &str, decoder_of: &str) -> ModelRegistry {
    ModelRegistry::new().register(cap, dir(decoder_of).join("decoder.onnx"))
}

/// The samples `expected.f32` holds: index i where i % every == 0 or i lies in an overlap.
fn selected(m: &Value) -> Vec<usize> {
    let n = m["n_samples"].as_u64().unwrap() as usize;
    let every = m["every"].as_u64().unwrap() as usize;
    let overlaps: Vec<(usize, usize)> = m["overlaps"].as_array().unwrap().iter()
        .map(|o| (o[0].as_u64().unwrap() as usize, o[1].as_u64().unwrap() as usize)).collect();
    (0..n).filter(|i| i % every == 0 || overlaps.iter().any(|&(a, b)| (a..b).contains(i))).collect()
}

fn invalid(r: Result<wai::neural::Dispatched, DecodeError>) -> String {
    match r {
        Err(DecodeError::InvalidPayload(why)) => why,
        Err(e) => panic!("refused, but not as an invalid payload: {e}"),
        Ok(_) => panic!("decoded"),
    }
}

#[test]
fn each_vector_decodes_to_the_samples_the_second_implementation_gives() {
    for name in VECTORS {
        let m = meta(name);
        let cap = m["capability"].as_str().unwrap();
        let channels = m["channels"].as_u64().unwrap() as usize;
        let n = m["n_samples"].as_u64().unwrap() as usize;
        let env = envelope(name);
        assert_eq!(env.capability(), cap, "{name}");
        // The envelope pins its decoder: the registry verifies the pin before the decode.
        let d = decode_envelope(&env, &registry(cap, name)).unwrap_or_else(|e| panic!("{name}: {e}"));
        assert_eq!(d.capability, cap);
        let Decoded::Audio(a) = d.output else { panic!("{name}: not audio") };
        assert_eq!(a.sample_rate as u64, m["sample_rate"].as_u64().unwrap(), "{name}");
        assert_eq!(a.channels as usize, channels, "{name}");
        assert_eq!(a.samples.len(), n * channels, "{name}: samples per channel");
        assert_eq!(decodes_to_several_channels(cap), channels > 1, "{name}");
        let raw = std::fs::read(dir(name).join("expected.f32")).unwrap();
        let want: Vec<f32> = raw.chunks(4).map(|b| f32::from_le_bytes(b.try_into().unwrap())).collect();
        let at = selected(&m);
        assert_eq!(want.len(), at.len() * channels, "{name}: expected.f32");
        let mut worst = 0f32;
        for (k, &i) in at.iter().enumerate() {
            for c in 0..channels {
                worst = worst.max((a.samples[i * channels + c] - want[k * channels + c]).abs());
            }
        }
        println!("  {name}: {n} samples × {channels} @ {} Hz, max |Δ| = {worst:.2e} over {} values", a.sample_rate, want.len());
        assert!(worst <= 1e-5, "{name}: the route's decode differs from the second implementation's by {worst}");
    }
}

#[test]
fn a_payload_outside_its_model_is_refused_before_the_decoder_runs() {
    let reg = registry("wai.neural.encodec24", "encodec24");
    let with = |f: &dyn Fn(&mut Vec<u8>)| {
        let mut e = envelope("encodec24");
        f(&mut e.payload);
        decode_envelope(&e, &reg)
    };
    // Three codebooks is no operating point of EnCodec 24 kHz.
    assert!(invalid(with(&|p| p[0..4].copy_from_slice(&3u32.to_le_bytes()))).contains("operating point"));
    // The value that marks a codebook not used is no code.
    assert!(invalid(with(&|p| p[12..14].copy_from_slice(&1024u16.to_le_bytes()))).contains("outside a codebook"));
    // More samples than the frames give.
    assert!(invalid(with(&|p| p[8..12].copy_from_slice(&1601u32.to_le_bytes()))).contains("samples per channel"));
    assert!(invalid(with(&|p| { p.pop(); })).contains("its header gives"));
    // EnCodec 48 kHz: a chunk's scale is finite and above 0.
    let mut e = envelope("encodec48");
    e.payload[16..20].copy_from_slice(&0f32.to_le_bytes());
    assert!(invalid(decode_envelope(&e, &registry("wai.neural.encodec48", "encodec48"))).contains("scale"));
    // DAC 16 kHz: its decoder gives 8 samples fewer than the frames' 320 each.
    let mut e = envelope("dac16");
    e.payload[8..12].copy_from_slice(&(4u32 * 320 - 7).to_le_bytes());
    assert!(invalid(decode_envelope(&e, &registry("wai.neural.dac16", "dac16"))).contains("samples per channel"));
}

#[test]
fn a_decoder_the_pin_does_not_name_is_refused() {
    // DAC 24 kHz's tiny decoder has EnCodec 24 kHz's input shape, but not its digest.
    let r = decode_envelope(&envelope("encodec24"), &registry("wai.neural.encodec24", "dac24"));
    assert!(matches!(r, Err(DecodeError::PriorMismatch(_))), "{:?}", r.map(|d| d.capability));
}

#[test]
fn the_candidate_soundstream_is_never_dispatched() {
    assert!(is_candidate(CAP_NEURAL_SOUNDSTREAM));
    assert!(!is_dispatchable(CAP_NEURAL_SOUNDSTREAM));
    // An envelope that names it, at a sink that holds a file for it, is inert there.
    let mut env = envelope("encodec24");
    env.manifest.model_requirement.capability = CAP_NEURAL_SOUNDSTREAM.to_string();
    env.manifest.conditioning.prior = None;
    let reg = registry(CAP_NEURAL_SOUNDSTREAM, "encodec24");
    assert!(matches!(decode_envelope(&env, &reg), Err(DecodeError::ModelNotRegistered(_))));
    // And the route has no decoder for it.
    let params = wai::neural::ParamSource::Unpinned(dir("encodec24").join("decoder.onnx"));
    assert!(matches!(decode_payload(CAP_NEURAL_SOUNDSTREAM, &env.payload, &params), Err(DecodeError::UnknownCapability(_))));
}

/// The C ABI has no form for two channels: the stereo capability is refused before
/// anything is decoded, never returned as mono; a mono one decodes.
#[cfg(feature = "ffi")]
#[test]
fn the_c_abi_refuses_two_channels_and_decodes_one() {
    use std::ffi::{CStr, CString};
    use wai::ffi::*;
    use wai::ffi_neural::*;
    let call = |name: &str| {
        let env = std::fs::read(dir(name).join("vector.wai")).unwrap();
        let path = CString::new(dir(name).join("decoder.onnx").to_str().unwrap()).unwrap();
        let mut out = WaiBuffer { data: std::ptr::null_mut(), len: 0, cap: 0 };
        let mut sr = 0u32;
        let code = unsafe { wai_neural_decode_audio(env.as_ptr(), env.len(), path.as_ptr(), &mut out, &mut sr) };
        // No error is set after a decode that succeeds: the pointer is null then.
        let last = wai_last_error();
        let err = if last.is_null() { String::new() } else { unsafe { CStr::from_ptr(last) }.to_string_lossy().into_owned() };
        let n = out.len / 4;
        if code == WAI_OK {
            unsafe { wai_buffer_free(out) };
        }
        (code, err, n, sr)
    };
    let (code, err, _, _) = call("encodec48");
    assert_eq!(code, WAI_ERR_CODEC, "{err}");
    assert!(err.contains("no wai_decoded_t kind"), "{err}");
    let (code, err, n, sr) = call("encodec24");
    assert_eq!(code, WAI_OK, "{err}");
    assert_eq!((n, sr), (1523, 24_000));
}