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()); +}