From d470ff7b71be799f2bb4dfb8ffaa499a57b42825 Mon Sep 17 00:00:00 2001 From: Jeremy Date: Wed, 21 Oct 2026 23:17:47 +0100 Subject: [PATCH] feat(ai-features): list SAM 2.1 on the built-in rten engine, reading the published ONNX files (#84) * feat(ai-core): share the image.segment request and result contract * fix(ai-core): take only an inline PNG as the selection mask * feat(ai-engine-rten): bounded picture decoding, mask encoding and resampling for selection models * feat(ai-engine-rten): run SAM 2.1 through the shared selection contract * fix(ai-engine-rten): bound a thin picture's memory, refuse extreme sides, and make the mask size fallible * feat(ai-features): let each built-in engine follow its own recipe A built-in engine now plans, offers, registers and checks files from the first usable recipe that lists it (`Feature::engine_recipe`), instead of every engine acting on the first usable recipe. A first recipe that lists one engine no longer hides a later recipe that lists another. The plain built-in offer and engine advice still start from the first usable recipe, and `Feature::engines` lists every engine some usable recipe names, once each, in catalogue order. The catalogue is unchanged, so no behaviour changes today; the tests build the two-recipe shape from the real SAM 3 recipe. * feat(ai-features): list SAM 2.1 on the built-in rten engine, reading the published ONNX files SAM 2.1 now lists a `rten` route beside the ComfyUI one: the two published ONNX files (`cstr/sam2.1-hiera-tiny-ONNX` at a pinned revision, size and SHA-256 pinned), read as they are, with no conversion and nothing hosted by us. A built-in route may now name `.onnx` weights as well as `.safetensors`, in the catalogue and when an app is given a weights path. SAM 2.1 stays second in the list, so nothing changes for a person who has already chosen SAM 3. The model is in the Apache-2.0 licence of Meta's SAM 2.1; the publisher of the conversion declares the same. Real-model checks (ignored unless the files are given): two apps get a mask by asking for Select subject, and a hover through the ordinary job path. * docs(ai-features): the pinned SAM 2.1 export is the tiny size, not the small one A review finding: the doc comment on the ONNX file list said small, but the pinned repository is cstr/sam2.1-hiera-tiny-ONNX and the pinned sizes match the tiny variant. SAM 2.1 also has a separate small variant, so the wrong word could send a maintainer to the wrong file. * fix(ai-engine-rten): judge a match on the mask handed on, not on the raw logits A review finding: matched was set from the model's raw logits before resampling and 8-bit quantization, so one positive logit that a small picture never samples reported a match with an all-black mask. The match is now whether any pixel of the emitted mask is above half, using the same quantization as the PNG. Written first and seen to fail (matched stayed true for a speck at logit 8 in a 10 by 6 picture). The runner's raw-logit test stays as a shortcut that skips the resampling; forcing it open gives the same answer, so its mutation case is retired with the reason in the group header. * docs(mutation): point the usage text at GROUPS instead of a list that was out of date A review finding: the usage line named only the early groups and left out rt, which is new here, and also seg, gs and others. A list in the usage text goes stale each time a group is added, so it now says to take the group from GROUPS and README.md. --- crates/ai-client/src/local.rs | 10 +-- crates/ai-client/tests/engine_advice.rs | 5 +- crates/ai-client/tests/local.rs | 8 +- crates/ai-client/tests/local_native.rs | 48 ++++++++++- crates/ai-client/tests/picker_native.rs | 26 +++--- crates/ai-client/tests/rten_hover_live.rs | 86 ++++++++++++++++++++ crates/ai-client/tests/rten_live.rs | 99 +++++++++++++++++++++++ crates/ai-features/src/catalogue/sam21.rs | 52 +++++++++++- crates/ai-features/src/recipe.rs | 14 ++-- crates/ai-features/tests/native_onnx.rs | 87 ++++++++++++++++++++ tools/mutation/mutation_check.py | 6 +- tools/mutation/mutation_native.py | 3 + 12 files changed, 410 insertions(+), 34 deletions(-) create mode 100644 crates/ai-client/tests/rten_hover_live.rs create mode 100644 crates/ai-client/tests/rten_live.rs create mode 100644 crates/ai-features/tests/native_onnx.rs diff --git a/crates/ai-client/src/local.rs b/crates/ai-client/src/local.rs index d968ad4..f4f6cb6 100644 --- a/crates/ai-client/src/local.rs +++ b/crates/ai-client/src/local.rs @@ -322,18 +322,18 @@ fn local_connection( } } -/// An absolute path to an existing `.safetensors` file, as text. +/// An absolute path to an existing `.safetensors` or `.onnx` file, as text. fn weights_file(path: &Path) -> Result { let text = path .to_str() .filter(|t| !t.contains('\0')) .ok_or_else(|| AiError::invalid("the weights path is not valid text"))?; - let safetensors = path + let known = path .extension() - .is_some_and(|e| e.eq_ignore_ascii_case("safetensors")); - if !path.is_absolute() || !safetensors || !path.is_file() { + .is_some_and(|e| e.eq_ignore_ascii_case("safetensors") || e.eq_ignore_ascii_case("onnx")); + if !path.is_absolute() || !known || !path.is_file() { return Err(AiError::invalid( - "the weights must be an existing .safetensors file, given by its full path", + "the weights must be an existing .safetensors or .onnx file, given by its full path", )); } Ok(text.to_owned()) diff --git a/crates/ai-client/tests/engine_advice.rs b/crates/ai-client/tests/engine_advice.rs index 6e2162f..ddb38b5 100644 --- a/crates/ai-client/tests/engine_advice.rs +++ b/crates/ai-client/tests/engine_advice.rs @@ -52,11 +52,14 @@ impl World { /// A client of an app that has not enabled any built-in engine, whatever cargo features are on. fn without_engines() -> ProviderFactory { let mut factory = ProviderFactory::standard(); + factory.unregister("rten"); factory.unregister("candle"); factory.unregister("burn"); factory } +/// An app with no built-in engine is advised ComfyUI, and every engine some usable recipe lists +/// (rten for SAM 2.1, Candle and Burn for SAM 3) is listed, neither usable nor recommended. #[test] fn an_app_that_only_has_comfyui_is_advised_comfyui_and_the_built_in_ones_are_listed_unusable() { let w = World::new(); @@ -64,7 +67,7 @@ fn an_app_that_only_has_comfyui_is_advised_comfyui_and_the_built_in_ones_are_lis .client() .engine_advice(&select_subject(), &windows_nvidia()); let providers: Vec<&str> = advice.iter().map(|a| a.provider.as_str()).collect(); - assert_eq!(providers, ["comfyui", "candle", "burn"]); + assert_eq!(providers, ["comfyui", "candle", "burn", "rten"]); assert!(advice[0].recommended && advice[0].usable); assert!(advice[1..].iter().all(|a| !a.usable && !a.recommended)); } diff --git a/crates/ai-client/tests/local.rs b/crates/ai-client/tests/local.rs index 92ba7d6..0a9a2f7 100644 --- a/crates/ai-client/tests/local.rs +++ b/crates/ai-client/tests/local.rs @@ -215,9 +215,9 @@ fn the_picker_then_lists_sam_3_as_ready_not_as_something_to_install() { .unwrap(); let options = craft_ai_features::setup_options(&craft_ai_features::builtin()); let view = w.client("photocraft").picker(&role(), &options).unwrap(); - // SAM 3 is ready; SAM 2.1 is not installed, so it keeps its way in; the built-in engine is a - // separate option and this app has no adapter for it. - assert_eq!(view.groups.len(), 3); + // SAM 3 is ready; SAM 2.1 is not installed, so it keeps its way in; the two built-in engines + // (Candle for SAM 3, rten for SAM 2.1) are separate options. + assert_eq!(view.groups.len(), 4); let comfy = |state| { view.groups .iter() @@ -235,7 +235,7 @@ fn the_picker_then_lists_sam_3_as_ready_not_as_something_to_install() { let native = view .groups .iter() - .find(|g| g.title.contains("built in")) + .find(|g| g.title.starts_with("SAM 3, built in")) .unwrap(); if cfg!(feature = "candle") { assert_eq!(native.state, GroupState::NotSetUp); diff --git a/crates/ai-client/tests/local_native.rs b/crates/ai-client/tests/local_native.rs index 9d5c6ea..f2f8ea8 100644 --- a/crates/ai-client/tests/local_native.rs +++ b/crates/ai-client/tests/local_native.rs @@ -52,6 +52,8 @@ impl World { } } +/// Registering the Candle route binds the tool and an app reads the exact weights file back, with +/// no address. #[test] fn the_native_route_is_registered_bound_and_carries_the_file() { let w = World::new(); @@ -73,6 +75,29 @@ fn the_native_route_is_registered_bound_and_carries_the_file() { assert_eq!(route.file.as_deref(), file.to_str()); } +/// SAM 2.1 on rten reads published `.onnx` files: the encoder file is accepted as the weights, +/// registered under the rten connection and bound, and an app reads that exact file back. +#[test] +fn the_native_route_is_registered_for_an_onnx_file() { + let w = World::new(); + let file = w.weights("encoder.onnx"); + let done = w + .desk("launcher") + .enable_native_with(&feature("select-subject"), "rten", &file) + .unwrap(); + assert!(done.bound); + assert_eq!(done.connection.as_str(), "local-select-subject-rten"); + assert_eq!(done.model.as_str(), "sam21"); + let route = w + .client("lightcraft") + .local_route_for(&role()) + .unwrap() + .unwrap(); + assert_eq!(route.provider, "rten"); + assert_eq!(route.file.as_deref(), file.to_str()); +} + +/// Two apps that share the store read the same registered weights file. #[test] fn two_apps_read_the_same_file() { let w = World::new(); @@ -89,6 +114,8 @@ fn two_apps_read_the_same_file() { } } +/// A ComfyUI choice records no weights file, and a built-in runtime lookup on it is refused in +/// plain words. #[test] fn the_comfyui_route_has_no_file_and_is_not_a_native_one() { let w = World::new(); @@ -107,6 +134,8 @@ fn the_comfyui_route_has_no_file_and_is_not_a_native_one() { ); } +/// An app that is bound to a built-in runtime gets the plain "can not run the model yet" refusal +/// when it asks for a ComfyUI address. #[test] fn a_native_route_is_not_offered_as_a_comfyui_address() { let w = World::new(); @@ -122,6 +151,7 @@ fn a_native_route_is_not_offered_as_a_comfyui_address() { ); } +/// An app asking for a built-in runtime file when nothing is bound gets none, not an error. #[test] fn nothing_bound_is_not_an_error() { let w = World::new(); @@ -133,6 +163,8 @@ fn nothing_bound_is_not_an_error() { ); } +/// Enabling the built-in runtime after a ComfyUI choice leaves that choice bound and registers both +/// connections. #[test] fn a_choice_already_made_is_left_alone_and_both_routes_are_registered() { let w = World::new(); @@ -156,6 +188,8 @@ fn a_choice_already_made_is_left_alone_and_both_routes_are_registered() { assert_eq!(store.connections.len(), 2); } +/// Registering Candle again with another file updates the one connection to the new file and +/// leaves the binding as it was. #[test] fn registering_again_moves_the_file_and_keeps_one_connection() { let w = World::new(); @@ -177,6 +211,7 @@ fn registering_again_moves_the_file_and_keeps_one_connection() { ); } +/// A missing file, a folder and a wrong extension are each refused and nothing is registered. #[test] fn a_file_that_cannot_be_the_weights_is_refused_and_nothing_is_written() { let w = World::new(); @@ -197,6 +232,7 @@ fn a_file_that_cannot_be_the_weights_is_refused_and_nothing_is_written() { assert!(w.shared.load().unwrap().connections.is_empty()); } +/// A relative weights path is refused even when a file is there, and nothing is registered. #[test] fn a_relative_path_is_refused_even_when_it_points_at_a_real_file() { let w = World::new(); @@ -211,18 +247,20 @@ fn a_relative_path_is_refused_even_when_it_points_at_a_real_file() { assert!(w.shared.load().unwrap().connections.is_empty()); } +/// With no recipe able to run, or none with a built-in route, the runtime is refused and nothing is +/// written. #[test] fn a_runtime_that_is_not_ready_to_use_is_refused() { let w = World::new(); let mut f = feature("select-subject"); - f.locals[0].usable = false; + f.locals.iter_mut().for_each(|r| r.usable = false); assert!( w.desk("launcher") .enable_native(&f, &w.weights("sam3.safetensors")) .is_err() ); let mut f = feature("select-subject"); - f.locals[0].native.clear(); + f.locals.iter_mut().for_each(|r| r.native.clear()); assert!( w.desk("launcher") .enable_native(&f, &w.weights("sam3.safetensors")) @@ -231,6 +269,7 @@ fn a_runtime_that_is_not_ready_to_use_is_refused() { assert!(w.shared.load().unwrap().connections.is_empty()); } +/// Every other tool refuses a built-in runtime and writes nothing. #[test] fn a_tool_without_a_built_in_runtime_is_refused() { let w = World::new(); @@ -248,6 +287,7 @@ fn a_tool_without_a_built_in_runtime_is_refused() { assert!(w.shared.load().unwrap().connections.is_empty()); } +/// A Candle weights file that has gone is an error naming that, not an unbound tool. #[test] fn a_file_that_vanished_is_an_error_not_unbound() { let w = World::new(); @@ -267,6 +307,8 @@ fn a_file_that_vanished_is_an_error_not_unbound() { ); } +/// With the ComfyUI model and the built-in runtime both registered, the picker lists two groups +/// with different titles, one of them saying "built in". #[test] fn the_picker_offers_the_two_runtimes_of_one_model_side_by_side() { let w = World::new(); @@ -286,6 +328,8 @@ fn the_picker_offers_the_two_runtimes_of_one_model_side_by_side() { assert!(titles.iter().any(|t| t.contains("built in")), "{titles:?}"); } +/// A model saved before `local_file` existed loads with none, and the field is not written out when +/// it is absent. #[test] fn old_stores_without_the_field_load_and_it_is_not_written_when_absent() { let json = r#"{"id":"sam3","label":"SAM 3"}"#; diff --git a/crates/ai-client/tests/picker_native.rs b/crates/ai-client/tests/picker_native.rs index aef3512..2e63fa7 100644 --- a/crates/ai-client/tests/picker_native.rs +++ b/crates/ai-client/tests/picker_native.rs @@ -27,7 +27,7 @@ fn the_builtin_model_only_offers_setup_when_this_app_has_the_engine() { let native = view .groups .iter() - .find(|g| g.title.contains("built in")) + .find(|g| g.title.starts_with("SAM 3, built in")) .unwrap(); if cfg!(feature = "candle") { assert_eq!(native.state, GroupState::NotSetUp); @@ -68,7 +68,7 @@ fn a_built_in_model_whose_file_is_gone_can_be_set_up_again() { Arc::new(MemoryStore::new()), ); ToolDesk::new(launcher, shared, None) - .enable_native(feature, &weights) + .enable_native_with(feature, "candle", &weights) .unwrap(); let options = craft_ai_features::setup_options(&features); let native = |client: &AiClient| { @@ -77,7 +77,7 @@ fn a_built_in_model_whose_file_is_gone_can_be_set_up_again() { .unwrap() .groups .into_iter() - .filter(|g| g.title.contains("built in") && g.state == GroupState::NotSetUp) + .filter(|g| g.title.starts_with("SAM 3, built in") && g.state == GroupState::NotSetUp) .collect::>() }; assert!( @@ -120,7 +120,7 @@ fn a_built_in_connection_for_another_model_does_not_hide_this_one() { Arc::new(MemoryStore::new()), ); ToolDesk::new(launcher, shared.clone(), None) - .enable_native(feature, &weights) + .enable_native_with(feature, "candle", &weights) .unwrap(); // The file is there, but the connection now belongs to some other model. let mut store = shared.load().unwrap(); @@ -129,9 +129,11 @@ fn a_built_in_connection_for_another_model_does_not_hide_this_one() { let options = craft_ai_features::setup_options(&features); let view = client.picker(&role, &options).unwrap(); assert!( - view.groups.iter().any(|g| g.title.contains("built in") - && g.state == GroupState::NotSetUp - && matches!(&g.setup, Some(SetupLink::New { option }) + view.groups + .iter() + .any(|g| g.title.starts_with("SAM 3, built in") + && g.state == GroupState::NotSetUp + && matches!(&g.setup, Some(SetupLink::New { option }) if option == "native:select-subject:candle")), "this model still has a way in" ); @@ -165,7 +167,7 @@ fn an_older_model_on_the_same_connection_does_not_hide_the_current_one() { Arc::new(MemoryStore::new()), ); ToolDesk::new(launcher, shared.clone(), None) - .enable_native(feature, &weights) + .enable_native_with(feature, "candle", &weights) .unwrap(); let mut store = shared.load().unwrap(); store.connections[0].models[0].id = ModelId::new("sam3-older").unwrap(); @@ -173,9 +175,11 @@ fn an_older_model_on_the_same_connection_does_not_hide_the_current_one() { let options = craft_ai_features::setup_options(&features); let view = client.picker(&role, &options).unwrap(); assert!( - view.groups.iter().any(|g| g.title.contains("built in") - && g.state == GroupState::NotSetUp - && g.setup.is_some()), + view.groups + .iter() + .any(|g| g.title.starts_with("SAM 3, built in") + && g.state == GroupState::NotSetUp + && g.setup.is_some()), "the current model still has a way in" ); } diff --git a/crates/ai-client/tests/rten_hover_live.rs b/crates/ai-client/tests/rten_hover_live.rs new file mode 100644 index 0000000..5dfce45 --- /dev/null +++ b/crates/ai-client/tests/rten_hover_live.rs @@ -0,0 +1,86 @@ +//! How fast hover can be through the ordinary job path: one picture, then a mask for each mouse +//! position, each a normal `plan`/`start`/`wait` round trip with the real SAM 2.1 files on rten. +//! It measures, and fails only if a hover is slower than a person would accept. +//! +//! Skips unless `CRAFT_TEST_SAM21_ONNX_DIR` and `CRAFT_TEST_SAM21_IMAGE` are set (see `rten_live`). +//! Run: `cargo test -p craft-ai-client --features rten --release --test rten_hover_live -- --nocapture`. +#![cfg(feature = "rten")] + +use std::path::PathBuf; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use craft_ai_client::{AiClient, ToolDesk}; +use craft_ai_core::{AppId, ConsentLedger, InputRole, JobState, Payload, Request, StoreFile}; +use craft_ai_native::secrets::MemoryStore; +use craft_ai_native::shared::SharedStore; +use serde_json::json; + +/// The encoder weights and the sample picture, or `None` when the environment does not name them. +fn files() -> Option<(PathBuf, Vec)> { + let dir = PathBuf::from(std::env::var_os("CRAFT_TEST_SAM21_ONNX_DIR")?); + let image = std::fs::read(std::env::var_os("CRAFT_TEST_SAM21_IMAGE")?).ok()?; + Some((dir.join("encoder.onnx"), image)) +} + +/// The time below which a hover feels live; past it the selection visibly trails the mouse. +const HOVER_BUDGET: Duration = Duration::from_millis(150); + +/// After the first look at the picture, every hover position through the ordinary job path +/// returns its mask within [`HOVER_BUDGET`] at the median, and the timings are printed. +#[test] +fn hover_through_the_job_path_is_fast_enough_after_the_first_look() { + let Some((weights, image)) = files() else { + eprintln!("skipped: set CRAFT_TEST_SAM21_ONNX_DIR and CRAFT_TEST_SAM21_IMAGE"); + return; + }; + let dir = tempfile::tempdir().unwrap(); + let shared = SharedStore::new(dir.path().join("connections.json")); + shared.replace(&StoreFile::default()).unwrap(); + let client = |app: &str| { + AiClient::new( + AppId::new(app).unwrap(), + shared.clone(), + Arc::new(MemoryStore::new()), + ) + }; + let feature = craft_ai_features::builtin() + .into_iter() + .find(|f| f.id == "select-subject") + .unwrap(); + let role = feature.role().unwrap(); + ToolDesk::new(client("launcher"), shared.clone(), None) + .enable_native_with(&feature, "rten", &weights) + .unwrap(); + let app = client("photocraft"); + let source = Payload::inline("image/png", image).unwrap(); + let ask = |x: f64, y: f64| { + let request = Request::new(feature.capability.clone()) + .with_input(InputRole::Source, source.clone()) + .with_param("points", json!([{"x": x, "y": y, "positive": true}])); + let started = Instant::now(); + let plan = app.plan(&role, request).unwrap(); + let done = app + .start(&plan, &ConsentLedger::default()) + .unwrap() + .wait(Duration::from_millis(1), Duration::from_secs(120)) + .unwrap(); + assert_eq!(done.state, JobState::Succeeded); + started.elapsed() + }; + let first = ask(0.29, 0.64); + let mut hovers: Vec = (0..30) + .map(|i| { + let t = f64::from(i) / 30.0; + ask(0.2 + 0.6 * t, 0.3 + 0.5 * t) + }) + .collect(); + hovers.sort(); + let median = hovers[hovers.len() / 2]; + let slowest = hovers[hovers.len() - 1]; + println!( + "first look {first:?}; hover median {median:?}, fastest {:?}, slowest {slowest:?}", + hovers[0] + ); + assert!(median <= HOVER_BUDGET, "hover median {median:?}"); +} diff --git a/crates/ai-client/tests/rten_live.rs b/crates/ai-client/tests/rten_live.rs new file mode 100644 index 0000000..480cc5e --- /dev/null +++ b/crates/ai-client/tests/rten_live.rs @@ -0,0 +1,99 @@ +//! The whole chain with the real model: the catalogue's route, the shared store, the cargo feature, +//! the rten engine and the publisher's ONNX files. An app writes none of it, and two apps read +//! the same choice. +//! +//! Skips unless `CRAFT_TEST_SAM21_ONNX_DIR` names a folder holding `encoder.onnx` and +//! `decoder.onnx`, and `CRAFT_TEST_SAM21_IMAGE` names a picture. Run: +//! `cargo test -p craft-ai-client --features rten --release --test rten_live`. +#![cfg(feature = "rten")] + +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; + +use craft_ai_client::{AiClient, ToolDesk}; +use craft_ai_core::{ + AppId, ConsentLedger, Event, InputRole, JobState, Payload, Request, StoreFile, segment, +}; +use craft_ai_native::secrets::MemoryStore; +use craft_ai_native::shared::SharedStore; +use serde_json::json; + +/// The encoder path and the sample picture named in the environment, or `None` when they are not +/// set. +fn files() -> Option<(PathBuf, Vec)> { + let dir = PathBuf::from(std::env::var_os("CRAFT_TEST_SAM21_ONNX_DIR")?); + let image = std::fs::read(std::env::var_os("CRAFT_TEST_SAM21_IMAGE")?).ok()?; + Some((dir.join("encoder.onnx"), image)) +} + +/// Two apps ask only for the capability and each gets the real model's mask of the suit, at the +/// picture's size, with the shared result saying it matched. +#[test] +fn two_apps_get_a_mask_from_the_real_model_by_asking_for_select_subject() { + let Some((weights, image)) = files() else { + eprintln!("skipped: set CRAFT_TEST_SAM21_ONNX_DIR and CRAFT_TEST_SAM21_IMAGE"); + return; + }; + let dir = tempfile::tempdir().unwrap(); + let shared = SharedStore::new(dir.path().join("connections.json")); + shared.replace(&StoreFile::default()).unwrap(); + let client = |app: &str| { + AiClient::new( + AppId::new(app).unwrap(), + shared.clone(), + Arc::new(MemoryStore::new()), + ) + }; + let feature = craft_ai_features::builtin() + .into_iter() + .find(|f| f.id == "select-subject") + .unwrap(); + let role = feature.role().unwrap(); + + // The launcher's part, once the files are in place: rten's recipe is SAM 2.1. + ToolDesk::new(client("launcher"), shared.clone(), None) + .enable_native_with(&feature, "rten", &weights) + .unwrap(); + + // Each app's part: ask for the capability. Neither is told which engine, model or file. + for app in ["lightcraft", "photocraft"] { + let app = client(app); + let request = Request::new(feature.capability.clone()) + .with_input( + InputRole::Source, + Payload::inline("image/png", image.clone()).unwrap(), + ) + .with_param("points", json!([{"x": 0.29, "y": 0.64, "positive": true}])); + let plan = app.plan(&role, request).unwrap(); + let done = app + .start(&plan, &ConsentLedger::default()) + .unwrap() + .wait(Duration::from_millis(50), Duration::from_secs(120)) + .unwrap(); + assert_eq!(done.state, JobState::Succeeded); + let outputs: Vec<_> = done + .events + .iter() + .filter_map(|e| match e { + Event::Output(o) => Some(o.clone()), + _ => None, + }) + .collect(); + // The mask, then the shared result saying it matched. + assert_eq!(outputs.len(), 2); + assert!(segment::read_found(&outputs).unwrap().matched); + let Some(Payload::Inline { media_type, data }) = segment::mask(&outputs) else { + panic!("expected the mask itself"); + }; + assert_eq!(media_type.as_str(), "image/png"); + let mut reader = png::Decoder::new(std::io::Cursor::new(data)) + .read_info() + .unwrap(); + let mut pixels = vec![0; reader.output_buffer_size().unwrap()]; + let frame = reader.next_frame(&mut pixels).unwrap(); + assert_eq!((frame.width, frame.height), (512, 512)); + let selected = pixels.iter().filter(|v| **v > 127).count() as f64 / (512.0 * 512.0); + assert!((0.15..0.35).contains(&selected), "selected {selected}"); + } +} diff --git a/crates/ai-features/src/catalogue/sam21.rs b/crates/ai-features/src/catalogue/sam21.rs index d3bd953..f017abb 100644 --- a/crates/ai-features/src/catalogue/sam21.rs +++ b/crates/ai-features/src/catalogue/sam21.rs @@ -1,7 +1,48 @@ //! Meta's SAM 2.1, run through ComfyUI with Florence-2-large finding the words of a prompt. use super::florence2::florence2_large_files; -use crate::recipe::{Component, Licence, LocalRecipe, ModelFile, NodePack, Source}; +use crate::recipe::{Component, Licence, LocalRecipe, ModelFile, NativeRoute, NodePack, Source}; + +/// The published ONNX export of SAM 2.1 (tiny size) at a pinned revision. The built-in rten +/// engine reads these two files as they are: no conversion, nothing hosted by us. Both sit in +/// one folder, `encoder.onnx` first (the weights file the connection records). +fn sam21_onnx_files() -> Vec { + const REVISION: &str = "12bdf0f712731bd0b52ccbff807683ce49b726b9"; + let file = |name: &str, size_bytes: u64, sha256: &str| { + ModelFile { + folder: "sam21_onnx".to_owned(), + filename: name.to_owned(), + size_bytes, + sha256: sha256.to_owned(), + source: Source::HuggingFace { + repo: "cstr/sam2.1-hiera-tiny-ONNX".to_owned(), + file: name.to_owned(), + revision: Some(REVISION.to_owned()), + }, + licence: Licence { + name: "Apache-2.0, ONNX conversion".to_owned(), + url: Some("https://huggingface.co/cstr/sam2.1-hiera-tiny-ONNX".to_owned()), + accept_required: false, + note: Some( + "This file is a format conversion of Meta's SAM 2.1 weights, which Meta releases under Apache 2.0. The publisher of the conversion declares the same licence." + .to_owned(), + ), + }, + } + }; + vec![ + file( + "encoder.onnx", + 109_473_919, + "4fda6e68561e27808cedffa8689749a622bba5f11ba66f4ca4e3a8a9f6360db5", + ), + file( + "decoder.onnx", + 20_655_795, + "0ae944e08e55814c19afa4d143ef07b1a3dbad2decd77c15ec3ab220b9f57247", + ), + ] +} /// SAM 2.1 for Select subject: the large checkpoint in the layout the /// `ComfyUI-segment-anything-2` node pack reads, at a pinned revision of a community conversion, @@ -48,7 +89,14 @@ pub(super) fn sam21() -> LocalRecipe { }, ], min_vram_mb: None, - native: Vec::new(), + // Built-in engines that run SAM 2.1. It selects by clicking; a description needs the + // ComfyUI route above (Florence-2 finds the words). + native: vec![NativeRoute { + provider: "rten".to_owned(), + model: "sam21".to_owned(), + label: "SAM 2.1, built in (the processor, on any system)".to_owned(), + models: sam21_onnx_files(), + }], components: components(), } } diff --git a/crates/ai-features/src/recipe.rs b/crates/ai-features/src/recipe.rs index 0904d17..05e7c1f 100644 --- a/crates/ai-features/src/recipe.rs +++ b/crates/ai-features/src/recipe.rs @@ -103,18 +103,20 @@ pub struct NativeRoute { pub label: String, /// The files this runtime reads, when they are not the recipe's own. A runtime ported from /// another framework usually wants that framework's file layout, which is a different file - /// from the ComfyUI one. The first `.safetensors` file here is the weights file recorded on - /// the connection; the rest sit beside it. Empty means the recipe's own files. + /// from the ComfyUI one. The first `.safetensors` or `.onnx` file here is the weights file + /// recorded on the connection; the rest sit beside it. Empty means the recipe's own files. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub models: Vec, } impl NativeRoute { - /// The weights file this runtime is given: the first `.safetensors` file of its own files. + /// The weights file this runtime is given: the first `.safetensors` or `.onnx` file of its + /// own files. pub fn weights(&self) -> Option<&ModelFile> { - self.models - .iter() - .find(|m| m.filename.to_ascii_lowercase().ends_with(".safetensors")) + self.models.iter().find(|m| { + let name = m.filename.to_ascii_lowercase(); + name.ends_with(".safetensors") || name.ends_with(".onnx") + }) } /// Bytes to download for this runtime's own files. diff --git a/crates/ai-features/tests/native_onnx.rs b/crates/ai-features/tests/native_onnx.rs new file mode 100644 index 0000000..eac8c5b --- /dev/null +++ b/crates/ai-features/tests/native_onnx.rs @@ -0,0 +1,87 @@ +//! A built-in engine that reads ONNX files (rten) is described by the catalogue the same way one +//! that reads safetensors files is: the first weights file is the one the connection records. + +use craft_ai_features::{Licence, ModelFile, NativeRoute, Source}; + +/// A one-byte catalogue file called `name` from a made-up source. +fn file(name: &str) -> ModelFile { + ModelFile { + folder: "x".to_owned(), + filename: name.to_owned(), + size_bytes: 1, + sha256: "0".repeat(64), + source: Source::HuggingFace { + repo: "a/b".to_owned(), + file: name.to_owned(), + revision: None, + }, + licence: Licence { + name: "test".to_owned(), + url: None, + accept_required: false, + note: None, + }, + } +} + +/// A route for the `rten` engine listing `files` in the order given. +fn route(files: &[&str]) -> NativeRoute { + NativeRoute { + provider: "rten".to_owned(), + model: "m".to_owned(), + label: "m".to_owned(), + models: files.iter().map(|f| file(f)).collect(), + } +} + +/// An `.onnx` file counts as weights, in any case, and the first weights file in the list is the +/// one recorded. +#[test] +fn an_onnx_file_is_a_weights_file_and_the_first_one_wins() { + let r = route(&[ + "README.md", + "Encoder.ONNX", + "model.safetensors", + "decoder.onnx", + ]); + assert_eq!(r.weights().unwrap().filename, "Encoder.ONNX"); + let r = route(&["model.safetensors", "encoder.onnx"]); + assert_eq!(r.weights().unwrap().filename, "model.safetensors"); +} + +/// Tokenizer files and an empty list give no weights file. +#[test] +fn files_that_are_not_weights_give_none() { + assert!(route(&["vocab.json", "merges.txt"]).weights().is_none()); + assert!(route(&[]).weights().is_none()); +} + +/// The SAM 2.1 recipe gives rten the encoder and decoder, each pinned by revision, size and +/// SHA-256, with no licence gate. +#[test] +fn the_sam21_recipe_offers_the_publishers_onnx_files_to_the_rten_engine() { + let feature = craft_ai_features::builtin() + .into_iter() + .find(|f| f.id == "select-subject") + .unwrap(); + let recipe = feature + .locals + .iter() + .find(|l| l.workflow == "sam21-segment") + .unwrap(); + let route = recipe.native.iter().find(|n| n.provider == "rten").unwrap(); + assert_eq!(route.model, "sam21"); + assert_eq!(route.weights().unwrap().filename, "encoder.onnx"); + let names: Vec<_> = route.models.iter().map(|m| m.filename.as_str()).collect(); + assert_eq!(names, ["encoder.onnx", "decoder.onnx"]); + // Pinned and checked: a revision, a size and a SHA-256 for each file, and no licence gate. + for m in &route.models { + assert!( + matches!(&m.source, Source::HuggingFace { revision: Some(r), .. } if r.len() == 40) + ); + assert_eq!(m.sha256.len(), 64); + assert!(m.size_bytes > 1_000_000); + assert!(!m.licence.accept_required); + } + assert_eq!(route.total_bytes(), 109_473_919 + 20_655_795); +} diff --git a/tools/mutation/mutation_check.py b/tools/mutation/mutation_check.py index 0611a1f..1cf15be 100644 --- a/tools/mutation/mutation_check.py +++ b/tools/mutation/mutation_check.py @@ -444,9 +444,9 @@ LR = [ # nr: the built-in runtime route reads one weights file (ai-client/tests/local_native.rs) NR = [ - ('relative-path-accepted', LOCAL_RS, '!path.is_absolute() || !safetensors || !path.is_file()', '!safetensors || !path.is_file()', CL, 'local_native', 'relative_path'), - ('any-extension-accepted', LOCAL_RS, '!path.is_absolute() || !safetensors || !path.is_file()', '!path.is_absolute() || !path.is_file()', CL, 'local_native', 'cannot_be_the_weights'), - ('missing-file-accepted', LOCAL_RS, '!path.is_absolute() || !safetensors || !path.is_file()', '!path.is_absolute() || !safetensors', CL, 'local_native', 'cannot_be_the_weights'), + ('relative-path-accepted', LOCAL_RS, '!path.is_absolute() || !known || !path.is_file()', '!known || !path.is_file()', CL, 'local_native', 'relative_path'), + ('any-extension-accepted', LOCAL_RS, '!path.is_absolute() || !known || !path.is_file()', '!path.is_absolute() || !path.is_file()', CL, 'local_native', 'cannot_be_the_weights'), + ('missing-file-accepted', LOCAL_RS, '!path.is_absolute() || !known || !path.is_file()', '!path.is_absolute() || !known', CL, 'local_native', 'cannot_be_the_weights'), ('unusable-recipe-accepted', LOCAL_RS, '.engine_recipe(provider)\n .and_then(|r| r.native.iter().find(|n| n.provider == provider))', '.locals\n .first()\n .and_then(|r| r.native.iter().find(|n| n.provider == provider))', CL, 'local_recipes', 'built_in_runtime_comes_from'), ('file-not-given-to-apps', LOCAL_RS, 'file: resolved.model.local_file.clone(),', 'file: None,', CL, 'local_native', 'registered_bound'), ('file-not-stored', LOCAL_RS, 'model.local_file = Some(file);', 'model.local_file = None;', CL, 'local_native', 'registered_bound'), diff --git a/tools/mutation/mutation_native.py b/tools/mutation/mutation_native.py index 4908a61..78059f5 100644 --- a/tools/mutation/mutation_native.py +++ b/tools/mutation/mutation_native.py @@ -24,4 +24,7 @@ NF = [ ('native-options-before-later-locals', F + 'options.rs', 'locals.extend(natives);', 'locals.splice(0..0, natives);', FE, 'native_setup', 'follow_every_local'), ('built-in-offer-uses-the-comfyui-terms', F + 'setup.rs', 'agreement: agreement_for(route),', 'agreement: agreement(),', FE, 'setup', 'agrees_to_its_own'), ('native-option-loses-its-model-id', F + 'options.rs', 'model: native.model.clone(),', 'model: String::new(),', FE, 'native_setup', 'candle_with_its_own_size'), + ('onnx-is-not-weights', F + 'recipe.rs', ' name.ends_with(".safetensors") || name.ends_with(".onnx")', ' name.ends_with(".safetensors")', FE, 'native_onnx', 'an_onnx_file'), + ('onnx-name-is-case-sensitive', F + 'recipe.rs', 'let name = m.filename.to_ascii_lowercase();', 'let name = m.filename.clone();', FE, 'native_onnx', 'an_onnx_file'), + ('onnx-path-refused-when-binding', 'crates/ai-client/src/local.rs', '.is_some_and(|e| e.eq_ignore_ascii_case("safetensors") || e.eq_ignore_ascii_case("onnx"));', '.is_some_and(|e| e.eq_ignore_ascii_case("safetensors"));', 'craft-ai-client', 'local_native', 'the_native_route_is_registered'), ]