From c4dc558a9830dc54d10c33e5bce0007843bf7d16 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 29 Oct 2026 21:33:33 +0100 Subject: [PATCH] Roto Brush: validate GPU layouts and align object outlines (#1539) --- crates/sam2/src/gpu/mod.rs | 15 ++++-- crates/sam2/src/gpu/tests.rs | 75 ++++++++++++++++++++++++++ crates/sam2/src/neck.rs | 18 ++++--- crates/sam2/src/session.rs | 2 +- crates/ui-egui/src/panels/roto_tool.rs | 62 ++++++++++++++++++--- docs/roto-brush-sam2.md | 5 +- 6 files changed, 159 insertions(+), 18 deletions(-) diff --git a/crates/sam2/src/gpu/mod.rs b/crates/sam2/src/gpu/mod.rs index fab8573e1..3193dcf80 100644 --- a/crates/sam2/src/gpu/mod.rs +++ b/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) { diff --git a/crates/sam2/src/gpu/tests.rs b/crates/sam2/src/gpu/tests.rs index 544df73e0..674bd3b8b 100644 --- a/crates/sam2/src/gpu/tests.rs +++ b/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 }; diff --git a/crates/sam2/src/neck.rs b/crates/sam2/src/neck.rs index 3a80ede9e..50733fe79 100644 --- a/crates/sam2/src/neck.rs +++ b/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 { + self.positions.get(self.pos_features, h, w) + } + pub fn load(w: &Weights, channels: &[usize], fpn: usize, top_down: &[usize]) -> Result { 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::().unwrap(), cache.get(2, 2, 4).unwrap().flatten_all().unwrap().to_vec1::().unwrap()); - cache.get(2, 3, 4).unwrap(); + assert_eq!(first.flatten_all().unwrap().to_vec1::().unwrap(), neck.position(2, 4).unwrap().flatten_all().unwrap().to_vec1::().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::() <= 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"); } diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 54a39d212..48ec72a31 100644 --- a/crates/sam2/src/session.rs +++ b/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 }); } diff --git a/crates/ui-egui/src/panels/roto_tool.rs b/crates/ui-egui/src/panels/roto_tool.rs index 0b42befbe..cfed82e9b 100644 --- a/crates/ui-egui/src/panels/roto_tool.rs +++ b/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 { + 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(¶ms, 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, 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, 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 = (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)); diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index a80052bab..f2a2cb20d 100644 --- a/docs/roto-brush-sam2.md +++ b/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