From 92ac189a49f10d8981894259934f61111f6094b4 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:44:45 +0200 Subject: [PATCH] S2.31: load SAM 2 through the install flow; SAM 2 as a MaskModel Sam2Mask wraps the image path behind the Roto Brush's MaskModel interface (cached frame encoding, bilinear mask resize). The engine loads the verified weights natively (not on wasm). Checked with the real file: install, notice, load, use by Roto Brush 3.0, deselect. Co-Authored-By: Claude Sonnet 5.5 --- Cargo.lock | 2 + crates/engine/Cargo.toml | 3 +- crates/engine/src/analysis/models.rs | 5 + crates/engine/src/tests/models.rs | 17 ++++ crates/sam2/Cargo.toml | 1 + crates/sam2/src/lib.rs | 1 + crates/sam2/src/mask.rs | 141 +++++++++++++++++++++++++++ crates/sam2/src/weights.rs | 6 ++ crates/sam2/tests/reference.rs | 36 +++++++ 9 files changed, 211 insertions(+), 1 deletion(-) create mode 100644 crates/sam2/src/mask.rs diff --git a/Cargo.lock b/Cargo.lock index afb66fd53..7d951deb7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1508,6 +1508,7 @@ dependencies = [ "effectcraft-psd", "effectcraft-raster", "effectcraft-render", + "effectcraft-sam2", "effectcraft-segment", "effectcraft-svg", "effectcraft-text", @@ -1832,6 +1833,7 @@ version = "0.18.0" dependencies = [ "candle-core", "candle-nn", + "effectcraft-segment", ] [[package]] diff --git a/crates/engine/Cargo.toml b/crates/engine/Cargo.toml index 4bb488569..dc3fe923c 100644 --- a/crates/engine/Cargo.toml +++ b/crates/engine/Cargo.toml @@ -40,6 +40,7 @@ image = { workspace = true } web-time = { workspace = true } exr = { workspace = true } -# Native only: the volume's own identity for a file, to tell that a hard link is the footage it leads to. +# Native file identity and SAM 2 inference; browsers keep the built-in engines. [target.'cfg(not(target_arch = "wasm32"))'.dependencies] same-file = { workspace = true } +effectcraft-sam2 = { workspace = true } diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index b4ddab00d..83d0d3c51 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -202,6 +202,11 @@ impl Session { fn load_file(path: &Path, id: &str) -> std::result::Result { let bytes = std::fs::read(path).map_err(|e| format!("cannot read {}: {e}", path.display()))?; + // SAM 2 runs on candle and is native only; every other model loads through the registry. + #[cfg(not(target_arch = "wasm32"))] + if id == seg::SAM2_TINY.id { + return effectcraft_sam2::mask::load(&bytes).map(|m| Loaded::Mask(m)); + } seg::load(id, &bytes) } diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index 786770ec5..0b714d717 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -71,6 +71,23 @@ fn models_are_listed_chosen_and_verified() { s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); assert!(crate::effects::roto::model_for(¶ms).is_none()); } + // EFFECTCRAFT_SAM2=path/to/model.safetensors (facebook/sam2.1-hiera-tiny): installed with its + // notice, loaded (verified, native only) and used by Roto Brush 3.0, then dropped. + #[cfg(not(target_arch = "wasm32"))] + if let Ok(path) = std::env::var("EFFECTCRAFT_SAM2") { + let r = s.execute_checked("roto.model.install", json!({"path": path})).unwrap(); + assert_eq!(r["id"], "sam2.1-hiera-tiny"); + let dir = s.models_dir().unwrap(); + let notice = std::fs::read_to_string(dir.join("sam2.1-hiera-tiny.safetensors.NOTICE.txt")).unwrap(); + assert!(notice.contains("Meta AI") && notice.contains("Apache-2.0"), "{notice}"); + s.execute_checked("roto.model.select", json!({"id": "sam2.1-hiera-tiny"})).unwrap(); + wait_loaded(&mut s, Task::Mask); + assert_eq!(s.models.status(Task::Mask).active.as_deref(), Some("sam2.1-hiera-tiny")); + let params = crate::effects::Params { values: [("version".to_string(), effectcraft_keyframe::Value::Enum(2))].into_iter().collect() }; + assert_eq!(crate::effects::roto::model_for(¶ms).map(|m| m.info().id), Some("sam2.1-hiera-tiny")); + s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); + assert!(crate::effects::roto::model_for(¶ms).is_none()); + } // EFFECTCRAFT_FACE_LANDMARKER=path/to/face_landmarker.task: installed with its notice, // loaded for face tracking, dropped when the classic tracker is chosen, removed. if let Ok(path) = std::env::var("EFFECTCRAFT_FACE_LANDMARKER") { diff --git a/crates/sam2/Cargo.toml b/crates/sam2/Cargo.toml index f32a4e4e6..4b6dd7e9f 100644 --- a/crates/sam2/Cargo.toml +++ b/crates/sam2/Cargo.toml @@ -10,6 +10,7 @@ rust-version.workspace = true workspace = true [dependencies] +effectcraft-segment = { workspace = true } # Native only (D1): candle is not built for wasm32, where this crate is empty. [target.'cfg(not(target_arch = "wasm32"))'.dependencies] diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 4cb6f6f3e..0d7700a0e 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -13,6 +13,7 @@ use std::fmt; pub mod decoder; pub mod hiera; +pub mod mask; pub mod memattn; pub mod memory; pub mod model; diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs new file mode 100644 index 000000000..8bb309da0 --- /dev/null +++ b/crates/sam2/src/mask.rs @@ -0,0 +1,141 @@ +// Modified work: wraps the ported SAM 2 image path behind EffectCraft's `MaskModel` interface +// (Apache-2.0, see NOTICE). +//! SAM 2 as a drop-in [`MaskModel`]: a frame and a prompt in, a foreground probability per pixel +//! out. The frame's encoding is cached, so further strokes on the same frame cost one decode. + +use std::sync::{Arc, Mutex, PoisonError}; + +use candle_core::{Device, Tensor}; +use effectcraft_segment::{MaskModel, ModelInfo, SAM2_TINY}; + +use crate::model::Encoded; +use crate::prompt::{INPUT_SIZE, Prompt}; +use crate::session::VideoModel; +use crate::{Error, Result}; + +/// Largest frame accepted (pixels per side and in total). +const MAX_SIDE: usize = 16_384; +const MAX_PIXELS: usize = 100_000_000; + +/// Resize `w x h` RGB (0..1) to the model input and normalise it: `[1, 3, 1024, 1024]`. +pub fn preprocess(rgb: &[[f32; 3]], w: usize, h: usize) -> Result { + let pixels = w.checked_mul(h).filter(|n| *n > 0 && *n == rgb.len() && w <= MAX_SIDE && h <= MAX_SIDE && *n <= MAX_PIXELS); + let Some(pixels) = pixels else { + return Err(Error::Input(format!("a {w} x {h} frame needs {} pixels (at most {MAX_PIXELS}), got {}", w.saturating_mul(h), rgb.len()))); + }; + const MEAN: [f32; 3] = [0.485, 0.456, 0.406]; + const STD: [f32; 3] = [0.229, 0.224, 0.225]; + let mut planar = vec![0f32; 3 * pixels]; + for (i, px) in rgb.iter().enumerate() { + for c in 0..3 { + let v = px.get(c).copied().unwrap_or(0.0); + if let Some(o) = planar.get_mut(c * pixels + i) { + *o = if v.is_finite() { v.clamp(0.0, 1.0) } else { 0.0 }; + } + } + } + let x = Tensor::from_vec(planar, (1, 3, h, w), &Device::Cpu)?.upsample_bilinear2d(INPUT_SIZE, INPUT_SIZE, false)?; + let mean = Tensor::from_slice(&MEAN, (1, 3, 1, 1), &Device::Cpu)?; + let std = Tensor::from_slice(&STD, (1, 3, 1, 1), &Device::Cpu)?; + Ok(x.broadcast_sub(&mean)?.broadcast_div(&std)?) +} + +fn fingerprint(rgb: &[[f32; 3]], w: usize, h: usize) -> u64 { + let mut hash = 0xcbf2_9ce4_8422_2325u64 ^ ((w as u64) << 32 | h as u64); + for px in rgb { + for v in px { + hash = (hash ^ u64::from(v.to_bits())).wrapping_mul(0x0100_0000_01b3); + } + } + hash +} + +/// SAM 2 for single frames. +pub struct Sam2Mask { + model: Arc, + last: Mutex)>>, +} + +impl Sam2Mask { + pub fn new(model: Arc) -> Sam2Mask { + Sam2Mask { model, last: Mutex::new(None) } + } + + /// The video model, for trackers built on the same weights. + pub fn video_model(&self) -> &Arc { + &self.model + } + + fn encoded(&self, rgb: &[[f32; 3]], w: usize, h: usize) -> Result> { + let key = fingerprint(rgb, w, h); + if let Some((k, e)) = self.last.lock().unwrap_or_else(PoisonError::into_inner).as_ref() + && *k == key + { + return Ok(e.clone()); + } + let enc = Arc::new(self.model.image.encode(&preprocess(rgb, w, h)?)?); + *self.last.lock().unwrap_or_else(PoisonError::into_inner) = Some((key, enc.clone())); + Ok(enc) + } + + fn probability(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &effectcraft_segment::Prompt) -> Result> { + let enc = self.encoded(rgb, w, h)?; + let (sx, sy) = (INPUT_SIZE as f32 / w as f32, INPUT_SIZE as f32 / h as f32); + let p = Prompt { + points: prompt.points.iter().map(|(c, fg)| ([c[0] * sx, c[1] * sy], *fg)).collect(), + bbox: prompt.bbox.map(|b| [b[0] * sx, b[1] * sy, b[2] * sx, b[3] * sy]), + }; + if p.points.is_empty() && p.bbox.is_none() { + return Err(Error::Input("SAM 2 needs at least one point or a box".into())); + } + let multimask = p.points.len() + 2 * usize::from(p.bbox.is_some()) <= 1; + let out = self.model.image.segment(&enc, &p, multimask)?; + let ious = out.iou.flatten_all()?.to_vec1::()?; + let best = ious.iter().enumerate().fold((0usize, f32::NEG_INFINITY), |b, (i, v)| if *v > b.1 { (i, *v) } else { b }).0; + let logits = out.masks.narrow(1, best, 1)?.upsample_bilinear2d(h, w, false)?; + Ok(candle_nn::ops::sigmoid(&logits)?.flatten_all()?.to_vec1::()?) + } +} + +impl MaskModel for Sam2Mask { + fn info(&self) -> &'static ModelInfo { + &SAM2_TINY + } + + fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &effectcraft_segment::Prompt) -> effectcraft_segment::Result> { + self.probability(rgb, w, h, prompt).map_err(|e| e.to_string()) + } +} + +/// Verify the installed file against the registry and build the model. +pub fn load(bytes: &[u8]) -> std::result::Result, String> { + let got = effectcraft_segment::sha256::hex(bytes); + if got != SAM2_TINY.sha256 { + return Err(format!("{}: the file does not match (SHA-256 {got}, expected {})", SAM2_TINY.name, SAM2_TINY.sha256)); + } + let w = crate::weights::Weights::from_bytes(bytes).map_err(|e| e.to_string())?; + let model = VideoModel::load(&w).map_err(|e| e.to_string())?; + Ok(Arc::new(Sam2Mask::new(Arc::new(model)))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn preprocess_rejects_bad_frames_and_cleans_pixels() { + assert!(preprocess(&[], 0, 0).is_err()); + assert!(preprocess(&[[0.0; 3]; 3], 2, 2).is_err()); + assert!(preprocess(&[[0.5; 3]; 4], usize::MAX, 2).is_err()); + let t = preprocess(&[[f32::NAN, 2.0, -1.0]; 4], 2, 2).unwrap(); + assert_eq!(t.dims(), [1, 3, INPUT_SIZE, INPUT_SIZE]); + let v = t.flatten_all().unwrap().to_vec1::().unwrap(); + assert!(v.iter().all(|x| x.is_finite())); + } + + #[test] + fn a_wrong_file_is_refused_before_parsing() { + let e = load(b"not the weights").err().unwrap_or_default(); + assert!(e.contains("does not match"), "{e}"); + } +} diff --git a/crates/sam2/src/weights.rs b/crates/sam2/src/weights.rs index c8b79cc70..cfbd6a506 100644 --- a/crates/sam2/src/weights.rs +++ b/crates/sam2/src/weights.rs @@ -22,6 +22,12 @@ impl Weights { Weights::from_map(map) } + /// Read a checkpoint held in memory (the installed file's bytes, already verified by the caller). + pub fn from_bytes(bytes: &[u8]) -> Result { + let map = candle_core::safetensors::load_buffer(bytes, &Device::Cpu).map_err(|e| Error::Weights(e.to_string()))?; + Weights::from_map(map) + } + /// Validate an already-loaded set of tensors. pub fn from_map(map: HashMap) -> Result { let mut out = HashMap::with_capacity(map.len()); diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index 70af32cd0..63c1a06ed 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -301,3 +301,39 @@ fn tracker_rejects_hostile_input_and_honours_cancel() { let mut bad = VideoTracker::new(model, Box::new(Bad)).unwrap(); assert!(matches!(bad.prompt(0, obj, &click), Err(Error::Input(_)))); } + +#[test] +fn sam2_works_as_a_mask_model() { + use effectcraft_segment::{MaskModel, Prompt}; + let Some((w, _)) = load() else { return }; + let (fw, fh) = (256usize, 192usize); + let mut rgb = vec![[0f32; 3]; fw * fh]; + let inside = |x: usize, y: usize| ((x as f32 - 120.0).powi(2) + (y as f32 - 96.0).powi(2)).sqrt() < 50.0; + for y in 0..fh { + for x in 0..fw { + let bg = 0.3 + 0.2 * (x as f32 / 9.0).sin() * (y as f32 / 7.0).cos(); + rgb[y * fw + x] = if inside(x, y) { [0.85 + 0.1 * (x as f32 / 5.0).sin(), 0.4, 0.15] } else { [bg, bg * 0.8 + 0.1, bg * 0.6 + 0.3] }; + } + } + let model = effectcraft_sam2::mask::Sam2Mask::new(std::sync::Arc::new(effectcraft_sam2::session::VideoModel::load(&w).unwrap())); + let prompt = Prompt { points: vec![([120.0, 96.0], true)], bbox: None, mask: None }; + let p = model.segment(&rgb, fw, fh, &prompt).unwrap(); + assert_eq!(p.len(), fw * fh); + assert!(p.iter().all(|v| (0.0..=1.0).contains(v))); + let (mut i, mut u) = (0usize, 0usize); + for y in 0..fh { + for x in 0..fw { + let (a, b) = (p[y * fw + x] > 0.5, inside(x, y)); + i += usize::from(a && b); + u += usize::from(a || b); + } + } + let j = i as f64 / u as f64; + eprintln!("SAM 2 as a MaskModel: disc IoU {j:.3}"); + assert!(j > 0.9, "IoU {j}"); + // The same frame again is served from the cached encoding, and bad prompts are errors. + let again = model.segment(&rgb, fw, fh, &Prompt { points: vec![([20.0, 20.0], false), ([120.0, 96.0], true)], bbox: None, mask: None }).unwrap(); + assert_eq!(again.len(), fw * fh); + assert!(model.segment(&rgb, fw, fh, &Prompt::default()).is_err()); + assert!(model.segment(&rgb, fw + 1, fh, &prompt).is_err()); +}