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 <noreply@anthropic.com>

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]]

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 }

crates/engine/src/analysis/models.rs

@@ -202,6 +202,11 @@ impl Session {
 
 fn load_file(path: &Path, id: &str) -> std::result::Result<Loaded, String> {
     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)
 }
 

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(&params).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(&params).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(&params).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") {

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]

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;

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<Tensor> {
+    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<VideoModel>,
+    last: Mutex<Option<(u64, Arc<Encoded>)>>,
+}
+
+impl Sam2Mask {
+    pub fn new(model: Arc<VideoModel>) -> Sam2Mask {
+        Sam2Mask { model, last: Mutex::new(None) }
+    }
+
+    /// The video model, for trackers built on the same weights.
+    pub fn video_model(&self) -> &Arc<VideoModel> {
+        &self.model
+    }
+
+    fn encoded(&self, rgb: &[[f32; 3]], w: usize, h: usize) -> Result<Arc<Encoded>> {
+        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<Vec<f32>> {
+        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::<f32>()?;
+        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::<f32>()?)
+    }
+}
+
+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<Vec<f32>> {
+        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<Arc<Sam2Mask>, 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::<f32>().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}");
+    }
+}

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<Weights> {
+        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<String, Tensor>) -> Result<Weights> {
         let mut out = HashMap::with_capacity(map.len());

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