diff --git a/Cargo.lock b/Cargo.lock index 86d4bca6..767cac84 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2903,6 +2903,17 @@ dependencies = [ "web-time", ] +[[package]] +name = "lightcraft-faces" +version = "0.2.1" +dependencies = [ + "proptest", + "serde", + "serde_json", + "sha2", + "thiserror 2.0.21", +] + [[package]] name = "lightcraft-fetch" version = "0.4.0" diff --git a/Cargo.toml b/Cargo.toml index 77b20666..b2a3cb40 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ repository = "https://github.com/storytold/lightcraft" [workspace.dependencies] # internal lightcraft-geom = { path = "crates/geom" } +lightcraft-faces = { path = "crates/faces" } lightcraft-sysmem = { path = "crates/sysmem" } lightcraft-color = { path = "crates/color" } lightcraft-raster = { path = "crates/raster" } @@ -34,6 +35,7 @@ lightcraft-segment = { path = "crates/segment" } # foundation thiserror = "2" +sha2 = "0.10" serde = { version = "1", features = ["derive", "rc"] } serde_json = { version = "1", features = ["float_roundtrip"] } log = "0.4" diff --git a/crates/faces/Cargo.toml b/crates/faces/Cargo.toml new file mode 100644 index 00000000..da04c3d0 --- /dev/null +++ b/crates/faces/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "lightcraft-faces" +version.workspace = true +edition.workspace = true +license.workspace = true +description = "LightCraft faces (L1): face-model manifests, validation, the known-model table and an ONNX shape probe" + +[lints] +workspace = true + +[dependencies] +serde = { workspace = true } +serde_json = { workspace = true } +sha2 = { workspace = true } +thiserror = { workspace = true } + +[dev-dependencies] +proptest = { workspace = true } diff --git a/crates/faces/src/hash.rs b/crates/faces/src/hash.rs new file mode 100644 index 00000000..50f83be3 --- /dev/null +++ b/crates/faces/src/hash.rs @@ -0,0 +1,46 @@ +//! SHA-256 of a model file, streamed (a model can be hundreds of megabytes). + +use std::io::{self, Read}; +use std::path::Path; + +use sha2::{Digest, Sha256}; + +/// Lowercase hex SHA-256 of everything `r` yields. +pub fn sha256_reader(mut r: R) -> io::Result { + let mut h = Sha256::new(); + let mut buf = vec![0u8; 1 << 20]; + loop { + let n = r.read(&mut buf)?; + if n == 0 { + break; + } + h.update(buf.get(..n).unwrap_or(&[])); + } + Ok(hex(&h.finalize())) +} + +/// Lowercase hex SHA-256 of a file. +pub fn sha256_file(path: &Path) -> io::Result { + sha256_reader(io::BufReader::new(std::fs::File::open(path)?)) +} + +/// Lowercase hex SHA-256 of `bytes`. +pub fn sha256_hex(bytes: &[u8]) -> String { + hex(&Sha256::digest(bytes)) +} + +fn hex(bytes: &[u8]) -> String { + bytes.iter().map(|b| format!("{b:02x}")).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn known_digests() { + assert_eq!(sha256_hex(b""), "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"); + assert_eq!(sha256_hex(b"abc"), "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"); + assert_eq!(sha256_reader(&b"abc"[..]).unwrap(), sha256_hex(b"abc")); + } +} diff --git a/crates/faces/src/known.rs b/crates/faces/src/known.rs new file mode 100644 index 00000000..2eda8753 --- /dev/null +++ b/crates/faces/src/known.rs @@ -0,0 +1,136 @@ +//! The models LightCraft recognises by their SHA-256, with what is honestly known about each. +//! +//! The licence and provenance texts here are what the user reads before enabling a model, so they say what +//! is *not* known too. Weights are never bundled except where the bundling decision says so (the YuNet +//! detector, 232 KB, MIT); everything else is an opt-in the user installs themselves. + +use crate::manifest::{Colour, Commercial, InputSpec, Licence, ModelManifest, OutputSpec, Resize, Role, Thresholds}; + +/// Detector output decoders built into LightCraft. +pub const DECODERS: &[&str] = &["yunet-v2"]; + +pub const YUNET_SHA256: &str = "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4"; +pub const SFACE_SHA256: &str = "0ba9fbfa01b5270c96627c4ef784da859931e02f04419c829e83484087c34e79"; +pub const AURAFACE_SHA256: &str = "a7933ea5330113b01c9b60351d8f4c33003f145d8470ac5f0e52ee2effe25c60"; + +fn licence(name: &str, commercial: Commercial, url: &str, notice: &str) -> Licence { + Licence { name: name.into(), commercial, url: Some(url.into()), notice: notice.into() } +} + +/// YuNet 2023mar (OpenCV Zoo): the face detector LightCraft bundles. +pub fn yunet() -> ModelManifest { + ModelManifest { + id: "yunet-2023mar".into(), + name: "YuNet (face detector)".into(), + version: "2023mar".into(), + role: Role::Detector, + licence: licence( + "MIT", + Commercial::Yes, + "https://github.com/opencv/opencv_zoo/blob/main/models/face_detection_yunet/LICENSE", + "MIT licence, copyright Shiqi Yu. Bundled with LightCraft.", + ), + source: Some("https://github.com/opencv/opencv_zoo/tree/main/models/face_detection_yunet".into()), + sha256: Some(YUNET_SHA256.into()), + size_bytes: Some(232_589), + provenance: + "Trained on WIDER FACE, whose terms are not settled by the weights' MIT licence. A detector: it finds faces, it does not identify anyone." + .into(), + input: InputSpec { width: 640, height: 640, colour: Colour::Bgr, mean: [0.0; 3], std: [1.0; 3], resize: Resize::Letterbox }, + output: OutputSpec::Detector { decoder: "yunet-v2".into() }, + thresholds: Thresholds { match_cosine: None, score: Some(0.6), nms_iou: Some(0.3) }, + } +} + +/// SFace 2021dec (OpenCV Zoo). +pub fn sface() -> ModelManifest { + ModelManifest { + id: "sface-2021dec".into(), + name: "SFace (face recogniser)".into(), + version: "2021dec".into(), + role: Role::Embedder, + licence: licence( + "Apache-2.0 (as labelled by OpenCV Zoo)", + Commercial::Unknown, + "https://github.com/opencv/opencv_zoo/tree/main/models/face_recognition_sface", + "Labelled Apache-2.0, but what it was trained on is not documented, and two questions about commercial use (opencv_zoo issues 313 and 318) are unanswered. Fine to try for yourself; do not redistribute.", + ), + source: Some("https://github.com/opencv/opencv_zoo/tree/main/models/face_recognition_sface".into()), + sha256: Some(SFACE_SHA256.into()), + size_bytes: Some(38_696_353), + provenance: "Undocumented. The original SFace repository mentions CASIA-WebFace, VGGFace2 and MS1MV2.".into(), + input: InputSpec { width: 112, height: 112, colour: Colour::Rgb, mean: [0.0; 3], std: [1.0; 3], resize: Resize::Stretch }, + output: OutputSpec::Embedding { dim: 128 }, + thresholds: Thresholds::default(), + } +} + +/// AuraFace v1 `glintr100` (fal.ai). +pub fn auraface() -> ModelManifest { + ModelManifest { + id: "auraface-v1".into(), + name: "AuraFace v1 (face recogniser)".into(), + version: "1".into(), + role: Role::Embedder, + licence: licence( + "Apache-2.0", + Commercial::Yes, + "https://huggingface.co/fal/AuraFace-v1", + "Apache-2.0. Its training data is described only as a commercial dataset, with no dataset named and no consent statement, and it covers some ethnicities less well. About 261 MB, and slow on a CPU.", + ), + source: Some("https://huggingface.co/fal/AuraFace-v1".into()), + sha256: Some(AURAFACE_SHA256.into()), + size_bytes: Some(260_694_151), + provenance: "Undisclosed: \"a commercial dataset comprising face images from various sources\".".into(), + input: InputSpec::default(), + output: OutputSpec::Embedding { dim: 512 }, + thresholds: Thresholds::default(), + } +} + +/// Every model we know by hash. +pub fn all() -> Vec { + vec![yunet(), sface(), auraface()] +} + +/// The known model whose file has this SHA-256 (lowercase hex). +pub fn lookup(sha256: &str) -> Option { + all().into_iter().find(|m| m.sha256.as_deref() == Some(sha256)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::manifest::validate; + + #[test] + fn known_models_are_valid_and_unique() { + let all = all(); + for m in &all { + assert_eq!(validate(m), Ok(()), "{}", m.id); + } + let mut ids: Vec<_> = all.iter().map(|m| &m.id).collect(); + ids.sort(); + ids.dedup(); + assert_eq!(ids.len(), all.len(), "unique ids"); + let mut hashes: Vec<_> = all.iter().filter_map(|m| m.sha256.as_deref()).collect(); + hashes.sort(); + hashes.dedup(); + assert_eq!(hashes.len(), all.len(), "unique hashes"); + } + + #[test] + fn lookup_finds_by_hash() { + assert_eq!(lookup(AURAFACE_SHA256).map(|m| m.id), Some("auraface-v1".to_string())); + assert_eq!(lookup(YUNET_SHA256).map(|m| m.role), Some(Role::Detector)); + assert!(lookup(&"0".repeat(64)).is_none()); + assert!(lookup("").is_none()); + } + + #[test] + fn unresolved_models_say_so() { + assert_eq!(sface().licence.commercial, Commercial::Unknown); + assert!(sface().provenance.to_lowercase().contains("undocumented")); + assert!(auraface().provenance.to_lowercase().contains("undisclosed")); + } +} diff --git a/crates/faces/src/lib.rs b/crates/faces/src/lib.rs new file mode 100644 index 00000000..0ac55b66 --- /dev/null +++ b/crates/faces/src/lib.rs @@ -0,0 +1,18 @@ +//! Face models for LightCraft: what a model is ([`ModelManifest`]), whether to trust a description of one +//! ([`manifest::validate`]), the models we know by hash ([`known`]), and a probe that reads an unknown +//! `.onnx` file's input and output shapes without running it ([`onnx`]), so installing a model of your +//! own can be one drop and one licence prompt. +//! +//! Everything here treats model files and manifests as hostile input: sizes are capped, numbers must be +//! finite and sane, and a malformed file is an error, never a panic. Running a model (the `tract` runtime) +//! is not part of this crate yet. +#![forbid(unsafe_code)] +#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] + +pub mod hash; +pub mod known; +pub mod manifest; +pub mod onnx; +pub mod suggest; + +pub use manifest::{Colour, Commercial, InputSpec, Licence, ManifestError, ModelManifest, OutputSpec, Resize, Role, Thresholds}; diff --git a/crates/faces/src/manifest.rs b/crates/faces/src/manifest.rs new file mode 100644 index 00000000..22d70ddf --- /dev/null +++ b/crates/faces/src/manifest.rs @@ -0,0 +1,333 @@ +//! A face model's manifest: what the model is, what it needs and who may use it, as plain data. +//! +//! A manifest comes from the known-model table, from a suggestion made for a file the user dropped in, or +//! from a `face-model.json` next to a model. All of those are untrusted: [`validate`] must pass before a +//! manifest is used, and [`parse`] applies it after reading JSON of bounded size. + +use serde::{Deserialize, Serialize}; + +/// Largest `face-model.json` we read. +pub const MAX_MANIFEST_BYTES: usize = 64 * 1024; +/// Largest model file we accept (bytes). +pub const MAX_MODEL_BYTES: u64 = 4 * 1024 * 1024 * 1024; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Role { + /// Finds faces (boxes, and landmarks when the model has them). + Detector, + /// Turns an aligned face into a vector whose distance says whether two faces are the same person. + Embedder, +} + +/// Whether the weights may be used commercially. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Commercial { + Yes, + /// Research or personal use only. Never bundled or redistributed by LightCraft. + No, + /// The licence or the training data is unclear: the user is told so before installing. + Unknown, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct Licence { + /// SPDX id or a short name ("MIT", "Apache-2.0", "InsightFace non-commercial research"). + pub name: String, + pub commercial: Commercial, + pub url: Option, + /// Plain-language terms shown before the model is enabled. + pub notice: String, +} + +impl Default for Licence { + fn default() -> Self { + Self { name: String::new(), commercial: Commercial::Unknown, url: None, notice: String::new() } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Colour { + Rgb, + Bgr, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Resize { + /// Stretch the crop to the input size (aligned face crops). + Stretch, + /// Keep the aspect ratio and pad (whole photos for a detector). + Letterbox, +} + +/// How an image becomes the model's input: NCHW float, `(pixel − mean) / std` per channel, pixel in 0..255. +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct InputSpec { + pub width: u32, + pub height: u32, + pub colour: Colour, + pub mean: [f32; 3], + pub std: [f32; 3], + pub resize: Resize, +} + +impl Default for InputSpec { + /// ArcFace-style: 112 × 112, RGB, `(x − 127.5) / 127.5`. + fn default() -> Self { + Self { width: 112, height: 112, colour: Colour::Rgb, mean: [127.5; 3], std: [127.5; 3], resize: Resize::Stretch } + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "camelCase")] +pub enum OutputSpec { + /// One vector of `dim` numbers per face. + Embedding { dim: u32 }, + /// A detector whose outputs a built-in decoder understands (`"yunet-v2"`). + Detector { decoder: String }, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct Thresholds { + /// Embedders: cosine similarity at or above which two faces are suggested as a match. + pub match_cosine: Option, + /// Detectors: lowest detection score kept. + pub score: Option, + /// Detectors: overlap above which two boxes are one face. + pub nms_iou: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ModelManifest { + /// Lowercase letters, digits, `.` `_` `-`. + pub id: String, + pub name: String, + pub version: String, + pub role: Role, + #[serde(default)] + pub licence: Licence, + /// Where the weights come from (an `https://` page or file), for the user's information. + #[serde(default)] + pub source: Option, + /// Lowercase hex SHA-256 of the `.onnx` file, when known. + #[serde(default)] + pub sha256: Option, + #[serde(default)] + pub size_bytes: Option, + /// What the model was trained on, in plain words; says "undisclosed" when it is. + #[serde(default)] + pub provenance: String, + #[serde(default)] + pub input: InputSpec, + pub output: OutputSpec, + #[serde(default)] + pub thresholds: Thresholds, +} + +#[derive(Debug, PartialEq, thiserror::Error)] +pub enum ManifestError { + #[error("`{0}` is not valid: {1}")] + Field(&'static str, String), + #[error("not a face-model manifest: {0}")] + Parse(String), + #[error("a manifest may be at most {MAX_MANIFEST_BYTES} bytes")] + TooLarge, +} + +fn bad(field: &'static str, why: impl Into) -> Result { + Err(ManifestError::Field(field, why.into())) +} + +fn text(field: &'static str, s: &str, min: usize, max: usize) -> Result<(), ManifestError> { + let n = s.chars().count(); + if n < min || n > max { + return bad(field, format!("must be {min} to {max} characters")); + } + if s.chars().any(|c| c.is_control() && c != '\n') { + return bad(field, "must not contain control characters"); + } + Ok(()) +} + +fn url(field: &'static str, u: &Option) -> Result<(), ManifestError> { + match u { + None => Ok(()), + Some(u) if u.len() <= 500 && u.starts_with("https://") && !u.chars().any(|c| c.is_control() || c.is_whitespace()) => Ok(()), + Some(_) => bad(field, "must be an https:// address of at most 500 characters"), + } +} + +fn ratio(field: &'static str, v: Option, lo: f32, hi: f32) -> Result<(), ManifestError> { + match v { + Some(v) if !(v.is_finite() && (lo..=hi).contains(&v)) => bad(field, format!("must be a number from {lo} to {hi}")), + _ => Ok(()), + } +} + +/// Check a manifest from any source. Everything the rest of LightCraft relies on (sizes it will allocate, +/// numbers it will divide by, text it will show) is bounded here. +pub fn validate(m: &ModelManifest) -> Result<(), ManifestError> { + let id_ok = + (1..=64).contains(&m.id.len()) && m.id.bytes().all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || matches!(b, b'.' | b'_' | b'-')); + if !id_ok { + return bad("id", "must be 1 to 64 characters: lowercase letters, digits, '.', '_' or '-'"); + } + text("name", &m.name, 1, 120)?; + text("version", &m.version, 1, 40)?; + text("licence.name", &m.licence.name, 0, 80)?; + text("licence.notice", &m.licence.notice, 0, 800)?; + text("provenance", &m.provenance, 0, 800)?; + url("licence.url", &m.licence.url)?; + url("source", &m.source)?; + if let Some(h) = &m.sha256 + && !(h.len() == 64 && h.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))) + { + return bad("sha256", "must be 64 lowercase hexadecimal digits"); + } + if m.size_bytes.is_some_and(|s| s == 0 || s > MAX_MODEL_BYTES) { + return bad("sizeBytes", format!("must be 1 to {MAX_MODEL_BYTES}")); + } + let i = &m.input; + if !(16..=4096).contains(&i.width) || !(16..=4096).contains(&i.height) { + return bad("input.width/height", "must be 16 to 4096"); + } + if i.mean.iter().any(|v| !v.is_finite() || v.abs() > 1000.0) { + return bad("input.mean", "must be finite numbers within ±1000"); + } + if i.std.iter().any(|v| !v.is_finite() || v.abs() < 1e-3 || v.abs() > 1000.0) { + return bad("input.std", "must be finite, not zero, within ±1000"); + } + match &m.output { + OutputSpec::Embedding { dim } => { + if m.role != Role::Embedder { + return bad("output", "an embedding output needs the embedder role"); + } + if !(1..=4096).contains(dim) { + return bad("output.dim", "must be 1 to 4096"); + } + } + OutputSpec::Detector { decoder } => { + if m.role != Role::Detector { + return bad("output", "a detector output needs the detector role"); + } + if !crate::known::DECODERS.contains(&decoder.as_str()) { + return bad("output.decoder", format!("unknown decoder `{}`", decoder.chars().take(40).collect::())); + } + } + } + ratio("thresholds.matchCosine", m.thresholds.match_cosine, -1.0, 1.0)?; + ratio("thresholds.score", m.thresholds.score, 0.0, 1.0)?; + ratio("thresholds.nmsIou", m.thresholds.nms_iou, 0.0, 1.0)?; + Ok(()) +} + +/// Read a `face-model.json` (at most [`MAX_MANIFEST_BYTES`]) and validate it. +pub fn parse(bytes: &[u8]) -> Result { + if bytes.len() > MAX_MANIFEST_BYTES { + return Err(ManifestError::TooLarge); + } + let m: ModelManifest = serde_json::from_slice(bytes).map_err(|e| ManifestError::Parse(e.to_string()))?; + validate(&m)?; + Ok(m) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn good() -> ModelManifest { + ModelManifest { + id: "example-r50".into(), + name: "Example R50".into(), + version: "1".into(), + role: Role::Embedder, + licence: Licence { + name: "MIT".into(), + commercial: Commercial::Yes, + url: Some("https://example.org/licence".into()), + notice: String::new(), + }, + source: Some("https://example.org/model.onnx".into()), + sha256: Some("a".repeat(64)), + size_bytes: Some(1234), + provenance: "synthetic".into(), + input: InputSpec::default(), + output: OutputSpec::Embedding { dim: 512 }, + thresholds: Thresholds { match_cosine: Some(0.4), ..Default::default() }, + } + } + + #[test] + fn a_good_manifest_passes_and_round_trips() { + let m = good(); + assert_eq!(validate(&m), Ok(())); + let json = serde_json::to_vec(&m).unwrap(); + assert_eq!(parse(&json).unwrap(), m); + } + + #[test] + fn every_field_is_bounded() { + let cases: Vec<(&str, Box)> = vec![ + ("empty id", Box::new(|m| m.id = String::new())), + ("uppercase id", Box::new(|m| m.id = "Bad".into())), + ("path in id", Box::new(|m| m.id = "../etc".into())), + ("long id", Box::new(|m| m.id = "a".repeat(65))), + ("empty name", Box::new(|m| m.name = String::new())), + ("control in name", Box::new(|m| m.name = "a\u{7}b".into())), + ("huge notice", Box::new(|m| m.licence.notice = "x".repeat(801))), + ("http source", Box::new(|m| m.source = Some("http://example.org/m.onnx".into()))), + ("file source", Box::new(|m| m.source = Some("file:///etc/passwd".into()))), + ("space in url", Box::new(|m| m.source = Some("https://a b".into()))), + ("short sha", Box::new(|m| m.sha256 = Some("abc".into()))), + ("uppercase sha", Box::new(|m| m.sha256 = Some("A".repeat(64)))), + ("zero size", Box::new(|m| m.size_bytes = Some(0))), + ("huge size", Box::new(|m| m.size_bytes = Some(u64::MAX))), + ("tiny input", Box::new(|m| m.input.width = 4)), + ("huge input", Box::new(|m| m.input.height = 1 << 20)), + ("nan mean", Box::new(|m| m.input.mean[1] = f32::NAN)), + ("inf mean", Box::new(|m| m.input.mean[0] = f32::INFINITY)), + ("zero std", Box::new(|m| m.input.std[2] = 0.0)), + ("nan std", Box::new(|m| m.input.std[0] = f32::NAN)), + ("zero dim", Box::new(|m| m.output = OutputSpec::Embedding { dim: 0 })), + ("huge dim", Box::new(|m| m.output = OutputSpec::Embedding { dim: 10_000_000 })), + ("role mismatch", Box::new(|m| m.role = Role::Detector)), + ( + "unknown decoder", + Box::new(|m| { + m.role = Role::Detector; + m.output = OutputSpec::Detector { decoder: "made-up".into() }; + }), + ), + ("nan threshold", Box::new(|m| m.thresholds.match_cosine = Some(f32::NAN))), + ("threshold range", Box::new(|m| m.thresholds.score = Some(1.5))), + ]; + for (what, f) in cases { + let mut m = good(); + f(&mut m); + assert!(validate(&m).is_err(), "{what} must be rejected"); + } + } + + #[test] + fn hostile_json_is_an_error() { + assert!(matches!(parse(&vec![b' '; MAX_MANIFEST_BYTES + 1]), Err(ManifestError::TooLarge))); + for bytes in [ + &b""[..], + b"null", + b"[]", + b"{", + b"{\"id\": 1}", + b"\xff\xfe\x00", + b"{\"id\":\"x\",\"name\":\"n\",\"version\":\"1\",\"role\":\"embedder\"}", + ] { + assert!(parse(bytes).is_err(), "{:?}", String::from_utf8_lossy(bytes)); + } + } +} diff --git a/crates/faces/src/onnx.rs b/crates/faces/src/onnx.rs new file mode 100644 index 00000000..22e04dcf --- /dev/null +++ b/crates/faces/src/onnx.rs @@ -0,0 +1,526 @@ +//! A streaming probe of an ONNX file: the graph's inputs and outputs (name, element type, shape), without +//! running or even fully loading the model. +//! +//! ONNX is protobuf. We read only the few fields we need (the producer, the opset, and the graph's input, +//! output and initializer names) and seek past everything else, so a 260 MB model costs a few kilobytes and +//! a few milliseconds. The file is hostile input: every length is checked against its parent message and +//! the file size, the number of fields read is capped, strings are bounded, and nothing indexes or unwraps. + +use std::collections::HashSet; +use std::io::{self, Read, Seek, SeekFrom}; +use std::path::Path; + +const MAX_FIELDS: u32 = 5_000_000; +const MAX_STRING: u64 = 4096; +const MAX_DIMS: usize = 16; +const MAX_TENSORS: usize = 64; +/// Older exporters list every weight as a graph input too; that many are read before the weights are filtered out. +const MAX_LISTED_INPUTS: usize = 200_000; +const MAX_INITIALIZERS: usize = 1_000_000; + +#[derive(Clone, Debug, PartialEq)] +pub enum Dim { + Fixed(u64), + /// A named or unknown size (typically the batch). + Dynamic(String), +} + +#[derive(Clone, Debug, PartialEq)] +pub struct TensorInfo { + pub name: String, + /// ONNX element type (1 = float32). + pub elem_type: u32, + pub shape: Vec, +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct OnnxInfo { + pub producer: String, + /// The default-domain opset version. + pub opset: Option, + /// Real inputs only (weights that older exporters also list as inputs are left out). + pub inputs: Vec, + pub outputs: Vec, +} + +#[derive(Debug, thiserror::Error)] +pub enum ProbeError { + #[error("not an ONNX model: {0}")] + Malformed(&'static str), + #[error("could not read the model: {0}")] + Io(#[from] io::Error), +} + +type Result = std::result::Result; + +enum Wire { + Varint(u64), + /// A length-delimited field: its bytes are `start..end` and the reader sits at `start`. + Len { + start: u64, + end: u64, + }, + /// Fixed-width data already skipped. + Skipped, +} + +struct Walker<'a, R> { + r: &'a mut R, + total: u64, + fields: u32, +} + +impl Walker<'_, R> { + fn pos(&mut self) -> Result { + Ok(self.r.stream_position()?) + } + + fn seek_to(&mut self, at: u64) -> Result<()> { + self.r.seek(SeekFrom::Start(at))?; + Ok(()) + } + + fn byte(&mut self) -> Result { + let mut b = [0u8; 1]; + self.r + .read_exact(&mut b) + .map_err(|e| if e.kind() == io::ErrorKind::UnexpectedEof { ProbeError::Malformed("file ends inside a field") } else { e.into() })?; + let [x] = b; + Ok(x) + } + + fn varint(&mut self) -> Result { + let mut v = 0u64; + for shift in (0..70u32).step_by(7) { + let b = self.byte()?; + if shift < 64 { + v |= u64::from(b & 0x7f) << shift; + } + if b & 0x80 == 0 { + return Ok(v); + } + } + Err(ProbeError::Malformed("number too long")) + } + + /// The next field of the message that ends at `end`, or `None` at its end. + fn next(&mut self, end: u64) -> Result> { + let here = self.pos()?; + if here >= end { + return Ok(None); + } + self.fields += 1; + if self.fields > MAX_FIELDS { + return Err(ProbeError::Malformed("too many fields")); + } + let tag = self.varint()?; + let number = u32::try_from(tag >> 3).map_err(|_| ProbeError::Malformed("field number too large"))?; + let wire = match tag & 7 { + 0 => Wire::Varint(self.varint()?), + 1 | 5 => { + let n = if tag & 7 == 1 { 8 } else { 4 }; + let at = self.pos()?.checked_add(n).ok_or(ProbeError::Malformed("offset overflow"))?; + if at > end { + return Err(ProbeError::Malformed("field runs past its message")); + } + self.seek_to(at)?; + Wire::Skipped + } + 2 => { + let len = self.varint()?; + let start = self.pos()?; + let stop = start.checked_add(len).ok_or(ProbeError::Malformed("offset overflow"))?; + if stop > end || stop > self.total { + return Err(ProbeError::Malformed("field runs past its message")); + } + Wire::Len { start, end: stop } + } + _ => return Err(ProbeError::Malformed("unsupported field type")), + }; + Ok(Some((number, wire))) + } + + /// A bounded UTF-8 string occupying `start..end` (empty when it is longer than we care about). + fn string(&mut self, start: u64, end: u64) -> Result { + let len = end.saturating_sub(start); + if len > MAX_STRING { + return Ok(String::new()); + } + let mut buf = Vec::new(); + self.r.by_ref().take(len).read_to_end(&mut buf)?; + Ok(String::from_utf8_lossy(&buf).into_owned()) + } +} + +fn dim(w: &mut Walker, end: u64) -> Result> { + let mut out = None; + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Varint(v)) => out = Some(Dim::Fixed(v)), + (2, Wire::Len { start, end: e }) => { + out = Some(Dim::Dynamic(w.string(start, e)?)); + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(out) +} + +fn shape(w: &mut Walker, end: u64) -> Result> { + let mut dims = Vec::new(); + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { end: e, .. }) => { + if dims.len() >= MAX_DIMS { + return Err(ProbeError::Malformed("too many dimensions")); + } + dims.push(dim(w, e)?.unwrap_or(Dim::Dynamic(String::new()))); + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(dims) +} + +fn tensor_type(w: &mut Walker, end: u64, info: &mut TensorInfo) -> Result<()> { + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Varint(v)) => info.elem_type = u32::try_from(v).unwrap_or(0), + (2, Wire::Len { end: e, .. }) => { + info.shape = shape(w, e)?; + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(()) +} + +fn type_proto(w: &mut Walker, end: u64, info: &mut TensorInfo) -> Result<()> { + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { end: e, .. }) => { + tensor_type(w, e, info)?; + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(()) +} + +fn value_info(w: &mut Walker, end: u64) -> Result { + let mut info = TensorInfo { name: String::new(), elem_type: 0, shape: Vec::new() }; + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { start, end: e }) => { + info.name = w.string(start, e)?; + w.seek_to(e)?; + } + (2, Wire::Len { end: e, .. }) => { + type_proto(w, e, &mut info)?; + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(info) +} + +/// The `name` (field 8) of a weight tensor; its bulk data is skipped, never read. +fn tensor_name(w: &mut Walker, end: u64) -> Result> { + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (8, Wire::Len { start, end: e }) => return Ok(Some(w.string(start, e)?)), + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(None) +} + +fn graph(w: &mut Walker, end: u64, info: &mut OnnxInfo, weights: &mut HashSet) -> Result<()> { + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (5, Wire::Len { end: e, .. }) => { + if weights.len() < MAX_INITIALIZERS + && let Some(name) = tensor_name(w, e)? + { + weights.insert(name); + } + w.seek_to(e)?; + } + (11, Wire::Len { end: e, .. }) => { + if info.inputs.len() >= MAX_LISTED_INPUTS { + return Err(ProbeError::Malformed("too many inputs")); + } + info.inputs.push(value_info(w, e)?); + w.seek_to(e)?; + } + (12, Wire::Len { end: e, .. }) => { + if info.outputs.len() >= MAX_TENSORS { + return Err(ProbeError::Malformed("too many outputs")); + } + info.outputs.push(value_info(w, e)?); + w.seek_to(e)?; + } + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(()) +} + +fn opset(w: &mut Walker, end: u64) -> Result> { + let (mut domain, mut version) = (String::new(), None); + while let Some((n, f)) = w.next(end)? { + match (n, f) { + (1, Wire::Len { start, end: e }) => { + domain = w.string(start, e)?; + w.seek_to(e)?; + } + (2, Wire::Varint(v)) => version = Some(v), + (_, Wire::Len { end: e, .. }) => w.seek_to(e)?, + _ => {} + } + } + Ok(if domain.is_empty() || domain == "ai.onnx" { version } else { None }) +} + +/// Read an ONNX model's graph inputs and outputs from any seekable source. +pub fn probe_reader(r: &mut R) -> Result { + let total = r.seek(SeekFrom::End(0))?; + r.seek(SeekFrom::Start(0))?; + if total == 0 { + return Err(ProbeError::Malformed("the file is empty")); + } + let mut w = Walker { r, total, fields: 0 }; + let (mut info, mut weights, mut saw_graph) = (OnnxInfo::default(), HashSet::new(), false); + while let Some((n, f)) = w.next(total)? { + match (n, f) { + (2, Wire::Len { start, end }) => { + info.producer = w.string(start, end)?; + w.seek_to(end)?; + } + (7, Wire::Len { end, .. }) => { + saw_graph = true; + graph(&mut w, end, &mut info, &mut weights)?; + w.seek_to(end)?; + } + (8, Wire::Len { end, .. }) => { + if let Some(v) = opset(&mut w, end)? { + info.opset = Some(v); + } + w.seek_to(end)?; + } + (_, Wire::Len { end, .. }) => w.seek_to(end)?, + _ => {} + } + } + if !saw_graph { + return Err(ProbeError::Malformed("no graph")); + } + info.inputs.retain(|t| !weights.contains(&t.name)); + if info.inputs.len() > MAX_TENSORS { + return Err(ProbeError::Malformed("too many inputs")); + } + if info.inputs.is_empty() || info.outputs.is_empty() { + return Err(ProbeError::Malformed("the graph has no inputs or no outputs")); + } + Ok(info) +} + +/// Read an ONNX file's graph inputs and outputs. +pub fn probe_path(path: &Path) -> Result { + probe_reader(&mut io::BufReader::new(std::fs::File::open(path)?)) +} + +#[cfg(test)] +pub(crate) mod tests { + use super::*; + use std::io::Cursor; + + fn varint(mut v: u64, out: &mut Vec) { + loop { + let b = (v & 0x7f) as u8; + v >>= 7; + if v == 0 { + out.push(b); + return; + } + out.push(b | 0x80); + } + } + fn len_field(n: u64, body: &[u8], out: &mut Vec) { + varint(n << 3 | 2, out); + varint(body.len() as u64, out); + out.extend_from_slice(body); + } + fn varint_field(n: u64, v: u64, out: &mut Vec) { + varint(n << 3, out); + varint(v, out); + } + fn value_info_bytes(name: &str, elem: u64, dims: &[std::result::Result]) -> Vec { + let mut shape = Vec::new(); + for d in dims { + let mut dim = Vec::new(); + match d { + Ok(v) => varint_field(1, *v, &mut dim), + Err(p) => len_field(2, p.as_bytes(), &mut dim), + } + len_field(1, &dim, &mut shape); + } + let mut tensor = Vec::new(); + varint_field(1, elem, &mut tensor); + len_field(2, &shape, &mut tensor); + let mut ty = Vec::new(); + len_field(1, &tensor, &mut ty); + let mut vi = Vec::new(); + len_field(1, name.as_bytes(), &mut vi); + len_field(2, &ty, &mut vi); + vi + } + + /// A small valid model: input `data` [N, 3, 112, 112] float, an old-style weight also listed as an + /// input, a weight with bulk data, output `emb` [N, 512], ai.onnx opset 17. + pub(crate) fn model_bytes(output_dim: u64) -> Vec { + let mut graph = Vec::new(); + len_field(11, &value_info_bytes("data", 1, &[Err("N"), Ok(3), Ok(112), Ok(112)]), &mut graph); + len_field(11, &value_info_bytes("w", 1, &[Ok(4)]), &mut graph); + let mut tensor = Vec::new(); + varint_field(1, 4, &mut tensor); + len_field(8, b"w", &mut tensor); + len_field(9, &[7u8; 5000], &mut tensor); + len_field(5, &tensor, &mut graph); + len_field(12, &value_info_bytes("emb", 1, &[Err("N"), Ok(output_dim)]), &mut graph); + let mut ops = Vec::new(); + len_field(1, b"", &mut ops); + varint_field(2, 17, &mut ops); + let mut model = Vec::new(); + varint_field(1, 8, &mut model); + len_field(2, b"pytorch", &mut model); + len_field(7, &graph, &mut model); + len_field(8, &ops, &mut model); + model + } + + #[test] + fn reads_inputs_outputs_and_skips_weights() { + let info = probe_reader(&mut Cursor::new(model_bytes(512))).unwrap(); + assert_eq!(info.producer, "pytorch"); + assert_eq!(info.opset, Some(17)); + assert_eq!(info.inputs.len(), 1, "the weight listed as an input is left out"); + assert_eq!(info.inputs[0].name, "data"); + assert_eq!(info.inputs[0].elem_type, 1); + assert_eq!(info.inputs[0].shape, vec![Dim::Dynamic("N".into()), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]); + assert_eq!(info.outputs[0].shape, vec![Dim::Dynamic("N".into()), Dim::Fixed(512)]); + } + + #[test] + fn old_style_files_listing_hundreds_of_weights_as_inputs_still_probe() { + let mut graph = Vec::new(); + len_field(11, &value_info_bytes("data", 1, &[Err("N"), Ok(3), Ok(112), Ok(112)]), &mut graph); + for i in 0..500 { + let name = format!("weight{i}"); + len_field(11, &value_info_bytes(&name, 1, &[Ok(4)]), &mut graph); + let mut tensor = Vec::new(); + len_field(8, name.as_bytes(), &mut tensor); + len_field(5, &tensor, &mut graph); + } + len_field(12, &value_info_bytes("emb", 1, &[Err("N"), Ok(128)]), &mut graph); + let mut m = Vec::new(); + len_field(7, &graph, &mut m); + let info = probe_reader(&mut Cursor::new(m)).unwrap(); + assert_eq!(info.inputs.len(), 1); + assert_eq!(info.inputs[0].name, "data"); + } + + #[test] + fn many_real_inputs_are_refused() { + let mut graph = Vec::new(); + for i in 0..70 { + len_field(11, &value_info_bytes(&format!("in{i}"), 1, &[Ok(4)]), &mut graph); + } + len_field(12, &value_info_bytes("emb", 1, &[Ok(4)]), &mut graph); + let mut m = Vec::new(); + len_field(7, &graph, &mut m); + assert!(probe_reader(&mut Cursor::new(m)).is_err()); + } + + #[test] + fn truncated_at_every_length_never_panics() { + let bytes = model_bytes(128); + for n in 0..bytes.len() { + let _ = probe_reader(&mut Cursor::new(bytes[..n].to_vec())); + } + assert!(probe_reader(&mut Cursor::new(bytes)).is_ok()); + } + + #[test] + fn hostile_files_are_errors() { + let mut lying = Vec::new(); + varint(7 << 3 | 2, &mut lying); + varint(1 << 40, &mut lying); // a graph 1 TB long in a 10-byte file + let mut overflow = Vec::new(); + varint(7 << 3 | 2, &mut overflow); + varint(u64::MAX, &mut overflow); + let mut long_number = vec![0x08]; + long_number.extend_from_slice(&[0xff; 11]); + for (what, bytes) in [ + ("empty", Vec::new()), + ("zeros", vec![0u8; 64]), + ("text", b"this is not an onnx model at all".to_vec()), + ("lying length", lying), + ("overflowing length", overflow), + ("endless number", long_number), + ("no graph", { + let mut m = Vec::new(); + len_field(2, b"x", &mut m); + m + }), + ("graph without io", { + let mut m = Vec::new(); + len_field(7, &[], &mut m); + m + }), + ] { + assert!(probe_reader(&mut Cursor::new(bytes)).is_err(), "{what}"); + } + } + + #[test] + fn too_many_dimensions_or_tensors_are_errors() { + let dims: Vec> = (0..40).map(Ok).collect(); + let mut graph = Vec::new(); + len_field(11, &value_info_bytes("data", 1, &dims), &mut graph); + len_field(12, &value_info_bytes("out", 1, &[Ok(1)]), &mut graph); + let mut m = Vec::new(); + len_field(7, &graph, &mut m); + assert!(probe_reader(&mut Cursor::new(m)).is_err()); + } + + proptest::proptest! { + #[test] + fn arbitrary_bytes_never_panic(bytes in proptest::collection::vec(proptest::prelude::any::(), 0..400)) { + let _ = probe_reader(&mut Cursor::new(bytes)); + } + + #[test] + fn a_valid_model_with_flipped_bytes_never_panics(flips in proptest::collection::vec((0usize..400, proptest::prelude::any::()), 1..6)) { + let mut bytes = model_bytes(512); + for (at, v) in flips { + if let Some(b) = bytes.get_mut(at) { + *b = v; + } + } + let _ = probe_reader(&mut Cursor::new(bytes)); + } + } +} diff --git a/crates/faces/src/suggest.rs b/crates/faces/src/suggest.rs new file mode 100644 index 00000000..eddd07ed --- /dev/null +++ b/crates/faces/src/suggest.rs @@ -0,0 +1,231 @@ +//! What to do with a model file the user drops in: recognise it by hash, or look at its shape and propose a +//! manifest the user confirms, or say plainly why it cannot be used yet. + +use crate::manifest::{Commercial, InputSpec, Licence, ModelManifest, OutputSpec, Role, Thresholds}; +use crate::onnx::{Dim, OnnxInfo, TensorInfo}; + +/// ONNX element type for float32. +const FLOAT32: u32 = 1; + +#[derive(Clone, Debug, PartialEq)] +pub enum Suggestion { + /// A model we know by its SHA-256: its manifest as it is. + Known(ModelManifest), + /// An unknown model shaped like a face embedder. `assumptions` is what we guessed, in plain words. + Draft { manifest: ModelManifest, assumptions: Vec }, + /// Not usable yet, and why. + Unsupported(String), +} + +fn fixed(d: &Dim) -> Option { + match d { + Dim::Fixed(v) => Some(*v), + Dim::Dynamic(_) => None, + } +} + +/// An id made of the file name and the start of the hash: lowercase letters, digits, `.` `_` `-`. +fn make_id(file_name: &str, sha256: &str) -> String { + let stem = file_name.rsplit_once('.').map_or(file_name, |(s, _)| s); + let clean: String = stem.chars().map(|c| if c.is_ascii_alphanumeric() { c.to_ascii_lowercase() } else { '-' }).collect(); + let clean: String = clean.trim_matches('-').chars().take(40).collect(); + let short: String = sha256.chars().take(8).collect(); + match (clean.is_empty(), short.is_empty()) { + (true, true) => "custom".into(), + (true, false) => format!("custom-{short}"), + (false, true) => format!("custom-{clean}"), + (false, false) => format!("custom-{clean}-{short}"), + } +} + +fn display_name(file_name: &str) -> String { + let stem = file_name.rsplit_once('.').map_or(file_name, |(s, _)| s); + let name: String = stem.chars().filter(|c| !c.is_control()).take(60).collect(); + if name.trim().is_empty() { "Custom face model".into() } else { format!("{} (your model)", name.trim()) } +} + +/// The input of a face embedder: float32, `[batch, 3, height, width]`. +fn face_input(t: &TensorInfo) -> Option<(u32, u32)> { + if t.elem_type != FLOAT32 || t.shape.len() != 4 || t.shape.get(1).and_then(fixed) != Some(3) { + return None; + } + let size = |d: Option<&Dim>| match d.and_then(fixed) { + Some(v) if (16..=1024).contains(&v) => u32::try_from(v).ok(), + Some(_) => None, + None => Some(112), // a dynamic size: the common face-crop size + }; + Some((size(t.shape.get(3))?, size(t.shape.get(2))?)) +} + +/// Propose what to do with a probed model. `sha256` and `size` describe the file, `file_name` names it. +pub fn suggest(info: &OnnxInfo, sha256: &str, size: u64, file_name: &str) -> Suggestion { + if let Some(m) = crate::known::lookup(sha256) { + return Suggestion::Known(m); + } + let [input] = info.inputs.as_slice() else { + return Suggestion::Unsupported(format!("the model has {} inputs; face models here take one image input", info.inputs.len())); + }; + let Some((width, height)) = face_input(input) else { + return Suggestion::Unsupported("its input is not a float image of shape [batch, 3, height, width]".into()); + }; + let [output] = info.outputs.as_slice() else { + return Suggestion::Unsupported( + "it has several outputs, like a face detector; only models that give one vector per face are supported for now".into(), + ); + }; + if output.elem_type != FLOAT32 || output.shape.len() != 2 { + return Suggestion::Unsupported("its output is not one vector per face (shape [batch, dimension])".into()); + } + let dim = match output.shape.get(1).and_then(fixed).map(u32::try_from) { + Some(Ok(d)) if (8..=4096).contains(&d) => d, + _ => return Suggestion::Unsupported("its output size is not a fixed number between 8 and 4096".into()), + }; + let manifest = ModelManifest { + id: make_id(file_name, sha256), + name: display_name(file_name), + version: "1".into(), + role: Role::Embedder, + licence: Licence { + name: "Unknown".into(), + commercial: Commercial::Unknown, + url: None, + notice: "LightCraft does not know where this model came from or what its licence allows. Use it only if its licence allows what you are doing, and do not redistribute it.".into(), + }, + source: None, + sha256: Some(sha256.to_string()).filter(|h| h.len() == 64), + size_bytes: Some(size).filter(|s| *s > 0), + provenance: "Unknown.".into(), + input: InputSpec { width, height, ..InputSpec::default() }, + output: OutputSpec::Embedding { dim }, + thresholds: Thresholds::default(), + }; + let assumptions = vec![ + format!("Faces are aligned crops of {width} × {height} pixels."), + "Pixels are RGB, normalised as (x − 127.5) / 127.5: the usual ArcFace convention, which InsightFace models use.".into(), + format!("Each face becomes a vector of {dim} numbers; two faces match when their cosine similarity is high."), + ]; + Suggestion::Draft { manifest, assumptions } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::manifest::validate; + use crate::onnx::probe_reader; + use std::io::Cursor; + + fn tensor(elem: u32, shape: &[Dim]) -> TensorInfo { + TensorInfo { name: "t".into(), elem_type: elem, shape: shape.to_vec() } + } + fn dynamic() -> Dim { + Dim::Dynamic("N".into()) + } + fn info(input: TensorInfo, output: TensorInfo) -> OnnxInfo { + OnnxInfo { producer: "x".into(), opset: Some(17), inputs: vec![input], outputs: vec![output] } + } + const SHA: &str = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"; + + #[test] + fn a_known_hash_wins() { + let i = info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(10), Dim::Fixed(10)]), tensor(1, &[dynamic(), Dim::Fixed(5)])); + assert!(matches!(suggest(&i, crate::known::AURAFACE_SHA256, 1, "whatever.onnx"), Suggestion::Known(m) if m.id == "auraface-v1")); + } + + #[test] + fn an_unknown_embedder_gets_a_draft_that_validates() { + let i = info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(512)])); + let Suggestion::Draft { manifest, assumptions } = suggest(&i, SHA, 5000, "My Model (r50).onnx") else { unreachable_in_test() }; + assert_eq!(validate(&manifest), Ok(())); + assert_eq!(manifest.id, "custom-my-model--r50-01234567"); + assert_eq!(manifest.output, OutputSpec::Embedding { dim: 512 }); + assert_eq!((manifest.input.width, manifest.input.height), (112, 112)); + assert_eq!(manifest.licence.commercial, Commercial::Unknown); + assert!(assumptions.iter().any(|a| a.contains("127.5"))); + } + + #[allow(clippy::panic)] + fn unreachable_in_test() -> ! { + panic!("expected a draft"); + } + + #[test] + fn dynamic_sizes_default_and_odd_shapes_are_refused() { + let ok = info(tensor(1, &[dynamic(), Dim::Fixed(3), dynamic(), dynamic()]), tensor(1, &[dynamic(), Dim::Fixed(128)])); + assert!(matches!(suggest(&ok, SHA, 1, "a.onnx"), Suggestion::Draft { manifest, .. } if manifest.input.width == 112)); + let refused = [ + info(tensor(1, &[dynamic(), Dim::Fixed(1), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), // grey + info(tensor(7, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), // int64 input + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), // rank 3 + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(1_000_000), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(257), Dim::Fixed(384)])), // tokens + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), dynamic()])), + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(2)])), + info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(1_000_000)])), + ]; + for r in refused { + assert!(matches!(suggest(&r, SHA, 1, "a.onnx"), Suggestion::Unsupported(_)), "{r:?}"); + } + let mut two = ok.clone(); + two.outputs.push(two.outputs[0].clone()); + assert!(matches!(suggest(&two, SHA, 1, "a.onnx"), Suggestion::Unsupported(_)), "several outputs"); + } + + #[test] + fn file_names_cannot_break_the_manifest() { + let i = info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(512)])); + for name in ["", ".onnx", "../../etc/passwd", "\u{7}\u{0}x.onnx", &"x".repeat(5000), "ÅÄÖ 🙂.onnx", "UPPER lower.ONNX"] { + let Suggestion::Draft { manifest, .. } = suggest(&i, SHA, 1, name) else { unreachable_in_test() }; + assert_eq!(validate(&manifest), Ok(()), "{name:?}"); + } + // and with no usable hash at all + let Suggestion::Draft { manifest, .. } = suggest(&i, "", 0, "x.onnx") else { unreachable_in_test() }; + assert_eq!(validate(&manifest), Ok(())); + assert!(manifest.sha256.is_none() && manifest.size_bytes.is_none()); + } + + /// Opt-in: `LC_FACE_MODELS=/folder/with/onnx/files cargo test -p lightcraft-faces -- --ignored --nocapture` + /// probes real models (never committed) and prints what would be suggested, with timings. + #[test] + #[ignore = "needs real model files: set LC_FACE_MODELS"] + fn probing_real_models() { + let Some(dir) = std::env::var_os("LC_FACE_MODELS") else { return }; + for entry in std::fs::read_dir(dir).unwrap().flatten() { + let path = entry.path(); + if path.extension().is_none_or(|e| e != "onnx") { + continue; + } + let t = std::time::Instant::now(); + let sha = crate::hash::sha256_file(&path).unwrap(); + let hashed = t.elapsed(); + let t = std::time::Instant::now(); + let info = crate::onnx::probe_path(&path); + let probed = t.elapsed(); + let size = std::fs::metadata(&path).unwrap().len(); + let name = path.file_name().unwrap().to_string_lossy().into_owned(); + let verdict = match &info { + Ok(i) => match suggest(i, &sha, size, &name) { + Suggestion::Known(m) => format!("KNOWN {} ({:?})", m.id, m.licence.commercial), + Suggestion::Draft { manifest, .. } => format!("DRAFT {} {:?}", manifest.id, manifest.output), + Suggestion::Unsupported(why) => format!("UNSUPPORTED: {why}"), + }, + Err(e) => format!("PROBE ERROR: {e}"), + }; + println!("{name}: {} MB | hash {hashed:?} | probe {probed:?} | {verdict}", size / 1_048_576); + if let Ok(i) = info { + println!( + " in {:?} + out {:?}", + i.inputs, i.outputs + ); + } + } + } + + #[test] + fn probing_then_suggesting_end_to_end() { + let probed = probe_reader(&mut Cursor::new(crate::onnx::tests::model_bytes(512))).unwrap(); + assert!( + matches!(suggest(&probed, SHA, 9, "m.onnx"), Suggestion::Draft { manifest, .. } if manifest.output == OutputSpec::Embedding { dim: 512 }) + ); + } +} diff --git a/xtask/src/layers.rs b/xtask/src/layers.rs index 28d52677..44d60473 100644 --- a/xtask/src/layers.rs +++ b/xtask/src/layers.rs @@ -43,6 +43,7 @@ pub const TABLE: &[(&str, Class)] = &[ ("codecs", Class::Layer(1)), ("meta", Class::Layer(1)), ("develop", Class::Layer(1)), + ("faces", Class::Layer(1)), ("scenes", Class::Layer(1)), ("pipeline", Class::Layer(2)), ("gpu", Class::Layer(3)),