Roto Brush: validate GPU layouts and align object outlines (#1539)

crates/sam2/src/gpu/mod.rs

@@ -306,11 +306,20 @@ impl Gpu {
         if op.m == 0 || op.n == 0 || op.batch == 0 {
             return Ok(());
         }
+        if sc.col != 1 {
+            return Err(Error::Input(format!("matmul output needs a column stride of 1, got {}", sc.col)));
+        }
         // bounds: the last element each operand touches must lie inside its buffer
-        let last = |s: Strides, rows: usize, cols: usize| s.offset + (op.batch - 1) * s.batch + (rows - 1) * s.row + (cols.max(1) - 1) * s.col;
+        let last = |s: Strides, rows: usize, cols: usize| {
+            s.offset.checked_add((op.batch - 1).checked_mul(s.batch)?)?.checked_add((rows - 1).checked_mul(s.row)?)?.checked_add((cols - 1).checked_mul(s.col)?)
+        };
         for (name, buf, s, rows, cols) in [("A", a, sa, op.m, op.k), ("B", b, sb, op.k, op.n), ("C", c, sc, op.m, op.n)] {
-            if op.k > 0 && last(s, rows, cols) >= buf.len {
-                return Err(Error::Input(format!("matmul operand {name} reaches element {} of a {}-element buffer", last(s, rows, cols), buf.len)));
+            if name != "C" && op.k == 0 {
+                continue;
+            }
+            let last = last(s, rows, cols).ok_or_else(|| Error::Input(format!("matmul operand {name} layout overflows its element index")))?;
+            if last >= buf.len {
+                return Err(Error::Input(format!("matmul operand {name} reaches element {last} of a {}-element buffer", buf.len)));
             }
         }
         if op.bias.is_some_and(|bias| bias.len < op.n) || op.residual.is_some_and(|r| r.len < c.len) {

crates/sam2/src/gpu/tests.rs

@@ -86,6 +86,81 @@ fn matmul_matches_the_definition_for_awkward_sizes_and_strides() {
     }
 }
 
+#[test]
+fn matmul_rejects_nonunit_output_columns() {
+    let Some(g) = gpu() else { return };
+    let value = g.upload(&[1.0]);
+    let c = g.zeros(8).unwrap();
+    for col in [0, 2] {
+        let broadcast = Strides { offset: 0, batch: 0, row: 0, col: 0 };
+        let out = Strides { offset: 0, batch: 0, row: 4, col };
+        let op = Matmul {
+            m: 2,
+            n: 2,
+            k: 1,
+            batch: 1,
+            a: (&value, broadcast),
+            b: (&value, broadcast),
+            c: (&c, out),
+            bias: None,
+            residual: None,
+            alpha: 1.0,
+            act: Act::None,
+        };
+        assert!(matches!(g.matmul(&op), Err(Error::Input(_))), "output column stride {col}");
+    }
+}
+
+#[test]
+fn matmul_zero_inner_dimension_checks_output_and_writes_bias() {
+    let Some(g) = gpu() else { return };
+    let empty = g.upload(&[]);
+    let tiny = g.zeros(1).unwrap();
+    let bias = g.upload(&[3.0, 5.0]);
+    let mut op = Matmul {
+        m: 2,
+        n: 2,
+        k: 0,
+        batch: 1,
+        a: (&empty, Strides::row_major(0)),
+        b: (&empty, Strides::row_major(2)),
+        c: (&tiny, Strides::row_major(2)),
+        bias: Some(&bias),
+        residual: None,
+        alpha: 1.0,
+        act: Act::None,
+    };
+    assert!(matches!(g.matmul(&op), Err(Error::Input(_))), "C is written even when K is zero");
+    let c = g.zeros(4).unwrap();
+    op.c = (&c, Strides::row_major(2));
+    g.matmul(&op).unwrap();
+    g.flush().unwrap();
+    assert_eq!(g.download(&c).unwrap(), [3.0, 5.0, 3.0, 5.0]);
+}
+
+#[test]
+fn matmul_rejects_overflowing_operand_offsets() {
+    let Some(g) = gpu() else { return };
+    let value = g.upload(&[1.0]);
+    let c = g.zeros(4).unwrap();
+    let broadcast = Strides { offset: 0, batch: 0, row: 0, col: 0 };
+    let out = Strides { offset: usize::MAX, batch: 0, row: 1, col: 1 };
+    let op = Matmul {
+        m: 2,
+        n: 1,
+        k: 1,
+        batch: 1,
+        a: (&value, broadcast),
+        b: (&value, broadcast),
+        c: (&c, out),
+        bias: None,
+        residual: None,
+        alpha: 1.0,
+        act: Act::None,
+    };
+    assert!(matches!(g.matmul(&op), Err(Error::Input(_))));
+}
+
 #[test]
 fn matmul_with_transposed_operands_bias_activation_and_residual() {
     let Some(g) = gpu() else { return };

crates/sam2/src/neck.rs

@@ -110,6 +110,11 @@ pub struct NeckOutput {
 }
 
 impl Neck {
+    /// Position tables shared by the CPU and GPU paths, released with this neck.
+    pub(crate) fn position(&self, h: usize, w: usize) -> Result<Tensor> {
+        self.positions.get(self.pos_features, h, w)
+    }
+
     pub fn load(w: &Weights, channels: &[usize], fpn: usize, top_down: &[usize]) -> Result<Neck> {
         let convs = channels
             .iter()
@@ -144,7 +149,7 @@ impl Neck {
                 }
                 _ => lateral,
             };
-            positions.push(self.positions.get(self.pos_features, cur.dim(2)?, cur.dim(3)?)?);
+            positions.push(self.position(cur.dim(2)?, cur.dim(3)?)?);
             features.push(cur.clone());
             prev = Some(cur);
         }
@@ -199,16 +204,17 @@ mod tests {
 
     #[test]
     fn position_cache_bounds_retained_elements_and_preserves_the_encoding() {
-        let cache = PositionCache::new(64);
-        let first = cache.get(2, 2, 4).unwrap();
+        let neck = Neck { convs: vec![], top_down: vec![], pos_features: 2, positions: PositionCache::new(64) };
+        let cache = &neck.positions;
+        let first = neck.position(2, 4).unwrap();
         assert_eq!(first.dims(), &[1, 4, 2, 4]);
-        assert_eq!(first.flatten_all().unwrap().to_vec1::<f32>().unwrap(), cache.get(2, 2, 4).unwrap().flatten_all().unwrap().to_vec1::<f32>().unwrap());
-        cache.get(2, 3, 4).unwrap();
+        assert_eq!(first.flatten_all().unwrap().to_vec1::<f32>().unwrap(), neck.position(2, 4).unwrap().flatten_all().unwrap().to_vec1::<f32>().unwrap());
+        neck.position(3, 4).unwrap();
         let entries = cache.entries.lock().unwrap();
         assert_eq!(entries.len(), 1, "older tables are evicted to fit the budget");
         assert!(entries.values().map(Tensor::elem_count).sum::<usize>() <= 64);
         drop(entries);
-        let large = cache.get(2, 8, 8).unwrap();
+        let large = neck.position(8, 8).unwrap();
         assert_eq!(large.dims(), &[1, 4, 8, 8]);
         assert_eq!(cache.entries.lock().unwrap().len(), 1, "a table over budget is returned without retention");
     }

crates/sam2/src/session.rs

@@ -203,7 +203,7 @@ impl VideoModel {
             match g.encoder.encode(pixels) {
                 Ok(f) => {
                     let none = || Tensor::zeros((1, 0, 0, 0), DType::F32, &candle_core::Device::Cpu);
-                    let pos = crate::neck::sine_position(128, GRID, GRID)?;
+                    let pos = self.image.vision.neck.position(GRID, GRID)?;
                     let features = ImageFeatures { features: vec![none()?, none()?, f.low], positions: vec![none()?, none()?, pos] };
                     return Ok(Encoded { features, high: f.high });
                 }

crates/ui-egui/src/panels/roto_tool.rs

@@ -32,6 +32,31 @@ pub fn object_color(o: u32, base: [f64; 3]) -> Color32 {
     Color32::from_rgb(r, g, b)
 }
 
+/// Sample the layer-grid matte at each input pixel centre, including cropping and preview scale.
+fn object_mask_in_input(seg: &fx::FrameSeg, inp: &effectcraft_engine::effects::Buf) -> Vec<bool> {
+    let width = inp.img.width as usize;
+    if width == 0 || !inp.scale.is_finite() || inp.scale <= 0.0 || inp.offset.iter().any(|v| !v.is_finite()) {
+        return vec![false; inp.img.data.len()];
+    }
+    inp.img
+        .data
+        .iter()
+        .enumerate()
+        .map(|(i, _)| {
+            let x = ((i % width) as f64 + 0.5 - inp.offset[0]) / inp.scale;
+            let y = ((i / width) as f64 + 0.5 - inp.offset[1]) / inp.scale;
+            if x < 0.0 || y < 0.0 || x >= seg.w as f64 || y >= seg.h as f64 {
+                return false;
+            }
+            (y.floor() as usize)
+                .checked_mul(seg.w)
+                .and_then(|row| row.checked_add(x.floor() as usize))
+                .and_then(|index| seg.matte.get(index))
+                .is_some_and(|m| *m != 0)
+        })
+        .collect()
+}
+
 /// A stroke being drawn.
 #[derive(Clone, Debug)]
 struct Drag {
@@ -113,10 +138,7 @@ pub fn draw_overlay(
                     let f = fx::frame_of(layer.layer_time(time).seconds(), comp.frame_rate.as_f64());
                     fx::object_mattes(&params, size, comp.frame_rate.as_f64(), f)
                         .into_iter()
-                        .map(|(obj, seg)| {
-                            let seg = fx::resample(&seg, w, h);
-                            (object_color(obj, o.boundary_color), seg.matte.iter().map(|m| *m != 0).collect())
-                        })
+                        .map(|(obj, seg)| (object_color(obj, o.boundary_color), object_mask_in_input(&seg, &inp)))
                         .collect()
                 }
                 _ => vec![],
@@ -148,8 +170,18 @@ pub fn draw_overlay(
                                 if edge { bc } else { Color32::TRANSPARENT }
                             } else {
                                 // the outline of every object whose own matte has an edge here
-                                let inside_of =
-                                    |m: &Vec<bool>, x: i64, y: i64| x >= 0 && y >= 0 && x < w as i64 && y < h as i64 && m[y as usize * w + x as usize];
+                                let inside_of = |m: &Vec<bool>, x: i64, y: i64| {
+                                    x >= 0
+                                        && y >= 0
+                                        && x < w as i64
+                                        && y < h as i64
+                                        && (y as usize)
+                                            .checked_mul(w)
+                                            .and_then(|row| row.checked_add(x as usize))
+                                            .and_then(|index| m.get(index))
+                                            .copied()
+                                            .unwrap_or(false)
+                                };
                                 per_object
                                     .iter()
                                     .rev()
@@ -468,6 +500,24 @@ pub fn bar_buttons(app: &mut EffectcraftApp, ui: &mut egui::Ui, p: &egui::Painte
 mod tests {
     use super::*;
 
+    #[test]
+    fn per_object_outlines_follow_cropped_padded_and_scaled_input_coordinates() {
+        use effectcraft_engine::{effects::Buf, raster::Image};
+        let seg = fx::FrameSeg { w: 4, h: 2, matte: vec![0, 0, 255, 255, 0, 0, 255, 255], refine: vec![0; 8], ..Default::default() };
+        for (offset, scale, expected) in [
+            ([-2.0, 0.0], 1.0, vec![true, true]),
+            ([0.0, 0.0], 0.5, vec![false, true]),
+            ([-1.0, 0.0], 0.5, vec![true, false]),
+            ([2.0, 0.0], 1.0, vec![false, false]),
+        ] {
+            let inp = Buf { img: Image::new(2, 1), offset, scale };
+            assert_eq!(object_mask_in_input(&seg, &inp), expected, "offset={offset:?}, scale={scale}");
+        }
+        let inp = Buf { img: Image::new(6, 4), offset: [1.0, 1.0], scale: 1.0 };
+        let expected: Vec<bool> = (0..24).map(|i| matches!((i % 6, i / 6), (3 | 4, 1 | 2))).collect();
+        assert_eq!(object_mask_in_input(&seg, &inp), expected);
+    }
+
     #[test]
     fn confidence_colours_run_from_green_to_red_and_absent_is_slate() {
         let (g, a, r, s) = (trust_color(0.99, false), trust_color(0.6, false), trust_color(0.1, false), trust_color(0.99, true));

docs/roto-brush-sam2.md

@@ -4,9 +4,10 @@ Roto Brush can follow an object through a clip with Meta's **SAM 2.1 Tiny** (Hie
 Apache-2.0). It is optional: nothing changes until you install the model and choose it.
 
 ## Using it
-1. Settings ▸ Roto Brush ▸ SAM 2.1 Tiny ▸ **Install** (downloads `model.safetensors`, 155.9 MB, from the pinned
+1. Settings ▸ Roto Brush ▸ SAM 2.1 Tiny ▸ **Download** (downloads `model.safetensors`, 155.9 MB, from the pinned
    Hugging Face commit; the SHA-256 is checked and the licence notice is kept next to the file). Or install from a
-   file you downloaded. Agents: `roto.model.install {id: "sam2.1-hiera-tiny"}`.
+   file you downloaded. Agents: `roto.model.download {id: "sam2.1-hiera-tiny", wait: true}`, or
+   `roto.model.install {id: "sam2.1-hiera-tiny", path: "model.safetensors"}` for a file you already have.
 2. Choose it (`roto.model.select {id: "sam2.1-hiera-tiny"}`).
 3. Paint strokes on the base frame as usual, then **Propagate**. One tracking pass follows the object through the
    whole span in both directions, using SAM 2's streaming memory, so it keeps tracking through occlusion and when the