From 71e44deb84248476f946ed1ed7b773a40bde0d1c Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 00:51:46 +0200 Subject: [PATCH 01/77] S2.50: effectcraft-cli bench --roto on synthetic clips Synthetic Roto Brush clips with known mattes (movement, lighting change, occlusion, leaving and re-entering) and a runner that compares the classical engine with MobileSAM when installed. Needs no weights; SAM 2 plugs in later. Co-Authored-By: Claude Sonnet 5.5 --- Cargo.lock | 2 + apps/effectcraft-cli/Cargo.toml | 2 + apps/effectcraft-cli/src/main.rs | 40 +++++ crates/track/src/roto.rs | 1 + crates/track/src/roto/bench.rs | 264 +++++++++++++++++++++++++++++++ 5 files changed, 309 insertions(+) create mode 100644 crates/track/src/roto/bench.rs diff --git a/Cargo.lock b/Cargo.lock index 07564b7df..6a694299d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1394,7 +1394,9 @@ dependencies = [ "effectcraft-gpu", "effectcraft-host", "effectcraft-render", + "effectcraft-segment", "effectcraft-time", + "effectcraft-track", "effectcraft-ui-egui", "image", "serde_json", diff --git a/apps/effectcraft-cli/Cargo.toml b/apps/effectcraft-cli/Cargo.toml index b5784f4eb..7d8fe96ff 100644 --- a/apps/effectcraft-cli/Cargo.toml +++ b/apps/effectcraft-cli/Cargo.toml @@ -20,6 +20,8 @@ effectcraft-ui-egui = { workspace = true } effectcraft-render = { workspace = true } effectcraft-gpu = { workspace = true } effectcraft-time = { workspace = true } +effectcraft-track = { workspace = true } +effectcraft-segment = { workspace = true } serde_json = { workspace = true } image = { workspace = true } cpu-time = "1" diff --git a/apps/effectcraft-cli/src/main.rs b/apps/effectcraft-cli/src/main.rs index fd77162ed..3de4449b6 100644 --- a/apps/effectcraft-cli/src/main.rs +++ b/apps/effectcraft-cli/src/main.rs @@ -27,6 +27,8 @@ //! (--gpu: CPU vs GPU ms/frame for every comp at Full and Half) //! effectcraft-cli bench --ops [--small] [--layers N] [--comps N] [--footage N] everyday-operation timings //! on a large generated project (open, save, auto-save, undo/redo, timeline, Project panel) +//! effectcraft-cli bench --roto [--frames N] [--json] Roto Brush trackers on synthetic clips with known mattes +//! (classical engine, MobileSAM when installed; IoU, flicker, re-appearance, seconds per frame) //! effectcraft-cli script FILE.jsx [F.ecproj] | --eval CODE run an After Effects-style script //! effectcraft-cli mcp [--autosave | --bridge PORT] MCP server on stdio //! @@ -124,6 +126,8 @@ const USAGE: &str = "usage: effectcraft-cli Result<(), Failure> { if args.flag("--ops") { return bench_ops(args); } + if args.flag("--roto") { + return bench_roto(args); + } let mut s = effectcraft_host::session(); match &args.project { Some(p) => s.execute("file.open", json!({"path": p})), @@ -998,6 +1007,37 @@ fn bench_cmd(args: &Args) -> Result<(), Failure> { Ok(()) } +/// `bench --roto`: the Roto Brush's trackers on synthetic clips whose true mattes are known. +fn bench_roto(args: &Args) -> Result<(), Failure> { + use effectcraft_segment::{Loaded, MOBILE_SAM, MaskModel}; + use effectcraft_track::roto::bench; + let frames = args.num("--frames")?.unwrap_or(24.0).clamp(12.0, 240.0) as usize; + let session = effectcraft_host::session(); + let mobilesam: Option> = + session.model_path(&MOBILE_SAM).and_then(|p| std::fs::read(p).ok()).and_then(|bytes| effectcraft_segment::load(MOBILE_SAM.id, &bytes).ok()).and_then( + |m| match m { + Loaded::Mask(m) => Some(m), + _ => None, + }, + ); + if mobilesam.is_none() { + note!("bench --roto: MobileSAM is not installed; comparing the classical engine only"); + } + let mut runs = vec![]; + for sc in bench::scenarios(frames) { + runs.push(bench::run_flow(&sc, "classical", None)); + if let Some(m) = &mobilesam { + runs.push(bench::run_flow(&sc, "mobilesam", Some(m.as_ref()))); + } + } + if args.flag("--json") { + say!("{}", serde_json::to_string_pretty(&runs).map_err(|e| Failure::Error(e.to_string()))?); + } else { + say!("{}", bench::table(&runs)); + } + Ok(()) +} + /// `bench --ops`: everyday operations on a large generated project (200 comps, 5,000 layers in /// the main comp, 300 footage items, 20-deep precomp nesting, expressions): startup to first /// frame, open, save, auto-save, small and large edits with undo/redo, first frame, timeline diff --git a/crates/track/src/roto.rs b/crates/track/src/roto.rs index 103d695dd..abcf28aa0 100644 --- a/crates/track/src/roto.rs +++ b/crates/track/src/roto.rs @@ -31,6 +31,7 @@ //! deterministic function of its *chain*: the base frame's strokes, then every frame's strokes //! on the way out to it ([`RotoData::chain_keys`]). +pub mod bench; pub mod gmm; pub mod matting; pub mod maxflow; diff --git a/crates/track/src/roto/bench.rs b/crates/track/src/roto/bench.rs new file mode 100644 index 000000000..d52836979 --- /dev/null +++ b/crates/track/src/roto/bench.rs @@ -0,0 +1,264 @@ +//! Roto Brush benchmark: synthetic clips with known mattes (movement, colour change, occlusion, +//! leaving and re-entering the frame) and the metrics that compare trackers on them. +//! +//! Needs no weights: the classical engine always runs, and a trained [`MaskModel`] (MobileSAM, +//! later SAM 2) is added when the caller has one. `effectcraft-cli bench --roto` prints the table. + +use effectcraft_raster::Image; +use effectcraft_segment::MaskModel; +use serde::Serialize; +use web_time::Instant; + +use super::{FrameSeg, SegOpts, Stroke, StrokeKind, iou, propagate_with, segment_with}; + +pub const W: u32 = 200; +pub const H: u32 = 150; +const R: f64 = 32.0; + +/// A clip with the true matte of the visible object on every frame. +pub struct Scenario { + pub name: &'static str, + pub what: &'static str, + pub frames: Vec, + /// Visible-object matte per frame (`W×H`, 0/1); empty when the object is out of view. + pub truth: Vec>, + /// Frame the strokes are painted on. + pub base: usize, + pub strokes: Vec, +} + +/// One tracker's result on one scenario. +#[derive(Clone, Debug, Serialize)] +pub struct Run { + pub scenario: &'static str, + pub tracker: String, + pub mean_iou: f64, + pub min_iou: f64, + /// Mean IoU between consecutive mattes minus the same for the truth, as a loss (0 is steady). + pub flicker: f64, + pub secs_per_frame: f64, + /// Mean IoU over the frames after the object's last absence (re-appearance), if it had one. + pub after_return: Option, + pub per_frame: Vec, +} + +fn noise(x: f64, y: f64, seed: u32) -> f64 { + fn h(i: i64, j: i64, s: u32) -> f64 { + let mut v = (i.wrapping_mul(73856093) ^ j.wrapping_mul(19349663) ^ i64::from(s).wrapping_mul(83492791)) as u64; + v ^= v >> 13; + v = v.wrapping_mul(0x5bd1e995); + v ^= v >> 15; + (v % 10_000) as f64 / 10_000.0 + } + let mut out = 0.0; + for (cell, amp) in [(9.0, 0.6), (3.5, 0.4)] { + let (fx, fy) = (x / cell, y / cell); + let (i, j) = (fx.floor() as i64, fy.floor() as i64); + let (tx, ty) = (fx - i as f64, fy - j as f64); + let (sx, sy) = (tx * tx * (3.0 - 2.0 * tx), ty * ty * (3.0 - 2.0 * ty)); + let a = h(i, j, seed) + (h(i + 1, j, seed) - h(i, j, seed)) * sx; + let b = h(i, j + 1, seed) + (h(i + 1, j + 1, seed) - h(i, j + 1, seed)) * sx; + out += amp * (a + (b - a) * sy); + } + out +} + +/// What a frame shows: the disk centre, an optional occluding bar `[x0, y0, x1, y1]` and a +/// brightness gain applied to the whole frame. +#[derive(Clone, Copy)] +struct Look { + c: [f64; 2], + bar: Option<[f64; 4]>, + gain: f32, +} + +fn render(l: Look) -> (Image, Vec) { + let mut img = Image::new(W, H); + let mut gt = vec![0u8; (W * H) as usize]; + for y in 0..H { + for x in 0..W { + let p = [f64::from(x) + 0.5, f64::from(y) + 0.5]; + let d = (p[0] - l.c[0]).hypot(p[1] - l.c[1]); + let a = (R + 0.5 - d).clamp(0.0, 1.0) as f32; + let n = noise(p[0] - l.c[0], p[1] - l.c[1], 11); + let f = [(0.75 + 0.25 * n) as f32, (0.35 + 0.3 * n) as f32, (0.1 + 0.15 * n) as f32]; + let (bn, bm) = (noise(p[0], p[1], 5), noise(p[0] * 1.7, p[1] * 1.7, 9)); + let b = [(0.05 + 0.25 * bm) as f32, (0.3 + 0.4 * bn) as f32, (0.45 + 0.4 * bm) as f32]; + let mut px = [0.0f32; 3]; + for k in 0..3 { + px[k] = f[k] * a + b[k] * (1.0 - a); + } + let under_bar = l.bar.is_some_and(|r| p[0] >= r[0] && p[0] < r[2] && p[1] >= r[1] && p[1] < r[3]); + if under_bar { + let t = noise(p[0] * 0.8, p[1] * 0.8, 23) as f32; + px = [0.15 + 0.1 * t, 0.15 + 0.1 * t, 0.2 + 0.1 * t]; + } + let g = l.gain; + img.set(x, y, [(px[0] * g).min(1.0), (px[1] * g).min(1.0), (px[2] * g).min(1.0), 1.0]); + if let Some(v) = gt.get_mut((y * W + x) as usize) { + *v = (d <= R && !under_bar) as u8; + } + } + } + (img, gt) +} + +fn line(kind: StrokeKind, frame: i64, a: [f64; 2], b: [f64; 2], radius: f64) -> Stroke { + let pt = |k: u32| { + let t = f64::from(k) / 10.0; + [a[0] + (b[0] - a[0]) * t, a[1] + (b[1] - a[1]) * t] + }; + Stroke { kind, frame, radius, points: (0..=10).map(pt).collect() } +} + +/// The strokes an artist would paint on the base frame: one through the disk, two around it. +fn strokes_for(c: [f64; 2]) -> Vec { + vec![ + line(StrokeKind::Fg, 0, [c[0] - 18.0, c[1] - 6.0], [c[0] + 18.0, c[1] + 6.0], 5.0), + line(StrokeKind::Bg, 0, [c[0] - 50.0, c[1] - 45.0], [c[0] + 50.0, c[1] - 45.0], 5.0), + line(StrokeKind::Bg, 0, [c[0] + 50.0, c[1] + 45.0], [c[0] - 45.0, c[1] + 45.0], 5.0), + ] +} + +fn scenario(name: &'static str, what: &'static str, n: usize, look: impl Fn(usize) -> Look) -> Scenario { + let (mut frames, mut truth) = (Vec::with_capacity(n), Vec::with_capacity(n)); + for f in 0..n { + let (img, gt) = render(look(f)); + frames.push(img); + truth.push(gt); + } + Scenario { name, what, frames, truth, base: 0, strokes: strokes_for(look(0).c) } +} + +/// The benchmark clips, `n` frames each (at least 12, so every event happens). +pub fn scenarios(n: usize) -> Vec { + let n = n.clamp(12, 240); + let k = n as f64 / 24.0; + let drift = move |f: usize| { + let t = f as f64 / k; + [60.0 + 3.1 * t, 70.0 + 1.3 * t + 4.0 * (t * 0.3).sin()] + }; + vec![ + scenario("move", "disk drifts across a textured background", n, move |f| Look { c: drift(f), bar: None, gain: 1.0 }), + scenario("light", "same motion while the whole frame brightens by 40 %", n, move |f| Look { + c: drift(f), + bar: None, + gain: 1.0 + 0.4 * f as f32 / n as f32, + }), + scenario("occlusion", "a bar crosses and hides most of the disk, then clears it", n, move |f| { + let t = f as f64 / n as f64; + let c = drift(f); + let x = -40.0 + 320.0 * ((t - 0.3) / 0.4).clamp(0.0, 1.0); + Look { c, bar: (0.3..0.7).contains(&t).then_some([x - 40.0, 0.0, x + 40.0, f64::from(H)]), gain: 1.0 } + }), + scenario("reenter", "disk leaves through the right edge and comes back", n, move |f| { + let t = f as f64 / n as f64; + let x = 60.0 + 190.0 * (1.0 - (2.0 * t - 1.0).powi(2)); + Look { c: [x, 70.0 + 4.0 * (t * 6.0).sin()], bar: None, gain: 1.0 } + }), + ] +} + +fn metrics(sc: &Scenario, tracker: &str, mattes: &[Vec], secs: f64) -> Run { + let per_frame: Vec = mattes.iter().zip(&sc.truth).map(|(m, t)| iou(m, t)).collect(); + let n = per_frame.len().max(1) as f64; + let steady = |v: &[Vec]| { + let pairs: Vec = v.windows(2).filter_map(|w| Some(iou(w.first()?, w.get(1)?))).collect(); + pairs.iter().sum::() / pairs.len().max(1) as f64 + }; + let last_gap = sc.truth.iter().rposition(|t| t.iter().all(|v| *v == 0)); + let after_return = last_gap.and_then(|g| { + let tail = per_frame.get(g + 1..)?; + (!tail.is_empty()).then(|| tail.iter().sum::() / tail.len() as f64) + }); + Run { + scenario: sc.name, + tracker: tracker.to_string(), + mean_iou: per_frame.iter().sum::() / n, + min_iou: per_frame.iter().copied().fold(1.0, f64::min), + flicker: (steady(&sc.truth) - steady(mattes)).max(0.0), + secs_per_frame: secs / n, + after_return, + per_frame, + } +} + +/// Segment the base frame from the strokes and propagate to the end of the clip, as the Roto +/// Brush does, with `model` (or none: the classical engine). +pub fn run_flow(sc: &Scenario, tracker: &str, model: Option<&dyn MaskModel>) -> Run { + let opts = SegOpts::default(); + let start = Instant::now(); + let refs: Vec<&Stroke> = sc.strokes.iter().collect(); + let mut segs: Vec> = vec![None; sc.frames.len()]; + if let Some(first) = sc.frames.get(sc.base) { + let base = segment_with(first, &refs, None, &opts, 0.0, model); + if let Some(slot) = segs.get_mut(sc.base) { + *slot = Some(base); + } + } + for f in sc.base + 1..sc.frames.len() { + let (Some(prev_img), Some(img)) = (sc.frames.get(f - 1), sc.frames.get(f)) else { break }; + let Some(Some(prev)) = segs.get(f - 1) else { break }; + let next = propagate_with(prev_img, prev, img, &[], &opts, 0.0, model); + if let Some(slot) = segs.get_mut(f) { + *slot = Some(next); + } + } + let secs = start.elapsed().as_secs_f64(); + let empty = vec![0u8; (W * H) as usize]; + let mattes: Vec> = segs.into_iter().map(|s| s.map_or_else(|| empty.clone(), |s| s.matte)).collect(); + metrics(sc, tracker, &mattes, secs) +} + +/// A fixed-width table of the runs, grouped by scenario. +pub fn table(runs: &[Run]) -> String { + let mut out = format!("{:<10} {:<12} {:>8} {:>8} {:>8} {:>10} {:>9}\n", "clip", "tracker", "meanIoU", "minIoU", "flicker", "after ret.", "s/frame"); + for r in runs { + let back = r.after_return.map_or_else(|| "-".to_string(), |v| format!("{v:.3}")); + out.push_str(&format!( + "{:<10} {:<12} {:>8.3} {:>8.3} {:>8.3} {:>10} {:>9.3}\n", + r.scenario, r.tracker, r.mean_iou, r.min_iou, r.flicker, back, r.secs_per_frame + )); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + fn area(v: &[u8]) -> usize { + v.iter().filter(|x| **x != 0).count() + } + + #[test] + fn scenarios_have_the_events_they_claim() { + let all = scenarios(24); + assert_eq!(all.len(), 4); + for s in &all { + assert_eq!((s.frames.len(), s.truth.len()), (24, 24), "{}", s.name); + } + let full = all.first().map(|s| area(&s.truth[0])).unwrap_or(0); + assert!(full > 3000, "{full}"); + let occ = all.iter().find(|s| s.name == "occlusion").map(|s| s.truth.iter().map(|t| area(t)).min().unwrap_or(0)).unwrap_or(full); + assert!(occ < full / 2, "occluded minimum {occ} of {full}"); + let gone = all.iter().find(|s| s.name == "reenter").map(|s| s.truth.iter().filter(|t| area(t) == 0).count()).unwrap_or(0); + assert!(gone > 0, "the disk never leaves"); + assert!(all.iter().find(|s| s.name == "reenter").is_some_and(|s| area(s.truth.last().map_or(&[][..], |v| v)) > 0), "it must come back"); + } + + #[test] + fn classical_tracks_plain_motion_and_the_table_lists_it() { + let all = scenarios(24); + let Some(sc) = all.first() else { return }; + let run = run_flow(sc, "classical", None); + assert!(run.mean_iou > 0.9, "mean IoU {}", run.mean_iou); + assert_eq!(run.per_frame.len(), 24); + assert!(table(&[run]).contains("classical")); + } + + #[test] + fn tiny_clip_requests_are_clamped() { + assert!(scenarios(0).iter().all(|s| s.frames.len() == 12)); + } +} From 9cb48399c364083922f18aa5a907917f752387e9 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:06:47 +0200 Subject: [PATCH 02/77] S2.0: record the SAM 2 decisions D1-D6 Co-Authored-By: Claude Sonnet 5.5 --- docs/adr/0002-sam2-roto-tracker.md | 19 +++++++++++++++++++ tools/requirements-sam2-reference.txt | 4 ++++ 2 files changed, 23 insertions(+) create mode 100644 docs/adr/0002-sam2-roto-tracker.md create mode 100644 tools/requirements-sam2-reference.txt diff --git a/docs/adr/0002-sam2-roto-tracker.md b/docs/adr/0002-sam2-roto-tracker.md new file mode 100644 index 000000000..3807d84a2 --- /dev/null +++ b/docs/adr/0002-sam2-roto-tracker.md @@ -0,0 +1,19 @@ +# ADR 0002: SAM 2 video tracker for the Roto Brush + +Status: accepted (maintainer answers, 2026-10-15). Plan: `plan/roto-sam2.md` (maintainer-local). + +## Context +The Roto Brush's classical engine loses an object that is occluded or leaves and re-enters the frame +(`effectcraft-cli bench --roto`: IoU 0.49 on the occlusion clip, 0.62 on re-entry, 0 after the object returns). +Meta's SAM 2.1 (Apache-2.0 code and checkpoints) has streaming memory built for this. + +## Decisions +- **D1** A native-only `candle` dependency may enter the workspace, in its own crate (`effectcraft-sam2`), excluded from wasm32. +- **D2** `effectcraft-sam2` is licensed Apache-2.0 only (it derives from Apache-2.0 Hugging Face `transformers` code); `effectcraft-segment` stays MIT OR Apache-2.0. +- **D3** Shared candle helpers are vendored with a NOTICE for now; a shared crate with LightCraft can follow. +- **D4** Weights are never hosted or bundled by us: direct from Hugging Face, a mirror, or a user folder, with the SHA-256 pinned in the registry. +- **D5** SAM 2 is opt-in at first; it becomes the default for tracking only after the benchmark shows it should. +- **D6** Hand-written kernels (the MobileSAM way) remain the Phase 6 route for AMD GPUs if candle cannot get there. + +## Consequences +Without the model installed nothing changes. Every ported file carries Apache-2.0 §4(b) modified-work headers. diff --git a/tools/requirements-sam2-reference.txt b/tools/requirements-sam2-reference.txt new file mode 100644 index 000000000..8fab732ff --- /dev/null +++ b/tools/requirements-sam2-reference.txt @@ -0,0 +1,4 @@ +# Pinned versions the SAM 2 reference fixtures were generated with (CPU). +torch==2.7.0 +transformers==5.19.0 +safetensors==0.8.0 From fd2aedb0008a16b1281677e2aadaf582495960aa Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:07:28 +0200 Subject: [PATCH 03/77] S2.2: reference script that dumps per-stage SAM 2 tensors Not yet run: it needs the weights. Co-Authored-By: Claude Sonnet 5.5 --- tools/sam2_reference.py | 65 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 tools/sam2_reference.py diff --git a/tools/sam2_reference.py b/tools/sam2_reference.py new file mode 100644 index 000000000..e40019134 --- /dev/null +++ b/tools/sam2_reference.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python +"""Dump per-stage SAM 2 reference tensors (CPU, float32) for the Rust port's reference tests. + +Usage: python -I tools/sam2_reference.py MODEL_DIR OUT.safetensors +MODEL_DIR holds the facebook/sam2.1-hiera-tiny files (config.json, model.safetensors, ...). +Versions are pinned in tools/requirements-sam2-reference.txt. The input is synthetic, so no +media is committed; the weights are never copied into the repo. +""" +import sys + +import numpy as np +import torch +from safetensors.torch import save_file +from transformers import Sam2Model, Sam2Processor + + +def synthetic_image(size=256): + """A textured disc over a textured background (deterministic).""" + rng = np.random.default_rng(7) + yy, xx = np.mgrid[0:size, 0:size].astype(np.float32) + bg = 0.3 + 0.2 * np.sin(xx / 9.0) * np.cos(yy / 7.0) + disc = np.hypot(xx - size * 0.45, yy - size * 0.5) < size * 0.22 + img = np.stack([bg, bg * 0.8 + 0.1, bg * 0.6 + 0.3], -1) + img[disc] = np.stack([0.85 + 0.1 * np.sin(xx / 5), 0.4 + 0.1 * np.cos(yy / 5), 0.15 + 0 * xx], -1)[disc] + img += rng.normal(0, 0.01, img.shape).astype(np.float32) + return (img.clip(0, 1) * 255).astype(np.uint8) + + +def main(model_dir, out): + torch.manual_seed(0) + torch.set_num_threads(1) + model = Sam2Model.from_pretrained(model_dir).eval() + proc = Sam2Processor.from_pretrained(model_dir) + image = synthetic_image() + inputs = proc(images=image, input_points=[[[[115.0, 128.0]]]], input_labels=[[[1]]], return_tensors="pt") + taps = {} + + def tap(name): + def hook(_m, _i, o): + outs = o if isinstance(o, (tuple, list)) else [o] + for k, t in enumerate(outs): + if isinstance(t, torch.Tensor): + taps[f"{name}.{k}"] = t.detach().float().contiguous() + return hook + + # Stage taps by module path: patch embed, each Hiera block, the neck, prompt encoder, decoder. + for name, mod in model.named_modules(): + parts = name.split(".") + if name in ("vision_encoder.backbone.patch_embed", "vision_encoder.neck", "prompt_encoder", "mask_decoder") or ( + len(parts) == 4 and parts[:3] == ["vision_encoder", "backbone", "blocks"] + ): + mod.register_forward_hook(tap(name)) + with torch.no_grad(): + res = model(**inputs, multimask_output=True) + taps["input.pixel_values"] = inputs["pixel_values"].float() + taps["out.iou_scores"] = res.iou_scores.float() + taps["out.pred_masks"] = res.pred_masks.float() + save_file({k: v for k, v in taps.items()}, out) + print(f"wrote {len(taps)} tensors to {out}") + + +if __name__ == "__main__": + if len(sys.argv) != 3: + sys.exit(__doc__) + main(sys.argv[1], sys.argv[2]) From 48997fcd3d42a133198340b62bf7817003bcbc37 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:10:28 +0200 Subject: [PATCH 04/77] S2.2: make the reference script run without torchvision Preprocess by hand, tap the neck's ModelOutput, report forward time. Co-Authored-By: Claude Sonnet 5.5 --- tools/sam2_reference.py | 36 ++++++++++++++++++++++++++++-------- 1 file changed, 28 insertions(+), 8 deletions(-) diff --git a/tools/sam2_reference.py b/tools/sam2_reference.py index e40019134..2cfd03fde 100644 --- a/tools/sam2_reference.py +++ b/tools/sam2_reference.py @@ -11,7 +11,8 @@ import sys import numpy as np import torch from safetensors.torch import save_file -from transformers import Sam2Model, Sam2Processor +import torch.nn.functional as F +from transformers import Sam2Model def synthetic_image(size=256): @@ -30,17 +31,33 @@ def main(model_dir, out): torch.manual_seed(0) torch.set_num_threads(1) model = Sam2Model.from_pretrained(model_dir).eval() - proc = Sam2Processor.from_pretrained(model_dir) image = synthetic_image() - inputs = proc(images=image, input_points=[[[[115.0, 128.0]]]], input_labels=[[[1]]], return_tensors="pt") + # Preprocess by hand (no torchvision needed): resize to 1024x1024, ImageNet mean/std. + x = torch.from_numpy(image).permute(2, 0, 1)[None].float() / 255.0 + x = F.interpolate(x, size=(1024, 1024), mode="bilinear", align_corners=False) + mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + scale = 1024.0 / image.shape[0] + inputs = { + "pixel_values": (x - mean) / std, + "input_points": torch.tensor([[[[115.0 * scale, 128.0 * scale]]]]), + "input_labels": torch.tensor([[[1]]], dtype=torch.int32), + } taps = {} def tap(name): def hook(_m, _i, o): - outs = o if isinstance(o, (tuple, list)) else [o] - for k, t in enumerate(outs): - if isinstance(t, torch.Tensor): - taps[f"{name}.{k}"] = t.detach().float().contiguous() + def flat(x): + if isinstance(x, torch.Tensor): + yield x + elif isinstance(x, (tuple, list)): + for y in x: + yield from flat(y) + elif hasattr(x, "to_tuple"): # a ModelOutput (the neck) + yield from flat(x.to_tuple()) + + for k, t in enumerate(flat(o)): + taps[f"{name}.{k}"] = t.detach().float().clone().contiguous() return hook # Stage taps by module path: patch embed, each Hiera block, the neck, prompt encoder, decoder. @@ -50,12 +67,15 @@ def main(model_dir, out): len(parts) == 4 and parts[:3] == ["vision_encoder", "backbone", "blocks"] ): mod.register_forward_hook(tap(name)) + import time with torch.no_grad(): + t0 = time.perf_counter() res = model(**inputs, multimask_output=True) + print(f"forward (1 thread, first run): {time.perf_counter() - t0:.2f} s") taps["input.pixel_values"] = inputs["pixel_values"].float() taps["out.iou_scores"] = res.iou_scores.float() taps["out.pred_masks"] = res.pred_masks.float() - save_file({k: v for k, v in taps.items()}, out) + save_file({k: v.clone().contiguous() for k, v in taps.items()}, out) print(f"wrote {len(taps)} tensors to {out}") From 36a22733a2226485e2aee543bf961b37383a7cfd Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:11:11 +0200 Subject: [PATCH 05/77] S2.4: effectcraft-sam2 crate skeleton Native-only, Apache-2.0, forbid(unsafe_code), panic lints denied, NOTICE, layer 3. Empty on wasm32. Co-Authored-By: Claude Sonnet 5.5 --- Cargo.lock | 4 +++ Cargo.toml | 1 + crates/sam2/Cargo.toml | 12 +++++++ crates/sam2/NOTICE | 11 ++++++ crates/sam2/src/lib.rs | 78 ++++++++++++++++++++++++++++++++++++++++++ xtask/src/main.rs | 1 + 6 files changed, 107 insertions(+) create mode 100644 crates/sam2/Cargo.toml create mode 100644 crates/sam2/NOTICE create mode 100644 crates/sam2/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index 6a694299d..abf1f898b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1765,6 +1765,10 @@ dependencies = [ "web-time", ] +[[package]] +name = "effectcraft-sam2" +version = "0.9.2" + [[package]] name = "effectcraft-script" version = "0.18.0" diff --git a/Cargo.toml b/Cargo.toml index afc59257b..28a011f24 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,7 @@ effectcraft-text = { path = "crates/text" } effectcraft-effects = { path = "crates/effects" } effectcraft-track = { path = "crates/track" } effectcraft-segment = { path = "crates/segment" } +effectcraft-sam2 = { path = "crates/sam2" } effectcraft-model = { path = "crates/model" } effectcraft-render = { path = "crates/render" } effectcraft-gpu = { path = "crates/gpu" } diff --git a/crates/sam2/Cargo.toml b/crates/sam2/Cargo.toml new file mode 100644 index 000000000..24c317bc1 --- /dev/null +++ b/crates/sam2/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "effectcraft-sam2" +description = "SAM 2.1 (Hiera) video object tracker for the Roto Brush: a Rust port of the Hugging Face transformers implementation, native only, weights downloaded by the user" +version.workspace = true +edition.workspace = true +license = "Apache-2.0" +rust-version.workspace = true + +[lints] +workspace = true + +[dependencies] diff --git a/crates/sam2/NOTICE b/crates/sam2/NOTICE new file mode 100644 index 000000000..2dcb613fa --- /dev/null +++ b/crates/sam2/NOTICE @@ -0,0 +1,11 @@ +EffectCraft SAM 2 (effectcraft-sam2) +Copyright the EffectCraft contributors. + +This crate is a Rust port of the SAM 2 implementation in Hugging Face `transformers` +(`models/sam2` and `models/sam2_video`), Copyright 2025 The HuggingFace Inc. team, which in turn +implements "SAM 2: Segment Anything in Images and Videos" by Meta AI (facebookresearch/sam2). +Both are licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0). + +Each ported source file states that it is a modified work (Apache-2.0 section 4(b)). +Model weights are not part of this crate; they are downloaded by the user from the official +`facebook/sam2.1-hiera-*` repositories (Apache-2.0). diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs new file mode 100644 index 000000000..d41de59fa --- /dev/null +++ b/crates/sam2/src/lib.rs @@ -0,0 +1,78 @@ +//! SAM 2.1 video object tracker (Hiera backbone) for the Roto Brush. +//! +//! A Rust port of the Hugging Face `transformers` SAM 2 implementation (Apache-2.0; see +//! `NOTICE`). Native only: the crate is empty on wasm32 so L0–L4 still build for the web. +//! Nothing here loads weights unless asked to; weights are never bundled. +//! +//! Modules are added stage by stage, each checked against reference tensors produced by +//! `tools/sam2_reference.py`. +#![cfg(not(target_arch = "wasm32"))] +#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] + +use std::fmt; + +/// Everything that can go wrong in this crate. Never a panic. +#[derive(Debug, Clone, PartialEq)] +pub enum Error { + /// The weights file is missing a tensor, has the wrong shape, or contains non-finite values. + Weights(String), + /// The input (frame size, prompt, frame index) is unusable. + Input(String), + /// The work was cancelled by the caller. + Cancelled, + /// A time or memory budget was exceeded. + Budget(String), +} + +impl fmt::Display for Error { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Error::Weights(m) => write!(f, "SAM 2 weights: {m}"), + Error::Input(m) => write!(f, "SAM 2 input: {m}"), + Error::Cancelled => write!(f, "SAM 2: cancelled"), + Error::Budget(m) => write!(f, "SAM 2 budget: {m}"), + } + } +} + +impl std::error::Error for Error {} + +pub type Result = std::result::Result; + +/// The released SAM 2.1 sizes. Only Tiny is wired up first; the others are configuration. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Size { + Tiny, + Small, + BasePlus, + Large, +} + +impl Size { + /// The Hugging Face repository the official weights come from. + pub fn repo(self) -> &'static str { + match self { + Size::Tiny => "facebook/sam2.1-hiera-tiny", + Size::Small => "facebook/sam2.1-hiera-small", + Size::BasePlus => "facebook/sam2.1-hiera-base-plus", + Size::Large => "facebook/sam2.1-hiera-large", + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sizes_name_their_official_repositories() { + assert_eq!(Size::Tiny.repo(), "facebook/sam2.1-hiera-tiny"); + assert!(Size::Large.repo().ends_with("large")); + } + + #[test] + fn errors_say_what_failed() { + assert!(Error::Input("empty frame".into()).to_string().contains("empty frame")); + assert_eq!(Error::Cancelled.to_string(), "SAM 2: cancelled"); + } +} diff --git a/xtask/src/main.rs b/xtask/src/main.rs index ad03cada7..401ed861e 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -50,6 +50,7 @@ const LAYERS: &[(&str, u8)] = &[ ("expr", 3), ("export", 3), ("gpu", 3), + ("sam2", 3), ("lottie", 3), ("format", 3), ("plugin", 3), From 92b3a34732c03f7a9cc44df39568c3f228245a3d Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:15:51 +0200 Subject: [PATCH 06/77] S2.11: Hiera backbone matches PyTorch block by block Weights loader with shape/finite validation (S2.10), patch embed, bicubic positional embedding, window/global attention, query pooling. Reference test (skips without weights): max abs error 2.6e-5 over all 12 blocks. Co-Authored-By: Claude Sonnet 5.5 --- Cargo.lock | 341 ++++++++++++++++++++++++++++++- crates/sam2/Cargo.toml | 5 + crates/sam2/src/hiera.rs | 359 +++++++++++++++++++++++++++++++++ crates/sam2/src/lib.rs | 13 ++ crates/sam2/src/nn.rs | 32 +++ crates/sam2/src/weights.rs | 97 +++++++++ crates/sam2/tests/reference.rs | 79 ++++++++ 7 files changed, 916 insertions(+), 10 deletions(-) create mode 100644 crates/sam2/src/hiera.rs create mode 100644 crates/sam2/src/nn.rs create mode 100644 crates/sam2/src/weights.rs create mode 100644 crates/sam2/tests/reference.rs diff --git a/Cargo.lock b/Cargo.lock index abf1f898b..afb66fd53 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -582,7 +582,7 @@ dependencies = [ "oneshot", "pastey", "portable-atomic", - "rand", + "rand 0.10.3", "regress", "rustc-hash 2.1.3", "ryu-js", @@ -723,6 +723,12 @@ dependencies = [ "syn 3.0.6", ] +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "byteorder-lite" version = "0.1.0" @@ -786,6 +792,45 @@ dependencies = [ "wayland-client", ] +[[package]] +name = "candle-core" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c15b675b80d994b2eadb20a4bbe434eabeb454eac3ee5e2b4cf6f147ee9be091" +dependencies = [ + "byteorder", + "float8", + "gemm", + "half", + "libm", + "memmap2", + "num-traits", + "num_cpus", + "rand 0.9.5", + "rand_distr", + "rayon", + "safetensors", + "thiserror 2.0.21", + "yoke", + "zip 7.4.0", +] + +[[package]] +name = "candle-nn" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3045fa9e7aef8567d209a27d56b692f60b96f4d0569f4c3011f8ca6715c65e03" +dependencies = [ + "candle-core", + "half", + "libc", + "num-traits", + "rayon", + "safetensors", + "serde", + "thiserror 2.0.21", +] + [[package]] name = "cast" version = "0.3.0" @@ -839,7 +884,7 @@ checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" dependencies = [ "cfg-if", "cpufeatures 0.3.1", - "rand_core", + "rand_core 0.10.1", ] [[package]] @@ -1141,7 +1186,7 @@ dependencies = [ "verifysign", "winreg", "winsafe", - "zip", + "zip 9.1.3", ] [[package]] @@ -1310,6 +1355,22 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" +[[package]] +name = "dyn-stack" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c4713e43e2886ba72b8271aa66c93d722116acf7a75555cce11dcde84388fe8" +dependencies = [ + "bytemuck", + "dyn-stack-macros", +] + +[[package]] +name = "dyn-stack-macros" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d926b4d407d372f141f93bb444696142c29d32962ccbd3531117cf3aa0bfa9" + [[package]] name = "dynify" version = "0.1.2" @@ -1767,7 +1828,11 @@ dependencies = [ [[package]] name = "effectcraft-sam2" -version = "0.9.2" +version = "0.18.0" +dependencies = [ + "candle-core", + "candle-nn", +] [[package]] name = "effectcraft-script" @@ -2065,6 +2130,18 @@ version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "66b7e2430c6dff6a955451e2cfc438f09cea1965a9d6f87f7e3b90decc014099" +[[package]] +name = "enum-as-inner" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1e6a265c649f3f5979b601d26f1d05ada116434c87741c9493cb56218f76cbc" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "enumflags2" version = "0.7.12" @@ -2590,6 +2667,18 @@ dependencies = [ "zlib-rs", ] +[[package]] +name = "float8" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "719a903cc23e4a89e87962c2a80fdb45cdaad0983a89bd150bb57b4c8571a7d5" +dependencies = [ + "half", + "num-traits", + "rand 0.9.5", + "rand_distr", +] + [[package]] name = "fmv-macos-events" version = "0.1.0" @@ -2732,6 +2821,125 @@ dependencies = [ "slab", ] +[[package]] +name = "gemm" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa0673db364b12263d103b68337a68fbecc541d6f6b61ba72fe438654709eacb" +dependencies = [ + "dyn-stack", + "gemm-c32", + "gemm-c64", + "gemm-common", + "gemm-f16", + "gemm-f32", + "gemm-f64", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-c32" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "086936dbdcb99e37aad81d320f98f670e53c1e55a98bee70573e83f95beb128c" +dependencies = [ + "dyn-stack", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-c64" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20c8aeeeec425959bda4d9827664029ba1501a90a0d1e6228e48bef741db3a3f" +dependencies = [ + "dyn-stack", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-common" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88027625910cc9b1085aaaa1c4bc46bb3a36aad323452b33c25b5e4e7c8e2a3e" +dependencies = [ + "bytemuck", + "dyn-stack", + "half", + "libm", + "num-complex", + "num-traits", + "once_cell", + "paste", + "pulp", + "raw-cpuid", + "rayon", + "seq-macro", + "sysctl", +] + +[[package]] +name = "gemm-f16" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3df7a55202e6cd6739d82ae3399c8e0c7e1402859b30e4cb780e61525d9486e" +dependencies = [ + "dyn-stack", + "gemm-common", + "gemm-f32", + "half", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "rayon", + "seq-macro", +] + +[[package]] +name = "gemm-f32" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02e0b8c9da1fbec6e3e3ab2ce6bc259ef18eb5f6f0d3e4edf54b75f9fd41a81c" +dependencies = [ + "dyn-stack", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-f64" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "056131e8f2a521bfab322f804ccd652520c79700d81209e9d9275bbdecaadc6a" +dependencies = [ + "dyn-stack", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -2785,7 +2993,7 @@ dependencies = [ "js-sys", "libc", "r-efi 6.0.0", - "rand_core", + "rand_core 0.10.1", "wasm-bindgen", ] @@ -2933,9 +3141,12 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "num-traits", + "rand 0.9.5", + "rand_distr", "zerocopy", ] @@ -2975,6 +3186,8 @@ dependencies = [ "allocator-api2", "equivalent", "foldhash 0.2.0", + "serde", + "serde_core", ] [[package]] @@ -3610,6 +3823,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" dependencies = [ "libc", + "stable_deref_trait", ] [[package]] @@ -3861,6 +4075,16 @@ dependencies = [ "libm", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "num_enum" version = "0.7.6" @@ -4300,7 +4524,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7d8fae84b431384b68627d0f9b3b1245fcf9f46f6c0e3dc902e9dce64edd1967" dependencies = [ "libc", - "windows-sys 0.45.0", + "windows-sys 0.61.2", ] [[package]] @@ -4569,6 +4793,15 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + [[package]] name = "presser" version = "0.3.1" @@ -4611,7 +4844,7 @@ dependencies = [ "chacha20", "core_detect", "num-traits", - "rand", + "rand 0.10.3", "rand_xorshift", "regex-syntax", "rusty-fork", @@ -4699,6 +4932,16 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha", + "rand_core 0.9.5", +] + [[package]] name = "rand" version = "0.10.3" @@ -4707,7 +4950,26 @@ checksum = "65c9fb96cbc91e3478eaae79a69fcd3f1ae4ad052e471fe6732fff548984b4af" dependencies = [ "chacha20", "getrandom 0.4.3", - "rand_core", + "rand_core 0.10.1", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", ] [[package]] @@ -4716,13 +4978,23 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" +[[package]] +name = "rand_distr" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" +dependencies = [ + "num-traits", + "rand 0.9.5", +] + [[package]] name = "rand_xorshift" version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "60aa6af80be32871323012e02e6e65f8a7cc7890931ae421d217ad8fe0df2ccf" dependencies = [ - "rand_core", + "rand_core 0.10.1", ] [[package]] @@ -5044,6 +5316,17 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04d056b875a9d2e6cb9a61d127afee9ac5999b9f87bcb32079d1318e505be714" +[[package]] +name = "safetensors" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "675656c1eabb620b921efea4f9199f97fc86e36dd6ffd1fbbe48d0f59a4987f5" +dependencies = [ + "hashbrown 0.16.1", + "serde", + "serde_json", +] + [[package]] name = "same-file" version = "1.0.6" @@ -5119,6 +5402,12 @@ version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cd0b0ec5f1c1ca621c432a25813d8d60c88abe6d3e08a3eb9cf37d97a0fe3d73" +[[package]] +name = "seq-macro" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" + [[package]] name = "serde" version = "1.0.229" @@ -5646,6 +5935,20 @@ dependencies = [ "web-sys", ] +[[package]] +name = "sysctl" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01198a2debb237c62b6826ec7081082d951f46dbb64b0e8c7649a452230d1dfc" +dependencies = [ + "bitflags 2.13.2", + "byteorder", + "enum-as-inner", + "libc", + "thiserror 1.0.69", + "walkdir", +] + [[package]] name = "tag_ptr" version = "0.1.0" @@ -5937,6 +6240,12 @@ dependencies = [ "rustc-hash 2.1.3", ] +[[package]] +name = "typed-path" +version = "0.12.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54ec2755506471f54baf5e4bb5499d8e510cc1caf271270e778a736bd2d7336c" + [[package]] name = "typed-path" version = "0.13.0" @@ -7475,6 +7784,18 @@ dependencies = [ "syn 3.0.6", ] +[[package]] +name = "zip" +version = "7.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc12baa6db2b15a140161ce53d72209dacea594230798c24774139b54ecaa980" +dependencies = [ + "crc32fast", + "indexmap", + "memchr", + "typed-path 0.12.4", +] + [[package]] name = "zip" version = "9.1.3" @@ -7487,7 +7808,7 @@ dependencies = [ "indexmap", "memchr", "nix", - "typed-path", + "typed-path 0.13.0", "windows-sys 0.61.2", "zopfli", ] diff --git a/crates/sam2/Cargo.toml b/crates/sam2/Cargo.toml index 24c317bc1..f32a4e4e6 100644 --- a/crates/sam2/Cargo.toml +++ b/crates/sam2/Cargo.toml @@ -10,3 +10,8 @@ rust-version.workspace = true workspace = true [dependencies] + +# Native only (D1): candle is not built for wasm32, where this crate is empty. +[target.'cfg(not(target_arch = "wasm32"))'.dependencies] +candle-core = "=0.9.2" +candle-nn = "=0.9.2" diff --git a/crates/sam2/src/hiera.rs b/crates/sam2/src/hiera.rs new file mode 100644 index 000000000..5091d331b --- /dev/null +++ b/crates/sam2/src/hiera.rs @@ -0,0 +1,359 @@ +// Modified work: ported from Hugging Face `transformers` `models/sam2/modeling_sam2.py` +// (Sam2PatchEmbeddings, Sam2MultiScaleAttention, Sam2MultiScaleBlock, Sam2HieraDetModel; +// Apache-2.0, Copyright The HuggingFace Inc. team) to Rust on candle. See NOTICE. +//! The Hiera backbone: patch embedding, windowed positional embedding, and blocks with window or +//! global attention and query pooling at the stage changes. Activations are NHWC like the reference. + +use candle_core::{D, Tensor}; +use candle_nn::{LayerNorm, Linear, Module, ops::softmax}; + +use crate::nn::{Mlp, layer_norm, linear}; +use crate::weights::Weights; +use crate::{Error, Result}; + +/// Hiera hyper-parameters (the released SAM 2.1 sizes differ only here). +#[derive(Clone, Debug, PartialEq)] +pub struct HieraConfig { + pub embed_dims: Vec, + pub blocks: Vec, + pub heads: Vec, + pub windows: Vec, + pub global_blocks: Vec, + pub query_stride: usize, + pub query_pool_stages: usize, + pub mlp_ratio: usize, + pub eps: f64, + pub pos_background: usize, +} + +impl HieraConfig { + /// `facebook/sam2.1-hiera-tiny`. + pub fn tiny() -> HieraConfig { + HieraConfig { + embed_dims: vec![96, 192, 384, 768], + blocks: vec![1, 2, 7, 2], + heads: vec![1, 2, 4, 8], + windows: vec![8, 4, 14, 7], + global_blocks: vec![5, 7, 9], + query_stride: 2, + query_pool_stages: 3, + mlp_ratio: 4, + eps: 1e-6, + pos_background: 7, + } + } +} + +fn at(v: &[T], i: usize, what: &str) -> Result { + v.get(i).copied().ok_or_else(|| Error::Weights(format!("Hiera config has no {what} for stage {i}"))) +} + +/// Max pool a `[B, H, W, C]` tensor by `stride` (kernel = stride, no padding, floor). +fn do_pool(x: &Tensor, stride: usize) -> Result { + let x = x.permute((0, 3, 1, 2))?.contiguous()?.max_pool2d_with_stride(stride, stride)?; + Ok(x.permute((0, 2, 3, 1))?.contiguous()?) +} + +/// Split `[B, H, W, C]` into `[B * nH * nW, ws, ws, C]` windows, zero padding the bottom/right. +fn window_partition(x: &Tensor, ws: usize) -> Result<(Tensor, (usize, usize))> { + let (b, h, w, c) = x.dims4()?; + if ws == 0 { + return Err(Error::Input("window size 0".into())); + } + let (ph, pw) = ((ws - h % ws) % ws, (ws - w % ws) % ws); + let x = x.pad_with_zeros(2, 0, pw)?.pad_with_zeros(1, 0, ph)?; + let (nh, nw) = ((h + ph) / ws, (w + pw) / ws); + let x = x.reshape((b, nh, ws, nw, ws, c))?.permute((0, 1, 3, 2, 4, 5))?.contiguous()?.reshape((b * nh * nw, ws, ws, c))?; + Ok((x, (h + ph, w + pw))) +} + +fn window_unpartition(win: &Tensor, ws: usize, padded: (usize, usize), orig: (usize, usize)) -> Result { + if ws == 0 { + return Err(Error::Input("window size 0".into())); + } + let (nh, nw) = (padded.0 / ws, padded.1 / ws); + let (n, _, _, c) = win.dims4()?; + let b = n / (nh * nw).max(1); + let x = win.reshape((b, nh, nw, ws, ws, c))?.permute((0, 1, 3, 2, 4, 5))?.contiguous()?.reshape((b, padded.0, padded.1, c))?; + Ok(x.narrow(1, 0, orig.0)?.narrow(2, 0, orig.1)?.contiguous()?) +} + +struct Attention { + qkv: Linear, + proj: Linear, + heads: usize, + scale: f64, + query_stride: Option, +} + +impl Attention { + fn forward(&self, x: &Tensor) -> Result { + let (b, h, w, _) = x.dims4()?; + let qkv = self.qkv.forward(x)?.reshape((b, h * w, 3, self.heads, ()))?; + let part = |i: usize| -> Result { Ok(qkv.narrow(2, i, 1)?.squeeze(2)?) }; + let (mut q, k, v) = (part(0)?, part(1)?, part(2)?); + let (mut oh, mut ow) = (h, w); + if let Some(s) = self.query_stride { + let pooled = do_pool(&q.reshape((b, h, w, ()))?, s)?; + (oh, ow) = (pooled.dim(1)?, pooled.dim(2)?); + q = pooled.reshape((b, oh * ow, self.heads, ()))?; + } + let q = q.transpose(1, 2)?.contiguous()?; + let k = k.transpose(1, 2)?.contiguous()?; + let v = v.transpose(1, 2)?.contiguous()?; + let att = (q.matmul(&k.transpose(2, 3)?.contiguous()?)? * self.scale)?; + let att = softmax(&att, D::Minus1)?; + let out = att.matmul(&v)?.transpose(1, 2)?.contiguous()?.reshape((b, oh, ow, ()))?; + Ok(self.proj.forward(&out)?) + } +} + +/// One Hiera block: pre-norm attention (window or global, optionally pooling the queries) and MLP. +pub struct Block { + norm1: LayerNorm, + attn: Attention, + norm2: LayerNorm, + mlp: Mlp, + proj: Option, + window: usize, + query_stride: Option, +} + +impl Block { + fn load(w: &Weights, cfg: &HieraConfig, stage: usize, block_in_stage: usize, total: usize) -> Result { + let prefix = format!("vision_encoder.backbone.blocks.{total}"); + let first = stage > 0 && block_in_stage == 0; + let dim = at(&cfg.embed_dims, if first { stage - 1 } else { stage }, "embed dim")?; + let dim_out = at(&cfg.embed_dims, stage, "embed dim")?; + let window = if first { at(&cfg.windows, stage - 1, "window")? } else { at(&cfg.windows, stage, "window")? }; + let window = if cfg.global_blocks.contains(&total) { 0 } else { window }; + let query_stride = (stage > 0 && stage <= cfg.query_pool_stages && block_in_stage == 0).then_some(cfg.query_stride); + let heads = at(&cfg.heads, stage, "heads")?; + Ok(Block { + norm1: layer_norm(w, &format!("{prefix}.layer_norm1"), dim, cfg.eps)?, + attn: Attention { + qkv: linear(w, &format!("{prefix}.attn.qkv"), dim, dim_out * 3)?, + proj: linear(w, &format!("{prefix}.attn.proj"), dim_out, dim_out)?, + heads, + scale: ((dim_out / heads.max(1)) as f64).powf(-0.5), + query_stride, + }, + norm2: layer_norm(w, &format!("{prefix}.layer_norm2"), dim_out, cfg.eps)?, + mlp: Mlp::load(w, &format!("{prefix}.mlp"), dim_out, dim_out * cfg.mlp_ratio)?, + proj: (dim != dim_out).then(|| linear(w, &format!("{prefix}.proj"), dim, dim_out)).transpose()?, + window, + query_stride, + }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + let h = self.norm1.forward(x)?; + let residual = match (&self.proj, self.query_stride) { + (Some(p), Some(s)) => do_pool(&p.forward(&h)?, s)?, + (Some(p), None) => p.forward(&h)?, + (None, _) => x.clone(), + }; + let (hh, ww) = (h.dim(1)?, h.dim(2)?); + let mut ws = self.window; + let mut padded = (hh, ww); + let h = if ws > 0 { + let (win, pad) = window_partition(&h, ws)?; + padded = pad; + win + } else { + h + }; + let mut h = self.attn.forward(&h)?; + let mut orig = (hh, ww); + if let Some(s) = self.query_stride { + ws /= s.max(1); + orig = (residual.dim(1)?, residual.dim(2)?); + if ws > 0 { + padded = (orig.0 + (ws - orig.0 % ws) % ws, orig.1 + (ws - orig.1 % ws) % ws); + } + } + if ws > 0 { + h = window_unpartition(&h, ws, padded, orig)?; + } + let h = (residual + h)?; + let m = self.mlp.forward(&self.norm2.forward(&h)?)?; + Ok((h + m)?) + } +} + +/// PyTorch `F.interpolate(mode="bicubic", align_corners=False)` of `[C, in_h, in_w]` to `[C, out_h, out_w]`. +fn bicubic_resize(src: &[f32], c: usize, in_hw: (usize, usize), out_hw: (usize, usize)) -> Vec { + const A: f32 = -0.75; + let near = |x: f32| ((A + 2.0) * x - (A + 3.0)) * x * x + 1.0; + let far = |x: f32| ((A * x - 5.0 * A) * x + 8.0 * A) * x - 4.0 * A; + let axis = |n_in: usize, n_out: usize| -> Vec<([usize; 4], [f32; 4])> { + let scale = n_in as f32 / n_out.max(1) as f32; + (0..n_out) + .map(|o| { + let s = (o as f32 + 0.5) * scale - 0.5; + let i0 = s.floor(); + let t = s - i0; + let w = [far(t + 1.0), near(t), near(1.0 - t), far(2.0 - t)]; + let last = n_in as i64 - 1; + let idx = [-1i64, 0, 1, 2].map(|d| (i0 as i64 + d).clamp(0, last.max(0)) as usize); + (idx, w) + }) + .collect() + }; + let (ys, xs) = (axis(in_hw.0, out_hw.0), axis(in_hw.1, out_hw.1)); + let mut out = vec![0f32; c * out_hw.0 * out_hw.1]; + for ch in 0..c { + let plane = src.get(ch * in_hw.0 * in_hw.1..(ch + 1) * in_hw.0 * in_hw.1).unwrap_or(&[]); + for (oy, (yi, yw)) in ys.iter().enumerate() { + for (ox, (xi, xw)) in xs.iter().enumerate() { + let mut acc = 0f32; + for (a, wy) in yi.iter().zip(yw) { + for (b, wx) in xi.iter().zip(xw) { + acc += plane.get(a * in_hw.1 + b).copied().unwrap_or(0.0) * wy * wx; + } + } + if let Some(o) = out.get_mut((ch * out_hw.0 + oy) * out_hw.1 + ox) { + *o = acc; + } + } + } + } + out +} + +pub struct Hiera { + cfg: HieraConfig, + patch_w: Tensor, + patch_b: Tensor, + pos_embed: Vec, + pos_embed_window: Vec, + window0: usize, + hidden: usize, + pub blocks: Vec, + stage_ends: Vec, +} + +impl Hiera { + pub fn load(w: &Weights, cfg: &HieraConfig) -> Result { + let hidden = at(&cfg.embed_dims, 0, "embed dim")?; + let window0 = at(&cfg.windows, 0, "window")?; + let bg = cfg.pos_background; + let mut blocks = vec![]; + let mut stage_ends = vec![]; + for (stage, n) in cfg.blocks.iter().enumerate() { + for b in 0..*n { + blocks.push(Block::load(w, cfg, stage, b, blocks.len())?); + } + stage_ends.push(blocks.len().saturating_sub(1)); + } + let flat = |t: Tensor| -> Result> { Ok(t.flatten_all()?.to_vec1::()?) }; + Ok(Hiera { + cfg: cfg.clone(), + patch_w: w.get("vision_encoder.backbone.patch_embed.projection.weight", &[hidden, 3, 7, 7])?, + patch_b: w.get("vision_encoder.backbone.patch_embed.projection.bias", &[hidden])?, + pos_embed: flat(w.get("vision_encoder.backbone.pos_embed", &[1, hidden, bg, bg])?)?, + pos_embed_window: flat(w.get("vision_encoder.backbone.pos_embed_window", &[1, hidden, window0, window0])?)?, + window0, + hidden, + blocks, + stage_ends, + }) + } + + /// Patch embedding of `[B, 3, H, W]` pixels: `[B, H/4, W/4, C]` (before the positional embedding). + pub fn patch_embed(&self, pixels: &Tensor) -> Result { + let x = pixels.conv2d(&self.patch_w, 3, 4, 1, 1)?; + let x = x.broadcast_add(&self.patch_b.reshape((1, self.hidden, 1, 1))?)?; + Ok(x.permute((0, 2, 3, 1))?.contiguous()?) + } + + /// Background embedding resized to the grid plus the tiled window embedding, `[1, H, W, C]`. + fn pos(&self, h: usize, w: usize) -> Result { + let bg = self.cfg.pos_background; + let big = bicubic_resize(&self.pos_embed, self.hidden, (bg, bg), (h, w)); + let ws = self.window0.max(1); + let mut out = vec![0f32; h * w * self.hidden]; + for c in 0..self.hidden { + for y in 0..h { + for x in 0..w { + let tile = self.pos_embed_window.get((c * ws + y % ws) * ws + x % ws).copied().unwrap_or(0.0); + let base = big.get((c * h + y) * w + x).copied().unwrap_or(0.0); + if let Some(o) = out.get_mut((y * w + x) * self.hidden + c) { + *o = base + tile; + } + } + } + } + Ok(Tensor::from_vec(out, (1, h, w, self.hidden), &candle_core::Device::Cpu)?) + } + + /// Run all blocks. Returns the output of the last block of each stage (NHWC). + pub fn forward(&self, pixels: &Tensor) -> Result> { + let (_, _, ph, pw) = pixels.dims4()?; + if ph == 0 || pw == 0 || ph % 4 != 0 || pw % 4 != 0 { + return Err(Error::Input(format!("image size {pw}x{ph} must be a positive multiple of 4"))); + } + let mut x = self.patch_embed(pixels)?; + let (h, w) = (x.dim(1)?, x.dim(2)?); + if h % self.window0 != 0 || w % self.window0 != 0 { + return Err(Error::Input(format!("patch grid {w}x{h} must be a multiple of {}", self.window0))); + } + x = x.broadcast_add(&self.pos(h, w)?)?; + let mut stages = vec![]; + for (i, b) in self.blocks.iter().enumerate() { + x = b.forward(&x)?; + if self.stage_ends.contains(&i) { + stages.push(x.clone()); + } + } + Ok(stages) + } + + /// Run the blocks, returning every block's output (for the stage-by-stage reference test). + pub fn forward_all(&self, pixels: &Tensor) -> Result<(Tensor, Vec)> { + let emb = self.patch_embed(pixels)?; + let (h, w) = (emb.dim(1)?, emb.dim(2)?); + let mut x = emb.broadcast_add(&self.pos(h, w)?)?; + let mut outs = vec![]; + for b in &self.blocks { + x = b.forward(&x)?; + outs.push(x.clone()); + } + Ok((emb, outs)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use candle_core::Device; + + #[test] + fn bicubic_of_a_constant_is_that_constant() { + let out = bicubic_resize(&[2.5; 49], 1, (7, 7), (32, 32)); + assert!(out.iter().all(|v| (v - 2.5).abs() < 1e-5)); + } + + #[test] + fn bicubic_identity_size_is_identity() { + let src: Vec = (0..49).map(|i| (i as f32 * 0.37).sin()).collect(); + let out = bicubic_resize(&src, 1, (7, 7), (7, 7)); + assert!(src.iter().zip(&out).all(|(a, b)| (a - b).abs() < 1e-5)); + } + + #[test] + fn windows_round_trip_with_padding() { + let x = Tensor::arange(0f32, 60.0, &Device::Cpu).unwrap().reshape((1, 5, 6, 2)).unwrap(); + let (win, pad) = window_partition(&x, 4).unwrap(); + assert_eq!((win.dims(), pad), (&[4, 4, 4, 2][..], (8, 8))); + let back = window_unpartition(&win, 4, pad, (5, 6)).unwrap(); + let d = (back - x).unwrap().abs().unwrap().sum_all().unwrap().to_scalar::().unwrap(); + assert_eq!(d, 0.0); + } + + #[test] + fn zero_window_is_an_error_not_a_panic() { + let x = Tensor::zeros((1, 4, 4, 2), candle_core::DType::F32, &Device::Cpu).unwrap(); + assert!(window_partition(&x, 0).is_err()); + } +} diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index d41de59fa..246739e45 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -11,6 +11,10 @@ use std::fmt; +pub mod hiera; +pub mod nn; +pub mod weights; + /// Everything that can go wrong in this crate. Never a panic. #[derive(Debug, Clone, PartialEq)] pub enum Error { @@ -22,6 +26,8 @@ pub enum Error { Cancelled, /// A time or memory budget was exceeded. Budget(String), + /// A tensor operation failed (shape mismatch, out of memory, device error). + Compute(String), } impl fmt::Display for Error { @@ -31,12 +37,19 @@ impl fmt::Display for Error { Error::Input(m) => write!(f, "SAM 2 input: {m}"), Error::Cancelled => write!(f, "SAM 2: cancelled"), Error::Budget(m) => write!(f, "SAM 2 budget: {m}"), + Error::Compute(m) => write!(f, "SAM 2 compute: {m}"), } } } impl std::error::Error for Error {} +impl From for Error { + fn from(e: candle_core::Error) -> Error { + Error::Compute(e.to_string()) + } +} + pub type Result = std::result::Result; /// The released SAM 2.1 sizes. Only Tiny is wired up first; the others are configuration. diff --git a/crates/sam2/src/nn.rs b/crates/sam2/src/nn.rs new file mode 100644 index 000000000..ef334432a --- /dev/null +++ b/crates/sam2/src/nn.rs @@ -0,0 +1,32 @@ +// Modified work: small layers used by the ported SAM 2 modules (Apache-2.0, see NOTICE). +//! Linear, layer norm and a few tensor helpers on candle. + +use candle_core::Tensor; +use candle_nn::{LayerNorm, Linear, Module}; + +use crate::Result; +use crate::weights::Weights; + +pub fn linear(w: &Weights, prefix: &str, input: usize, output: usize) -> Result { + Ok(Linear::new(w.get(&format!("{prefix}.weight"), &[output, input])?, Some(w.get(&format!("{prefix}.bias"), &[output])?))) +} + +pub fn layer_norm(w: &Weights, prefix: &str, dim: usize, eps: f64) -> Result { + Ok(LayerNorm::new(w.get(&format!("{prefix}.weight"), &[dim])?, w.get(&format!("{prefix}.bias"), &[dim])?, eps)) +} + +/// Two-layer MLP with exact (erf) GELU, as `Sam2FeedForward` with `num_layers = 2`. +pub struct Mlp { + pub proj_in: Linear, + pub proj_out: Linear, +} + +impl Mlp { + pub fn load(w: &Weights, prefix: &str, dim: usize, hidden: usize) -> Result { + Ok(Mlp { proj_in: linear(w, &format!("{prefix}.proj_in"), dim, hidden)?, proj_out: linear(w, &format!("{prefix}.proj_out"), hidden, dim)? }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + Ok(self.proj_out.forward(&self.proj_in.forward(x)?.gelu_erf()?)?) + } +} diff --git a/crates/sam2/src/weights.rs b/crates/sam2/src/weights.rs new file mode 100644 index 000000000..c8b79cc70 --- /dev/null +++ b/crates/sam2/src/weights.rs @@ -0,0 +1,97 @@ +// Modified work: loads the safetensors checkpoint of Hugging Face `transformers` SAM 2 +// (Apache-2.0, Copyright The HuggingFace Inc. team). See NOTICE. +//! Checkpoint loading with validation: every tensor is f32 and finite, every lookup checks its +//! shape, so a truncated, mislabelled or poisoned file is an error, not a crash or garbage. + +use std::collections::HashMap; +use std::path::Path; + +use candle_core::{DType, Device, Tensor}; + +use crate::{Error, Result}; + +/// Named tensors of one checkpoint. +pub struct Weights { + map: HashMap, +} + +impl Weights { + /// Read a `.safetensors` file; rejects non-finite values and unsupported dtypes. + pub fn open(path: &Path) -> Result { + let map = candle_core::safetensors::load(path, &Device::Cpu).map_err(|e| Error::Weights(format!("{}: {e}", path.display())))?; + 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()); + for (name, t) in map { + let t = match t.dtype() { + DType::F32 => t, + DType::F16 | DType::BF16 | DType::F64 => t.to_dtype(DType::F32)?, + other => return Err(Error::Weights(format!("`{name}` has unsupported dtype {other:?}"))), + }; + let flat = t.flatten_all()?.to_vec1::()?; + if flat.iter().any(|v| !v.is_finite()) { + return Err(Error::Weights(format!("`{name}` contains NaN or infinity"))); + } + out.insert(name, t); + } + Ok(Weights { map: out }) + } + + pub fn len(&self) -> usize { + self.map.len() + } + + pub fn is_empty(&self) -> bool { + self.map.is_empty() + } + + /// The tensor `name`, which must have exactly `shape`. + pub fn get(&self, name: &str, shape: &[usize]) -> Result { + let t = self.map.get(name).ok_or_else(|| Error::Weights(format!("missing tensor `{name}`")))?; + if t.dims() != shape { + return Err(Error::Weights(format!("`{name}` has shape {:?}, expected {shape:?}", t.dims()))); + } + Ok(t.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn one(name: &str, t: Tensor) -> HashMap { + HashMap::from([(name.to_string(), t)]) + } + + #[test] + fn rejects_missing_wrong_shape_and_non_finite() { + let ok = Tensor::zeros((2, 3), DType::F32, &Device::Cpu).unwrap(); + let w = Weights::from_map(one("a", ok)).unwrap(); + assert!(w.get("a", &[2, 3]).is_ok()); + assert!(matches!(w.get("a", &[3, 2]), Err(Error::Weights(_)))); + assert!(matches!(w.get("b", &[1]), Err(Error::Weights(_)))); + let bad = Tensor::new(&[1.0f32, f32::NAN], &Device::Cpu).unwrap(); + assert!(matches!(Weights::from_map(one("n", bad)), Err(Error::Weights(_)))); + let inf = Tensor::new(&[f32::INFINITY], &Device::Cpu).unwrap(); + assert!(Weights::from_map(one("i", inf)).is_err()); + let ints = Tensor::new(&[1u32], &Device::Cpu).unwrap(); + assert!(Weights::from_map(one("u", ints)).is_err()); + } + + #[test] + fn a_truncated_file_is_an_error() { + let dir = std::env::temp_dir().join(format!("sam2-weights-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let p = dir.join("w.safetensors"); + let t = Tensor::zeros((4, 4), DType::F32, &Device::Cpu).unwrap(); + candle_core::safetensors::save(&one("x", t), &p).unwrap(); + let bytes = std::fs::read(&p).unwrap(); + std::fs::write(&p, &bytes[..bytes.len() - 8]).unwrap(); + assert!(matches!(Weights::open(&p), Err(Error::Weights(_)))); + assert!(matches!(Weights::open(&dir.join("absent")), Err(Error::Weights(_)))); + std::fs::remove_dir_all(&dir).ok(); + } +} diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs new file mode 100644 index 000000000..e5b7ae596 --- /dev/null +++ b/crates/sam2/tests/reference.rs @@ -0,0 +1,79 @@ +//! Stage-by-stage comparison with the PyTorch / transformers reference (`tools/sam2_reference.py`). +//! +//! Needs the `facebook/sam2.1-hiera-tiny` weights (default `/sam2-tiny`, or the folder in +//! `EFFECTCRAFT_SAM2_DIR`) and the fixture tensors (`/plan/sam2-ref/image.safetensors`, or +//! `EFFECTCRAFT_SAM2_REF`). Without them every test here skips, so CI stays green. +#![cfg(not(target_arch = "wasm32"))] + +use std::collections::HashMap; +use std::path::PathBuf; + +use candle_core::{Device, Tensor}; +use effectcraft_sam2::hiera::{Hiera, HieraConfig}; +use effectcraft_sam2::weights::Weights; + +fn repo() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../..") +} + +fn weights_file() -> PathBuf { + std::env::var_os("EFFECTCRAFT_SAM2_DIR").map_or_else(|| repo().join("sam2-tiny"), PathBuf::from).join("model.safetensors") +} + +fn fixture_file() -> PathBuf { + std::env::var_os("EFFECTCRAFT_SAM2_REF").map_or_else(|| repo().join("plan/sam2-ref/image.safetensors"), PathBuf::from) +} + +/// The weights and the reference tensors, or `None` (the test then skips). +fn load() -> Option<(Weights, HashMap)> { + let (wf, rf) = (weights_file(), fixture_file()); + if !wf.exists() || !rf.exists() { + eprintln!("skipped: SAM 2 weights or reference tensors not found ({} / {})", wf.display(), rf.display()); + return None; + } + let w = Weights::open(&wf).unwrap(); + let r = candle_core::safetensors::load(&rf, &Device::Cpu).unwrap(); + Some((w, r)) +} + +fn max_abs(a: &Tensor, b: &Tensor) -> f32 { + assert_eq!(a.dims(), b.dims(), "shape mismatch"); + (a - b).unwrap().abs().unwrap().flatten_all().unwrap().max(0).unwrap().to_scalar::().unwrap() +} + +/// f32 activations differ from PyTorch by accumulation order only; the Hiera blocks stay well +/// inside this bound (measured max 2.6e-5, 2026-10-15). +const HIERA_TOL: f32 = 1e-4; + +#[test] +fn hiera_matches_pytorch_block_by_block() { + let Some((w, r)) = load() else { return }; + let hiera = Hiera::load(&w, &HieraConfig::tiny()).unwrap(); + let pixels = r.get("input.pixel_values").unwrap(); + let (emb, blocks) = hiera.forward_all(pixels).unwrap(); + let e = max_abs(&emb, r.get("vision_encoder.backbone.patch_embed.0").unwrap()); + eprintln!("patch_embed max abs error {e:e}"); + assert!(e < 1e-4, "patch embed error {e}"); + assert_eq!(blocks.len(), 12); + for (i, b) in blocks.iter().enumerate() { + let want = r.get(&format!("vision_encoder.backbone.blocks.{i}.0")).unwrap(); + let e = max_abs(b, want); + eprintln!("block {i}: max abs error {e:e}"); + assert!(e < HIERA_TOL, "block {i} error {e} (tolerance {HIERA_TOL})"); + } +} + +#[test] +fn hiera_is_deterministic_and_rejects_bad_sizes() { + let Some((w, r)) = load() else { return }; + let hiera = Hiera::load(&w, &HieraConfig::tiny()).unwrap(); + let pixels = r.get("input.pixel_values").unwrap(); + let a = hiera.forward(pixels).unwrap(); + let b = hiera.forward(pixels).unwrap(); + assert_eq!(a.len(), 4); + for (x, y) in a.iter().zip(&b) { + assert_eq!(max_abs(x, y), 0.0); + } + let odd = pixels.narrow(2, 0, 1022).unwrap(); + assert!(hiera.forward(&odd).is_err()); +} From d2291926f018d7eb078e2bc7140a39dd46be35bd Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:16:53 +0200 Subject: [PATCH 07/77] S2.12: FPN neck and vision encoder match PyTorch Max abs error 9.7e-6 on the features, 6e-8 on the sine position encodings. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/src/lib.rs | 1 + crates/sam2/src/neck.rs | 141 +++++++++++++++++++++++++++++++++ crates/sam2/tests/reference.rs | 30 +++++++ 3 files changed, 172 insertions(+) create mode 100644 crates/sam2/src/neck.rs diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 246739e45..6515eff7b 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -12,6 +12,7 @@ use std::fmt; pub mod hiera; +pub mod neck; pub mod nn; pub mod weights; diff --git a/crates/sam2/src/neck.rs b/crates/sam2/src/neck.rs new file mode 100644 index 000000000..0a82ec7a9 --- /dev/null +++ b/crates/sam2/src/neck.rs @@ -0,0 +1,141 @@ +// Modified work: ported from Hugging Face `transformers` `models/sam2/modeling_sam2.py` +// (Sam2VisionNeck, Sam2SinePositionEmbedding, Sam2VisionModel; Apache-2.0, Copyright The +// HuggingFace Inc. team) to Rust on candle. See NOTICE. +//! The feature pyramid neck on top of Hiera, and the vision encoder that joins them. + +use candle_core::{Device, Tensor}; + +use crate::hiera::{Hiera, HieraConfig}; +use crate::weights::Weights; +use crate::{Error, Result}; + +/// Sine position encoding `[1, 2 * num_pos_features, h, w]`, normalised (`Sam2SinePositionEmbedding`). +pub fn sine_position(num_pos_features: usize, h: usize, w: usize) -> Result { + let eps = 1e-6f32; + let scale = 2.0 * std::f32::consts::PI; + let dim_t: Vec = (0..num_pos_features).map(|i| 10000f32.powf(2.0 * (i / 2) as f32 / num_pos_features as f32)).collect(); + let axis = |n: usize, k: usize| -> Vec { + // value of one coordinate at index k of an axis of length n, over the feature dimension + let e = (k + 1) as f32 / (n as f32 + eps) * scale; + dim_t.iter().enumerate().map(|(i, d)| if i % 2 == 0 { (e / d).sin() } else { (e / d).cos() }).collect() + }; + let (ys, xs): (Vec<_>, Vec<_>) = ((0..h).map(|k| axis(h, k)).collect(), (0..w).map(|k| axis(w, k)).collect()); + let mut out = vec![0f32; 2 * num_pos_features * h * w]; + for y in 0..h { + for x in 0..w { + for f in 0..num_pos_features { + let (py, px) = (ys.get(y).and_then(|v| v.get(f)), xs.get(x).and_then(|v| v.get(f))); + if let (Some(py), Some(px)) = (py, px) { + if let Some(o) = out.get_mut((f * h + y) * w + x) { + *o = *py; + } + if let Some(o) = out.get_mut(((num_pos_features + f) * h + y) * w + x) { + *o = *px; + } + } + } + } + } + Ok(Tensor::from_vec(out, (1, 2 * num_pos_features, h, w), &Device::Cpu)?) +} + +struct Conv1x1 { + w: Tensor, + b: Tensor, + out: usize, +} + +impl Conv1x1 { + fn forward(&self, x: &Tensor) -> Result { + Ok(x.conv2d(&self.w, 0, 1, 1, 1)?.broadcast_add(&self.b.reshape((1, self.out, 1, 1))?)?) + } +} + +/// The FPN neck: a 1x1 conv per stage, top-down additions on the two lowest-resolution levels. +pub struct Neck { + convs: Vec, + top_down: Vec, + pos_features: usize, +} + +/// What the neck produces, both in the neck's own order (lowest resolution first, all four +/// levels) as `[B, C, H, W]`. +pub struct NeckOutput { + pub features: Vec, + pub positions: Vec, +} + +impl Neck { + pub fn load(w: &Weights, channels: &[usize], fpn: usize, top_down: &[usize]) -> Result { + let convs = channels + .iter() + .enumerate() + .map(|(j, c)| { + Ok(Conv1x1 { + w: w.get(&format!("vision_encoder.neck.convs.{j}.weight"), &[fpn, *c, 1, 1])?, + b: w.get(&format!("vision_encoder.neck.convs.{j}.bias"), &[fpn])?, + out: fpn, + }) + }) + .collect::>>()?; + Ok(Neck { convs, top_down: top_down.to_vec(), pos_features: fpn / 2 }) + } + + /// `stages`: Hiera stage outputs, NHWC, highest resolution first. + pub fn forward(&self, stages: &[Tensor]) -> Result { + let n = self.convs.len().checked_sub(1).ok_or_else(|| Error::Weights("neck has no convolutions".into()))?; + if stages.len() != self.convs.len() { + return Err(Error::Input(format!("neck expects {} stages, got {}", self.convs.len(), stages.len()))); + } + let (mut features, mut positions) = (vec![], vec![]); + let mut prev: Option = None; + for i in (0..=n).rev() { + let stage = stages.get(i).ok_or_else(|| Error::Input("missing stage".into()))?; + let conv = self.convs.get(n - i).ok_or_else(|| Error::Weights("missing neck convolution".into()))?; + let lateral = conv.forward(&stage.permute((0, 3, 1, 2))?.contiguous()?)?; + let cur = match (&prev, self.top_down.contains(&i) && i != n) { + (Some(p), true) => { + let (h, w) = (p.dim(2)? * 2, p.dim(3)? * 2); + (lateral + p.upsample_nearest2d(h, w)?)? + } + _ => lateral, + }; + positions.push(sine_position(self.pos_features, cur.dim(2)?, cur.dim(3)?)?); + features.push(cur.clone()); + prev = Some(cur); + } + Ok(NeckOutput { features, positions }) + } +} + +/// Hiera plus neck: `Sam2VisionModel`. +pub struct VisionEncoder { + pub hiera: Hiera, + pub neck: Neck, + levels: usize, +} + +/// The multi-scale image features the decoder and memory modules use, highest resolution first +/// (256², 128², 64² for a 1024 input), `[B, C, H, W]`. +pub struct ImageFeatures { + pub features: Vec, + pub positions: Vec, +} + +impl VisionEncoder { + pub fn load_tiny(w: &Weights) -> Result { + let cfg = HieraConfig::tiny(); + Ok(VisionEncoder { hiera: Hiera::load(w, &cfg)?, neck: Neck::load(w, &[768, 384, 192, 96], 256, &[2, 3])?, levels: 3 }) + } + + pub fn forward(&self, pixels: &Tensor) -> Result { + let stages = self.hiera.forward(pixels)?; + let out = self.neck.forward(&stages)?; + // keep the last `levels` of the neck order and reverse: high to low resolution + let keep = |v: Vec| -> Vec { + let skip = v.len().saturating_sub(self.levels); + v.into_iter().skip(skip).rev().collect() + }; + Ok(ImageFeatures { features: keep(out.features), positions: keep(out.positions) }) + } +} diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index e5b7ae596..bae7d42e3 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -77,3 +77,33 @@ fn hiera_is_deterministic_and_rejects_bad_sizes() { let odd = pixels.narrow(2, 0, 1022).unwrap(); assert!(hiera.forward(&odd).is_err()); } + +/// The neck adds 1x1 convolutions and nearest-neighbour sums to the Hiera output; the sine +/// position encodings are exact functions of the grid size. +const NECK_TOL: f32 = 1e-4; + +#[test] +fn neck_and_vision_encoder_match_pytorch() { + let Some((w, r)) = load() else { return }; + let enc = effectcraft_sam2::neck::VisionEncoder::load_tiny(&w).unwrap(); + let pixels = r.get("input.pixel_values").unwrap(); + let stages = enc.hiera.forward(pixels).unwrap(); + let out = enc.neck.forward(&stages).unwrap(); + assert_eq!((out.features.len(), out.positions.len()), (4, 4)); + for (i, f) in out.features.iter().enumerate() { + let e = max_abs(f, r.get(&format!("vision_encoder.neck.{i}")).unwrap()); + eprintln!("neck feature {i}: max abs error {e:e}"); + assert!(e < NECK_TOL, "neck feature {i} error {e}"); + } + for (i, p) in out.positions.iter().enumerate() { + let e = max_abs(p, r.get(&format!("vision_encoder.neck.{}", i + 4)).unwrap()); + eprintln!("neck position {i}: max abs error {e:e}"); + assert!(e < NECK_TOL, "neck position {i} error {e}"); + } + // The vision encoder keeps the three highest-resolution levels, highest first. + let feats = enc.forward(pixels).unwrap(); + let dims: Vec = feats.features.iter().map(|f| f.dim(2).unwrap()).collect(); + assert_eq!(dims, vec![256, 128, 64]); + let e = max_abs(feats.features.first().unwrap(), r.get("vision_encoder.neck.3").unwrap()); + assert!(e < NECK_TOL); +} From 8a6ad6f5e9d5cbc104820919a4712a807cbfe192 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:19:14 +0200 Subject: [PATCH 08/77] S2.13: prompt encoder and two-way mask decoder match PyTorch A click on the synthetic image: mask logits within 1.1e-4 (range 11.6), IoU head 1e-6, object score 2e-5, prompt embeddings 1e-7. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/src/decoder.rs | 258 +++++++++++++++++++++++++++++++++ crates/sam2/src/lib.rs | 3 + crates/sam2/src/model.rs | 77 ++++++++++ crates/sam2/src/prompt.rs | 181 +++++++++++++++++++++++ crates/sam2/tests/reference.rs | 36 +++++ 5 files changed, 555 insertions(+) create mode 100644 crates/sam2/src/decoder.rs create mode 100644 crates/sam2/src/model.rs create mode 100644 crates/sam2/src/prompt.rs diff --git a/crates/sam2/src/decoder.rs b/crates/sam2/src/decoder.rs new file mode 100644 index 000000000..e42a7acd9 --- /dev/null +++ b/crates/sam2/src/decoder.rs @@ -0,0 +1,258 @@ +// Modified work: ported from Hugging Face `transformers` `models/sam2/modeling_sam2.py` +// (Sam2Attention, Sam2TwoWayAttentionBlock, Sam2TwoWayTransformer, Sam2MaskDecoder; Apache-2.0, +// Copyright The HuggingFace Inc. team) to Rust on candle. See NOTICE. +//! The two-way mask decoder: output tokens and prompt tokens attend to the image and back, then +//! hyper-networks turn the mask tokens into masks over the up-scaled image embedding. +//! One object per call (batch 1, one point batch). + +use candle_core::{D, Tensor}; +use candle_nn::{LayerNorm, Linear, Module, ops::sigmoid, ops::softmax}; + +use crate::nn::{layer_norm, linear}; +use crate::weights::Weights; +use crate::{Error, Result}; + +const DIM: usize = 256; +const MASK_TOKENS: usize = 4; + +struct Attention { + q: Linear, + k: Linear, + v: Linear, + o: Linear, + heads: usize, + scale: f64, +} + +impl Attention { + fn load(w: &Weights, prefix: &str, downsample: usize, heads: usize) -> Result { + let inner = DIM / downsample; + Ok(Attention { + q: linear(w, &format!("{prefix}.q_proj"), DIM, inner)?, + k: linear(w, &format!("{prefix}.k_proj"), DIM, inner)?, + v: linear(w, &format!("{prefix}.v_proj"), DIM, inner)?, + o: linear(w, &format!("{prefix}.o_proj"), inner, DIM)?, + heads, + scale: ((inner / heads) as f64).powf(-0.5), + }) + } + + /// `q`: `[1, nq, 256]`, `k` and `v`: `[1, nk, 256]`. + fn forward(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result { + let split = |t: Tensor| -> Result { Ok(t.reshape((1, (), self.heads, t.dim(2)? / self.heads))?.transpose(1, 2)?.contiguous()?) }; + let (q, k, v) = (split(self.q.forward(q)?)?, split(self.k.forward(k)?)?, split(self.v.forward(v)?)?); + let att = softmax(&(q.matmul(&k.transpose(2, 3)?.contiguous()?)? * self.scale)?, D::Minus1)?; + let out = att.matmul(&v)?.transpose(1, 2)?.contiguous()?; + let (n, h, d) = (out.dim(1)?, out.dim(2)?, out.dim(3)?); + Ok(self.o.forward(&out.reshape((1, n, h * d))?)?) + } +} + +/// Feed-forward with ReLU, `Sam2FeedForward`: `proj_in`, `layers`, `proj_out`. +pub struct FeedForward { + proj_in: Linear, + layers: Vec, + proj_out: Linear, + sigmoid_output: bool, +} + +impl FeedForward { + pub fn load(w: &Weights, prefix: &str, input: usize, hidden: usize, output: usize, num_layers: usize, sigmoid_output: bool) -> Result { + Ok(FeedForward { + proj_in: linear(w, &format!("{prefix}.proj_in"), input, hidden)?, + layers: (0..num_layers.saturating_sub(2)).map(|i| linear(w, &format!("{prefix}.layers.{i}"), hidden, hidden)).collect::>()?, + proj_out: linear(w, &format!("{prefix}.proj_out"), hidden, output)?, + sigmoid_output, + }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + let mut x = self.proj_in.forward(x)?.relu()?; + for l in &self.layers { + x = l.forward(&x)?.relu()?; + } + let x = self.proj_out.forward(&x)?; + Ok(if self.sigmoid_output { sigmoid(&x)? } else { x }) + } +} + +struct Block { + self_attn: Attention, + ln1: LayerNorm, + token_to_image: Attention, + ln2: LayerNorm, + mlp: FeedForward, + ln3: LayerNorm, + ln4: LayerNorm, + image_to_token: Attention, + skip_first_pe: bool, +} + +impl Block { + fn forward(&self, queries: &Tensor, keys: &Tensor, query_pe: &Tensor, key_pe: &Tensor) -> Result<(Tensor, Tensor)> { + let mut queries = if self.skip_first_pe { + self.self_attn.forward(queries, queries, queries)? + } else { + let q = (queries + query_pe)?; + (queries + self.self_attn.forward(&q, &q, queries)?)? + }; + queries = self.ln1.forward(&queries)?; + let (q, k) = ((&queries + query_pe)?, (keys + key_pe)?); + queries = self.ln2.forward(&(&queries + self.token_to_image.forward(&q, &k, keys)?)?)?; + queries = self.ln3.forward(&(&queries + self.mlp.forward(&queries)?)?)?; + let (q, k) = ((&queries + query_pe)?, (keys + key_pe)?); + let keys = self.ln4.forward(&(keys + self.image_to_token.forward(&k, &q, &queries)?)?)?; + Ok((queries, keys)) + } +} + +/// What the decoder returns for one object. +pub struct DecoderOutput { + /// Mask logits `[1, k, 4 * GRID, 4 * GRID]` (`k` is 3 for multimask output, else 1). + pub masks: Tensor, + /// Predicted IoU per mask `[1, k]`. + pub iou: Tensor, + /// The output tokens of the chosen masks `[1, k, 256]` (the object pointer is made from these). + pub sam_tokens: Tensor, + /// Object-presence logit `[1, 1]` (negative: the object is not in the frame). + pub object_score: Tensor, +} + +pub struct MaskDecoder { + obj_token: Tensor, + iou_token: Tensor, + mask_tokens: Tensor, + layers: Vec, + final_attn: Attention, + final_ln: LayerNorm, + up1: (Tensor, Tensor), + up2: (Tensor, Tensor), + up_ln: LayerNorm, + conv_s0: (Tensor, Tensor), + conv_s1: (Tensor, Tensor), + hyper: Vec, + iou_head: FeedForward, + obj_head: FeedForward, +} + +fn conv_w(w: &Weights, name: &str, i: usize, o: usize, k: usize, transposed: bool) -> Result<(Tensor, Tensor)> { + let shape = if transposed { [i, o, k, k] } else { [o, i, k, k] }; + Ok((w.get(&format!("{name}.weight"), &shape)?, w.get(&format!("{name}.bias"), &[o])?)) +} + +fn add_bias(x: Tensor, b: &Tensor) -> Result { + Ok(x.broadcast_add(&b.reshape((1, b.dim(0)?, 1, 1))?)?) +} + +impl MaskDecoder { + pub fn load(w: &Weights) -> Result { + let layers = (0..2) + .map(|i| { + let p = format!("mask_decoder.transformer.layers.{i}"); + Ok(Block { + self_attn: Attention::load(w, &format!("{p}.self_attn"), 1, 8)?, + ln1: layer_norm(w, &format!("{p}.layer_norm1"), DIM, 1e-5)?, + token_to_image: Attention::load(w, &format!("{p}.cross_attn_token_to_image"), 2, 8)?, + ln2: layer_norm(w, &format!("{p}.layer_norm2"), DIM, 1e-5)?, + mlp: FeedForward::load(w, &format!("{p}.mlp"), DIM, 2048, DIM, 2, false)?, + ln3: layer_norm(w, &format!("{p}.layer_norm3"), DIM, 1e-5)?, + ln4: layer_norm(w, &format!("{p}.layer_norm4"), DIM, 1e-5)?, + image_to_token: Attention::load(w, &format!("{p}.cross_attn_image_to_token"), 2, 8)?, + skip_first_pe: i == 0, + }) + }) + .collect::>>()?; + Ok(MaskDecoder { + obj_token: w.get("mask_decoder.obj_score_token.weight", &[1, DIM])?, + iou_token: w.get("mask_decoder.iou_token.weight", &[1, DIM])?, + mask_tokens: w.get("mask_decoder.mask_tokens.weight", &[MASK_TOKENS, DIM])?, + layers, + final_attn: Attention::load(w, "mask_decoder.transformer.final_attn_token_to_image", 2, 8)?, + final_ln: layer_norm(w, "mask_decoder.transformer.layer_norm_final_attn", DIM, 1e-5)?, + up1: conv_w(w, "mask_decoder.upscale_conv1", DIM, DIM / 4, 2, true)?, + up2: conv_w(w, "mask_decoder.upscale_conv2", DIM / 4, DIM / 8, 2, true)?, + up_ln: layer_norm(w, "mask_decoder.upscale_layer_norm", DIM / 4, 1e-6)?, + conv_s0: conv_w(w, "mask_decoder.conv_s0", DIM, DIM / 8, 1, false)?, + conv_s1: conv_w(w, "mask_decoder.conv_s1", DIM, DIM / 4, 1, false)?, + hyper: (0..MASK_TOKENS) + .map(|i| FeedForward::load(w, &format!("mask_decoder.output_hypernetworks_mlps.{i}"), DIM, DIM, DIM / 8, 3, false)) + .collect::>()?, + iou_head: FeedForward::load(w, "mask_decoder.iou_prediction_head", DIM, DIM, MASK_TOKENS, 3, true)?, + obj_head: FeedForward::load(w, "mask_decoder.pred_obj_score_head", DIM, DIM, 1, 3, false)?, + }) + } + + /// Project the two highest-resolution neck levels to the decoder's skip features + /// (`conv_s0` on the 4x level, `conv_s1` on the 2x level), `[1, 32, 4g, 4g]` and `[1, 64, 2g, 2g]`. + pub fn high_res(&self, level0: &Tensor, level1: &Tensor) -> Result<(Tensor, Tensor)> { + let s0 = add_bias(level0.conv2d(&self.conv_s0.0, 0, 1, 1, 1)?, &self.conv_s0.1)?; + let s1 = add_bias(level1.conv2d(&self.conv_s1.0, 0, 1, 1, 1)?, &self.conv_s1.1)?; + Ok((s0, s1)) + } + + /// Predict masks. `image` and `pos`: `[1, 256, g, g]` (the lowest-resolution level with the + /// no-memory embedding already added when there is no memory), `sparse`: `[1, n, 256]`, + /// `dense`: `[1, 256, g, g]`, `high`: the skip features from [`MaskDecoder::high_res`]. + pub fn forward(&self, image: &Tensor, pos: &Tensor, sparse: &Tensor, dense: &Tensor, high: &(Tensor, Tensor), multimask: bool) -> Result { + let (_, c, g, gw) = image.dims4()?; + if c != DIM || g != gw { + return Err(Error::Input(format!("image embedding must be 256 x g x g, got {:?}", image.dims()))); + } + let output_tokens = Tensor::cat(&[&self.obj_token, &self.iou_token, &self.mask_tokens], 0)?.unsqueeze(0)?; + let tokens = Tensor::cat(&[&output_tokens, sparse], 1)?; + let image = (image + dense)?; + let flat = |t: &Tensor| -> Result { Ok(t.flatten_from(2)?.transpose(1, 2)?.contiguous()?) }; + let (mut keys, key_pe) = (flat(&image)?, flat(pos)?); + let mut queries = tokens.clone(); + for l in &self.layers { + (queries, keys) = l.forward(&queries, &keys, &tokens, &key_pe)?; + } + let (q, k) = ((&queries + &tokens)?, (&keys + &key_pe)?); + let queries = self.final_ln.forward(&(&queries + self.final_attn.forward(&q, &k, &keys)?)?)?; + + let iou_out = queries.narrow(1, 1, 1)?.squeeze(1)?; + let mask_tokens_out = queries.narrow(1, 2, MASK_TOKENS)?; + let obj_out = queries.narrow(1, 0, 1)?.squeeze(1)?; + + let emb = keys.transpose(1, 2)?.contiguous()?.reshape((1, DIM, g, g))?; + let (s0, s1) = high; + let up = add_bias(emb.conv_transpose2d(&self.up1.0, 0, 0, 2, 1)?, &self.up1.1)?; + let up = (up + s1)?; + let up = self.up_ln.forward(&up.permute((0, 2, 3, 1))?)?.permute((0, 3, 1, 2))?.contiguous()?.gelu_erf()?; + let up = (add_bias(up.conv_transpose2d(&self.up2.0, 0, 0, 2, 1)?, &self.up2.1)? + s0)?.gelu_erf()?; + + let hyper: Vec = self.hyper.iter().enumerate().map(|(i, m)| m.forward(&mask_tokens_out.narrow(1, i, 1)?.squeeze(1)?)).collect::>()?; + let hyper = Tensor::stack(&hyper, 1)?; + let (_, uc, uh, uw) = up.dims4()?; + let masks = hyper.matmul(&up.reshape((1, uc, uh * uw))?)?.reshape((1, MASK_TOKENS, uh, uw))?; + let iou = self.iou_head.forward(&iou_out)?; + let object_score = self.obj_head.forward(&obj_out)?; + + let (masks, iou, sam_tokens) = if multimask { + (masks.narrow(1, 1, 3)?, iou.narrow(1, 1, 3)?, mask_tokens_out.narrow(1, 1, 3)?) + } else { + let (m, i) = dynamic_single(&masks, &iou)?; + (m, i, mask_tokens_out.narrow(1, 0, 1)?) + }; + Ok(DecoderOutput { masks, iou, sam_tokens, object_score }) + } +} + +/// Single-mask output with the stability fallback (`_dynamic_multimask_via_stability`): if the +/// first mask token's mask is unstable, use the best of the other three. +fn dynamic_single(masks: &Tensor, iou: &Tensor) -> Result<(Tensor, Tensor)> { + const DELTA: f32 = 0.05; + const THRESH: f32 = 0.98; + let single = masks.narrow(1, 0, 1)?; + let single_iou = iou.narrow(1, 0, 1)?; + let logits = single.flatten_all()?.to_vec1::()?; + let area_i = logits.iter().filter(|v| **v > DELTA).count() as f32; + let area_u = logits.iter().filter(|v| **v > -DELTA).count() as f32; + let stability = if area_u > 0.0 { area_i / area_u } else { 1.0 }; + if stability >= THRESH { + return Ok((single, single_iou)); + } + let multi_iou = iou.narrow(1, 1, 3)?.flatten_all()?.to_vec1::()?; + let best = multi_iou.iter().enumerate().fold((0usize, f32::NEG_INFINITY), |b, (i, v)| if *v > b.1 { (i, *v) } else { b }).0; + Ok((masks.narrow(1, 1 + best, 1)?, iou.narrow(1, 1 + best, 1)?)) +} diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 6515eff7b..0a67f8ad7 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -11,9 +11,12 @@ use std::fmt; +pub mod decoder; pub mod hiera; +pub mod model; pub mod neck; pub mod nn; +pub mod prompt; pub mod weights; /// Everything that can go wrong in this crate. Never a panic. diff --git a/crates/sam2/src/model.rs b/crates/sam2/src/model.rs new file mode 100644 index 000000000..13481fdb9 --- /dev/null +++ b/crates/sam2/src/model.rs @@ -0,0 +1,77 @@ +// Modified work: ported from Hugging Face `transformers` `Sam2Model` (Apache-2.0, Copyright The +// HuggingFace Inc. team) to Rust on candle. See NOTICE. +//! The single-image model: encode a frame once, then decode any number of prompts against it. + +use std::path::Path; + +use candle_core::{Device, Tensor}; + +use crate::decoder::{DecoderOutput, MaskDecoder}; +use crate::neck::{ImageFeatures, VisionEncoder}; +use crate::prompt::{INPUT_SIZE, Prompt, PromptEncoder}; +use crate::weights::Weights; +use crate::{Error, Result}; + +/// A frame after the vision encoder: reusable for every prompt on that frame. +pub struct Encoded { + /// The three neck levels, highest resolution first (`[1, 256, 256, 256]`, `128`, `64`). + pub features: ImageFeatures, + /// The decoder's skip features from the two highest-resolution levels. + pub high: (Tensor, Tensor), +} + +pub struct ImageModel { + pub vision: VisionEncoder, + pub prompt: PromptEncoder, + pub decoder: MaskDecoder, + no_memory: Tensor, + image_pos: Tensor, +} + +impl ImageModel { + pub fn load(w: &Weights) -> Result { + let prompt = PromptEncoder::load(w)?; + Ok(ImageModel { + vision: VisionEncoder::load_tiny(w)?, + decoder: MaskDecoder::load(w)?, + no_memory: w.get("no_memory_embedding", &[1, 1, 256])?.reshape((1, 256, 1, 1))?, + image_pos: prompt.image_positions()?, + prompt, + }) + } + + pub fn open(path: &Path) -> Result { + ImageModel::load(&Weights::open(path)?) + } + + /// Encode a normalised `[1, 3, 1024, 1024]` frame. + pub fn encode(&self, pixels: &Tensor) -> Result { + if pixels.dims() != [1, 3, INPUT_SIZE, INPUT_SIZE] { + return Err(Error::Input(format!("pixels must be [1, 3, {INPUT_SIZE}, {INPUT_SIZE}], got {:?}", pixels.dims()))); + } + let features = self.vision.forward(pixels)?; + let (l0, l1) = match (features.features.first(), features.features.get(1)) { + (Some(a), Some(b)) => (a, b), + _ => return Err(Error::Compute("vision encoder returned too few levels".into())), + }; + let high = self.decoder.high_res(l0, l1)?; + Ok(Encoded { features, high }) + } + + /// Prompt embeddings (sparse tokens, empty dense map) for a click/box prompt. + pub fn embed_prompt(&self, p: &Prompt) -> Result<(Tensor, Tensor)> { + Ok((self.prompt.sparse(p)?, self.prompt.dense_empty()?)) + } + + /// Decode masks for a prompt on an encoded frame. + pub fn segment(&self, enc: &Encoded, p: &Prompt, multimask: bool) -> Result { + let low = enc.features.features.get(2).ok_or_else(|| Error::Compute("vision encoder returned too few levels".into()))?; + let image = low.broadcast_add(&self.no_memory)?; + let (sparse, dense) = self.embed_prompt(p)?; + self.decoder.forward(&image, &self.image_pos, &sparse, &dense, &enc.high, multimask) + } + + pub fn device(&self) -> Device { + Device::Cpu + } +} diff --git a/crates/sam2/src/prompt.rs b/crates/sam2/src/prompt.rs new file mode 100644 index 000000000..95834c3cb --- /dev/null +++ b/crates/sam2/src/prompt.rs @@ -0,0 +1,181 @@ +// Modified work: ported from Hugging Face `transformers` `models/sam2/modeling_sam2.py` +// (Sam2PositionalEmbedding, Sam2MaskEmbedding, Sam2PromptEncoder; Apache-2.0, Copyright The +// HuggingFace Inc. team) to Rust on candle. See NOTICE. +//! The prompt encoder: clicks and boxes become sparse tokens, a mask or "no mask" becomes a dense map. + +use candle_core::{Device, Tensor}; +use candle_nn::{LayerNorm, Module}; + +use crate::nn::layer_norm; +use crate::weights::Weights; +use crate::{Error, Result}; + +/// Side of the (square) model input in pixels; prompts are given in these coordinates. +pub const INPUT_SIZE: usize = 1024; +/// Side of the image-embedding grid (input / 16). +pub const GRID: usize = INPUT_SIZE / 16; +const DIM: usize = 256; + +/// What the user pointed at, in model-input pixels (0..1024). +#[derive(Clone, Debug, Default, PartialEq)] +pub struct Prompt { + /// Clicks and whether each is foreground. + pub points: Vec<([f32; 2], bool)>, + /// A box `[x0, y0, x1, y1]`. + pub bbox: Option<[f32; 4]>, +} + +impl Prompt { + /// Reject non-finite or absurd coordinates (they come from the user or an agent). + pub fn validate(&self) -> Result<()> { + let bad = |v: f32| !v.is_finite() || v.abs() > 4.0 * INPUT_SIZE as f32; + if self.points.iter().any(|(p, _)| p.iter().any(|v| bad(*v))) || self.bbox.is_some_and(|b| b.iter().any(|v| bad(*v))) { + return Err(Error::Input("prompt coordinates must be finite and near the frame".into())); + } + if self.points.len() > 64 { + return Err(Error::Input("at most 64 points per prompt".into())); + } + Ok(()) + } +} + +struct MaskEmbed { + conv: [(Tensor, Tensor, usize, usize); 3], + ln1: LayerNorm, + ln2: LayerNorm, +} + +/// Sparse and dense prompt embeddings. +pub struct PromptEncoder { + gauss: Vec, + point_embed: Vec, + not_a_point: Vec, + no_mask: Tensor, + mask: MaskEmbed, +} + +fn rows(t: Tensor) -> Result> { + Ok(t.flatten_all()?.to_vec1::()?) +} + +impl PromptEncoder { + pub fn load(w: &Weights) -> Result { + let conv = |name: &str, i: usize, o: usize, k: usize, stride: usize| -> Result<(Tensor, Tensor, usize, usize)> { + Ok(( + w.get(&format!("prompt_encoder.mask_embed.{name}.weight"), &[o, i, k, k])?, + w.get(&format!("prompt_encoder.mask_embed.{name}.bias"), &[o])?, + k, + stride, + )) + }; + Ok(PromptEncoder { + gauss: rows(w.get("prompt_encoder.shared_embedding.positional_embedding", &[2, DIM / 2])?)?, + point_embed: rows(w.get("prompt_encoder.point_embed.weight", &[4, DIM])?)?, + not_a_point: rows(w.get("prompt_encoder.not_a_point_embed.weight", &[1, DIM])?)?, + no_mask: w.get("prompt_encoder.no_mask_embed.weight", &[1, DIM])?, + mask: MaskEmbed { + conv: [conv("conv1", 1, 4, 2, 2)?, conv("conv2", 4, 16, 2, 2)?, conv("conv3", 16, DIM, 1, 1)?], + ln1: layer_norm(w, "prompt_encoder.mask_embed.layer_norm1", 4, 1e-6)?, + ln2: layer_norm(w, "prompt_encoder.mask_embed.layer_norm2", 16, 1e-6)?, + }, + }) + } + + /// Random-Fourier position encoding of coordinates already normalised to 0..1: `[n * 256]`. + fn fourier(&self, coords01: &[[f32; 2]]) -> Vec { + let half = DIM / 2; + let mut out = vec![0f32; coords01.len() * DIM]; + for (n, c) in coords01.iter().enumerate() { + let (x, y) = (2.0 * c[0] - 1.0, 2.0 * c[1] - 1.0); + for j in 0..half { + let g0 = self.gauss.get(j).copied().unwrap_or(0.0); + let g1 = self.gauss.get(half + j).copied().unwrap_or(0.0); + let v = 2.0 * std::f32::consts::PI * (x * g0 + y * g1); + if let Some(o) = out.get_mut(n * DIM + j) { + *o = v.sin(); + } + if let Some(o) = out.get_mut(n * DIM + half + j) { + *o = v.cos(); + } + } + } + out + } + + /// The position encoding of the whole `GRID x GRID` embedding, `[1, 256, GRID, GRID]`. + pub fn image_positions(&self) -> Result { + let coords: Vec<[f32; 2]> = (0..GRID * GRID).map(|i| [((i % GRID) as f32 + 0.5) / GRID as f32, ((i / GRID) as f32 + 0.5) / GRID as f32]).collect(); + let flat = self.fourier(&coords); + Ok(Tensor::from_vec(flat, (GRID * GRID, DIM), &Device::Cpu)?.t()?.reshape((1, DIM, GRID, GRID))?.contiguous()?) + } + + fn point_row(&self, i: usize) -> &[f32] { + self.point_embed.get(i * DIM..(i + 1) * DIM).unwrap_or(&[]) + } + + /// Sparse tokens `[1, n, 256]` for a prompt (a click-only prompt gets one padding token). + pub fn sparse(&self, p: &Prompt) -> Result { + p.validate()?; + let norm = |c: [f32; 2]| [(c[0] + 0.5) / INPUT_SIZE as f32, (c[1] + 0.5) / INPUT_SIZE as f32]; + let mut tokens: Vec = vec![]; + let mut n = 0; + let mut push = |v: &[f32]| { + tokens.extend_from_slice(v); + n += 1; + }; + let pts: Vec<[f32; 2]> = p.points.iter().map(|(c, _)| norm(*c)).collect(); + let enc = self.fourier(&pts); + for (k, (_, fg)) in p.points.iter().enumerate() { + let e = enc.get(k * DIM..(k + 1) * DIM).unwrap_or(&[]); + let label_row = self.point_row(usize::from(*fg)); + let v: Vec = e.iter().zip(label_row).map(|(a, b)| a + b).collect(); + push(&v); + } + match p.bbox { + None => push(&self.not_a_point.clone()), + Some(b) => { + let corners = [norm([b[0], b[1]]), norm([b[2], b[3]])]; + let enc = self.fourier(&corners); + for (k, row) in [2usize, 3].into_iter().enumerate() { + let e = enc.get(k * DIM..(k + 1) * DIM).unwrap_or(&[]); + let v: Vec = e.iter().zip(self.point_row(row)).map(|(a, b)| a + b).collect(); + push(&v); + } + push(&self.not_a_point.clone()); + } + } + Ok(Tensor::from_vec(tokens, (1, n, DIM), &Device::Cpu)?) + } + + /// The dense embedding for "no mask prompt", `[1, 256, GRID, GRID]`. + pub fn dense_empty(&self) -> Result { + Ok(self.no_mask.reshape((1, DIM, 1, 1))?.broadcast_as((1, DIM, GRID, GRID))?.contiguous()?) + } + + /// The dense embedding of a mask prompt `[1, 1, 4 * GRID, 4 * GRID]` (low-resolution logits). + pub fn dense_mask(&self, mask: &Tensor) -> Result { + let (c1, c2, c3) = (&self.mask.conv[0], &self.mask.conv[1], &self.mask.conv[2]); + let step = |x: &Tensor, c: &(Tensor, Tensor, usize, usize)| -> Result { + Ok(x.conv2d(&c.0, 0, c.3, 1, 1)?.broadcast_add(&c.1.reshape((1, c.1.dim(0)?, 1, 1))?)?) + }; + let cl = |x: &Tensor, ln: &LayerNorm| -> Result { Ok(ln.forward(&x.permute((0, 2, 3, 1))?)?.permute((0, 3, 1, 2))?.contiguous()?) }; + let x = cl(&step(mask, c1)?, &self.mask.ln1)?.gelu_erf()?; + let x = cl(&step(&x, c2)?, &self.mask.ln2)?.gelu_erf()?; + step(&x, c3) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn hostile_prompts_are_rejected() { + let ok = Prompt { points: vec![([10.0, 20.0], true)], bbox: None }; + assert!(ok.validate().is_ok()); + assert!(Prompt { points: vec![([f32::NAN, 0.0], true)], bbox: None }.validate().is_err()); + assert!(Prompt { points: vec![], bbox: Some([0.0, 0.0, f32::INFINITY, 5.0]) }.validate().is_err()); + assert!(Prompt { points: vec![([1e9, 0.0], false)], bbox: None }.validate().is_err()); + assert!(Prompt { points: vec![([1.0, 1.0], true); 65], bbox: None }.validate().is_err()); + } +} diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index bae7d42e3..6536a9eb5 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -107,3 +107,39 @@ fn neck_and_vision_encoder_match_pytorch() { let e = max_abs(feats.features.first().unwrap(), r.get("vision_encoder.neck.3").unwrap()); assert!(e < NECK_TOL); } + +/// Click-prompted single image: prompt embeddings and decoder outputs against `Sam2Model`. +const DECODER_TOL: f32 = 1e-3; + +#[test] +fn prompt_encoder_and_decoder_match_pytorch() { + use effectcraft_sam2::model::ImageModel; + use effectcraft_sam2::prompt::Prompt; + let Some((w, r)) = load() else { return }; + let model = ImageModel::load(&w).unwrap(); + let pixels = r.get("input.pixel_values").unwrap(); + let prompt = Prompt { points: vec![([460.0, 512.0], true)], bbox: None }; + + let (sparse, dense) = model.embed_prompt(&prompt).unwrap(); + let want_sparse = r.get("prompt_encoder.0").unwrap().squeeze(0).unwrap(); + let e = max_abs(&sparse, &want_sparse); + eprintln!("sparse prompt: max abs error {e:e}"); + assert!(e < 1e-5, "sparse prompt error {e}"); + let e = max_abs(&dense, r.get("prompt_encoder.1").unwrap()); + assert!(e < 1e-6, "dense prompt error {e}"); + + let enc = model.encode(pixels).unwrap(); + let out = model.segment(&enc, &prompt, true).unwrap(); + let want_masks = r.get("mask_decoder.0").unwrap().squeeze(0).unwrap(); + let e = max_abs(&out.masks, &want_masks); + eprintln!("masks: max abs error {e:e} (logit range {:?})", want_masks.flatten_all().unwrap().max(0).unwrap().to_scalar::()); + assert!(e < DECODER_TOL, "mask logits error {e}"); + let e = max_abs(&out.iou, &r.get("mask_decoder.1").unwrap().squeeze(0).unwrap()); + eprintln!("iou: max abs error {e:e}"); + assert!(e < 1e-4, "iou error {e}"); + let e = max_abs(&out.object_score, &r.get("mask_decoder.3").unwrap().squeeze(0).unwrap()); + eprintln!("object score: max abs error {e:e}"); + assert!(e < 1e-4, "object score error {e}"); + let e = max_abs(&out.sam_tokens, &r.get("mask_decoder.2").unwrap().squeeze(0).unwrap()); + assert!(e < 1e-4, "sam token error {e}"); +} From 5aa52c5ba3e016232c18bca2df87775bd74662fd Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:26:16 +0200 Subject: [PATCH 09/77] S2.14: sam2_bench example and first CPU speedups Encode 5.0 s -> 3.4 s per 1024x1024 frame in release on 16 threads (fused softmax, queries scaled before the matmul, flat 2-D linear layers); a click on an encoded frame is about 120 ms. Reference tests unchanged. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/examples/sam2_bench.rs | 51 ++++++++++++++++++++++++++++++ crates/sam2/src/decoder.rs | 14 ++++---- crates/sam2/src/hiera.rs | 20 ++++++------ crates/sam2/src/nn.rs | 35 +++++++++++++++++--- crates/sam2/src/prompt.rs | 3 +- 5 files changed, 101 insertions(+), 22 deletions(-) create mode 100644 crates/sam2/examples/sam2_bench.rs diff --git a/crates/sam2/examples/sam2_bench.rs b/crates/sam2/examples/sam2_bench.rs new file mode 100644 index 000000000..7bc678b7c --- /dev/null +++ b/crates/sam2/examples/sam2_bench.rs @@ -0,0 +1,51 @@ +//! Timing of the SAM 2 image path on this machine: load, encode one frame, decode a click. +//! +//! `cargo run -p effectcraft-sam2 --example sam2_bench -- [MODEL_DIR]` (default `sam2-tiny`). +//! Prints a message and exits when the weights are not there. + +use std::time::Instant; + +use candle_core::{DType, Device, Tensor}; +use effectcraft_sam2::model::ImageModel; +use effectcraft_sam2::prompt::Prompt; + +fn main() { + if let Err(e) = run() { + eprintln!("sam2_bench: {e}"); + std::process::exit(1); + } +} + +fn run() -> Result<(), Box> { + let dir = std::env::args().nth(1).unwrap_or_else(|| "sam2-tiny".to_string()); + let file = std::path::Path::new(&dir).join("model.safetensors"); + if !file.exists() { + println!("no weights at {} (download facebook/sam2.1-hiera-tiny first)", file.display()); + return Ok(()); + } + let t = Instant::now(); + let model = ImageModel::open(&file)?; + println!("load: {:.2} s", t.elapsed().as_secs_f64()); + let pixels = Tensor::randn(0f32, 1.0, (1, 3, 1024, 1024), &Device::Cpu)?.to_dtype(DType::F32)?; + for run in 0..3 { + let t = Instant::now(); + let enc = model.encode(&pixels)?; + let encode = t.elapsed().as_secs_f64(); + let prompt = Prompt { points: vec![([512.0, 512.0], true)], bbox: None }; + let t = Instant::now(); + let out = model.segment(&enc, &prompt, true)?; + println!("run {run}: encode {encode:.2} s, click {:.0} ms, masks {:?}", t.elapsed().as_secs_f64() * 1e3, out.masks.dims()); + } + // Where the encode time goes: patch embedding, then each Hiera block. + let hiera = &model.vision.hiera; + let t = Instant::now(); + let mut x = hiera.patch_embed(&pixels)?; + println!("patch embed: {:.0} ms", t.elapsed().as_secs_f64() * 1e3); + for (i, b) in hiera.blocks.iter().enumerate() { + let t = Instant::now(); + x = b.forward(&x)?; + println!("block {i}: {:.0} ms, out {:?}", t.elapsed().as_secs_f64() * 1e3, x.dims()); + } + println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get())); + Ok(()) +} diff --git a/crates/sam2/src/decoder.rs b/crates/sam2/src/decoder.rs index e42a7acd9..d7a7a6dcf 100644 --- a/crates/sam2/src/decoder.rs +++ b/crates/sam2/src/decoder.rs @@ -6,9 +6,9 @@ //! One object per call (batch 1, one point batch). use candle_core::{D, Tensor}; -use candle_nn::{LayerNorm, Linear, Module, ops::sigmoid, ops::softmax}; +use candle_nn::{Linear, ops::sigmoid, ops::softmax}; -use crate::nn::{layer_norm, linear}; +use crate::nn::{LayerNorm, apply, layer_norm, linear}; use crate::weights::Weights; use crate::{Error, Result}; @@ -40,11 +40,11 @@ impl Attention { /// `q`: `[1, nq, 256]`, `k` and `v`: `[1, nk, 256]`. fn forward(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result { let split = |t: Tensor| -> Result { Ok(t.reshape((1, (), self.heads, t.dim(2)? / self.heads))?.transpose(1, 2)?.contiguous()?) }; - let (q, k, v) = (split(self.q.forward(q)?)?, split(self.k.forward(k)?)?, split(self.v.forward(v)?)?); + let (q, k, v) = (split(apply(&self.q, q)?)?, split(apply(&self.k, k)?)?, split(apply(&self.v, v)?)?); let att = softmax(&(q.matmul(&k.transpose(2, 3)?.contiguous()?)? * self.scale)?, D::Minus1)?; let out = att.matmul(&v)?.transpose(1, 2)?.contiguous()?; let (n, h, d) = (out.dim(1)?, out.dim(2)?, out.dim(3)?); - Ok(self.o.forward(&out.reshape((1, n, h * d))?)?) + apply(&self.o, &out.reshape((1, n, h * d))?) } } @@ -67,11 +67,11 @@ impl FeedForward { } pub fn forward(&self, x: &Tensor) -> Result { - let mut x = self.proj_in.forward(x)?.relu()?; + let mut x = apply(&self.proj_in, x)?.relu()?; for l in &self.layers { - x = l.forward(&x)?.relu()?; + x = apply(l, &x)?.relu()?; } - let x = self.proj_out.forward(&x)?; + let x = apply(&self.proj_out, &x)?; Ok(if self.sigmoid_output { sigmoid(&x)? } else { x }) } } diff --git a/crates/sam2/src/hiera.rs b/crates/sam2/src/hiera.rs index 5091d331b..90833f7da 100644 --- a/crates/sam2/src/hiera.rs +++ b/crates/sam2/src/hiera.rs @@ -4,10 +4,10 @@ //! The Hiera backbone: patch embedding, windowed positional embedding, and blocks with window or //! global attention and query pooling at the stage changes. Activations are NHWC like the reference. -use candle_core::{D, Tensor}; -use candle_nn::{LayerNorm, Linear, Module, ops::softmax}; +use candle_core::Tensor; +use candle_nn::{Linear, ops::softmax_last_dim}; -use crate::nn::{Mlp, layer_norm, linear}; +use crate::nn::{LayerNorm, Mlp, apply, layer_norm, linear}; use crate::weights::Weights; use crate::{Error, Result}; @@ -89,7 +89,7 @@ struct Attention { impl Attention { fn forward(&self, x: &Tensor) -> Result { let (b, h, w, _) = x.dims4()?; - let qkv = self.qkv.forward(x)?.reshape((b, h * w, 3, self.heads, ()))?; + let qkv = apply(&self.qkv, x)?.reshape((b, h * w, 3, self.heads, ()))?; let part = |i: usize| -> Result { Ok(qkv.narrow(2, i, 1)?.squeeze(2)?) }; let (mut q, k, v) = (part(0)?, part(1)?, part(2)?); let (mut oh, mut ow) = (h, w); @@ -101,10 +101,12 @@ impl Attention { let q = q.transpose(1, 2)?.contiguous()?; let k = k.transpose(1, 2)?.contiguous()?; let v = v.transpose(1, 2)?.contiguous()?; - let att = (q.matmul(&k.transpose(2, 3)?.contiguous()?)? * self.scale)?; - let att = softmax(&att, D::Minus1)?; + // Scaling the (small) queries first equals scaling the scores, and the fused softmax is + // much faster than the generic one on the 4096 x 4096 global-attention maps. + let att = (q * self.scale)?.matmul(&k.transpose(2, 3)?.contiguous()?)?; + let att = softmax_last_dim(&att)?; let out = att.matmul(&v)?.transpose(1, 2)?.contiguous()?.reshape((b, oh, ow, ()))?; - Ok(self.proj.forward(&out)?) + apply(&self.proj, &out) } } @@ -149,8 +151,8 @@ impl Block { pub fn forward(&self, x: &Tensor) -> Result { let h = self.norm1.forward(x)?; let residual = match (&self.proj, self.query_stride) { - (Some(p), Some(s)) => do_pool(&p.forward(&h)?, s)?, - (Some(p), None) => p.forward(&h)?, + (Some(p), Some(s)) => do_pool(&apply(p, &h)?, s)?, + (Some(p), None) => apply(p, &h)?, (None, _) => x.clone(), }; let (hh, ww) = (h.dim(1)?, h.dim(2)?); diff --git a/crates/sam2/src/nn.rs b/crates/sam2/src/nn.rs index ef334432a..f90958eb1 100644 --- a/crates/sam2/src/nn.rs +++ b/crates/sam2/src/nn.rs @@ -2,17 +2,31 @@ //! Linear, layer norm and a few tensor helpers on candle. use candle_core::Tensor; -use candle_nn::{LayerNorm, Linear, Module}; +use candle_nn::{Linear, Module}; -use crate::Result; use crate::weights::Weights; +use crate::{Error, Result}; pub fn linear(w: &Weights, prefix: &str, input: usize, output: usize) -> Result { Ok(Linear::new(w.get(&format!("{prefix}.weight"), &[output, input])?, Some(w.get(&format!("{prefix}.bias"), &[output])?))) } +/// Layer norm over the last dimension using candle's fused kernel (much faster than the +/// composed `candle_nn::LayerNorm` on large activations). +pub struct LayerNorm { + weight: Tensor, + bias: Tensor, + eps: f32, +} + +impl LayerNorm { + pub fn forward(&self, x: &Tensor) -> Result { + Ok(candle_nn::ops::layer_norm(&x.contiguous()?, &self.weight, &self.bias, self.eps)?) + } +} + pub fn layer_norm(w: &Weights, prefix: &str, dim: usize, eps: f64) -> Result { - Ok(LayerNorm::new(w.get(&format!("{prefix}.weight"), &[dim])?, w.get(&format!("{prefix}.bias"), &[dim])?, eps)) + Ok(LayerNorm { weight: w.get(&format!("{prefix}.weight"), &[dim])?, bias: w.get(&format!("{prefix}.bias"), &[dim])?, eps: eps as f32 }) } /// Two-layer MLP with exact (erf) GELU, as `Sam2FeedForward` with `num_layers = 2`. @@ -27,6 +41,19 @@ impl Mlp { } pub fn forward(&self, x: &Tensor) -> Result { - Ok(self.proj_out.forward(&self.proj_in.forward(x)?.gelu_erf()?)?) + apply(&self.proj_out, &apply(&self.proj_in, x)?.gelu_erf()?) } } + +/// Apply a linear layer to `[.., in]` through a flat 2-D matmul. Candle's `Linear::forward` +/// broadcasts the weight over every leading dimension (copying it per window), which is slow +/// on the 4-D activations of the backbone. +pub fn apply(l: &Linear, x: &Tensor) -> Result { + let dims = x.dims(); + let (&last, lead) = dims.split_last().ok_or_else(|| Error::Compute("linear on a scalar".into()))?; + let rows: usize = lead.iter().product(); + let y = l.forward(&x.reshape((rows, last))?)?; + let mut out: Vec = lead.to_vec(); + out.push(y.dim(1)?); + Ok(y.reshape(out)?) +} diff --git a/crates/sam2/src/prompt.rs b/crates/sam2/src/prompt.rs index 95834c3cb..55dd901a2 100644 --- a/crates/sam2/src/prompt.rs +++ b/crates/sam2/src/prompt.rs @@ -4,9 +4,8 @@ //! The prompt encoder: clicks and boxes become sparse tokens, a mask or "no mask" becomes a dense map. use candle_core::{Device, Tensor}; -use candle_nn::{LayerNorm, Module}; -use crate::nn::layer_norm; +use crate::nn::{LayerNorm, layer_norm}; use crate::weights::Weights; use crate::{Error, Result}; From e4f4fea0f238ee111be2b64163f4cfcc6c6c3dc1 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:30:33 +0200 Subject: [PATCH 10/77] S2.21: memory encoder matches PyTorch Mask down-sampler, feature fuser and projection, sine position encoding, occlusion embedding and bfloat16 rounding of stored memories. Max abs error 3.5e-6 on the memory features, 6e-8 on the position encoding. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/src/lib.rs | 1 + crates/sam2/src/memory.rs | 132 +++++++++++++++++++++++++++++++++ crates/sam2/tests/reference.rs | 30 ++++++++ tools/sam2_video_reference.py | 112 ++++++++++++++++++++++++++++ 4 files changed, 275 insertions(+) create mode 100644 crates/sam2/src/memory.rs create mode 100644 tools/sam2_video_reference.py diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 0a67f8ad7..18fe5a8ae 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 memory; pub mod model; pub mod neck; pub mod nn; diff --git a/crates/sam2/src/memory.rs b/crates/sam2/src/memory.rs new file mode 100644 index 000000000..51e3d3c92 --- /dev/null +++ b/crates/sam2/src/memory.rs @@ -0,0 +1,132 @@ +// Modified work: ported from Hugging Face `transformers` `models/sam2_video/modeling_sam2_video.py` +// (Sam2VideoMemoryEncoder, Sam2VideoMaskDownSampler, Sam2VideoMemoryFuser; Apache-2.0, Copyright +// The HuggingFace Inc. team) to Rust on candle. See NOTICE. +//! The memory encoder: a frame's mask is down-sampled and fused with the frame's image features +//! into a compact memory (64 channels on the 64 x 64 grid) plus its position encoding. + +use candle_core::Tensor; + +use crate::Result; +use crate::neck::sine_position; +use crate::nn::{LayerNorm, apply, layer_norm, linear}; +use crate::weights::Weights; +use candle_nn::Linear; + +struct Conv { + w: Tensor, + b: Tensor, + padding: usize, + stride: usize, + groups: usize, +} + +impl Conv { + fn load(w: &Weights, name: &str, shape: &[usize], padding: usize, stride: usize, groups: usize) -> Result { + let out = shape.first().copied().unwrap_or(0); + Ok(Conv { w: w.get(&format!("{name}.weight"), shape)?, b: w.get(&format!("{name}.bias"), &[out])?, padding, stride, groups }) + } + + fn forward(&self, x: &Tensor) -> Result { + let y = x.conv2d(&self.w, self.padding, self.stride, 1, self.groups)?; + Ok(y.broadcast_add(&self.b.reshape((1, self.b.dim(0)?, 1, 1))?)?) + } +} + +/// Layer norm over the channel dimension of `[B, C, H, W]`. +fn norm_channels_first(ln: &LayerNorm, x: &Tensor) -> Result { + ln.forward(&x.permute((0, 2, 3, 1))?.contiguous()?)?.permute((0, 3, 1, 2))?.contiguous().map_err(Into::into) +} + +struct FuserBlock { + depthwise: Conv, + norm: LayerNorm, + pw1: Linear, + pw2: Linear, + scale: Tensor, +} + +impl FuserBlock { + fn forward(&self, x: &Tensor) -> Result { + let h = norm_channels_first(&self.norm, &self.depthwise.forward(x)?)?; + let h = h.permute((0, 2, 3, 1))?.contiguous()?; + let h = apply(&self.pw2, &apply(&self.pw1, &h)?.gelu_erf()?)?; + let h = h.broadcast_mul(&self.scale)?.permute((0, 3, 1, 2))?.contiguous()?; + Ok((x + h)?) + } +} + +pub struct MemoryEncoder { + down: Vec<(Conv, LayerNorm)>, + final_conv: Conv, + feature_projection: Conv, + fuser: Vec, + projection: Conv, + occlusion: Tensor, +} + +/// A frame's memory: `features` is `[HW, 64]` rounded through bfloat16 exactly as the reference +/// stores it, `pos` is the `[HW, 64]` sine position encoding. +pub struct Memory { + pub features: Tensor, + pub pos: Tensor, +} + +impl MemoryEncoder { + pub fn load(w: &Weights) -> Result { + let mut down = vec![]; + let mut ch = 1; + for i in 0..4 { + let out = ch * 4; + let p = format!("memory_encoder.mask_downsampler.layers.{i}"); + down.push((Conv::load(w, &format!("{p}.conv"), &[out, ch, 3, 3], 1, 2, 1)?, layer_norm(w, &format!("{p}.layer_norm"), out, 1e-6)?)); + ch = out; + } + let fuser = (0..2) + .map(|i| { + let p = format!("memory_encoder.memory_fuser.layers.{i}"); + Ok(FuserBlock { + depthwise: Conv::load(w, &format!("{p}.depthwise_conv"), &[256, 1, 7, 7], 3, 1, 256)?, + norm: layer_norm(w, &format!("{p}.layer_norm"), 256, 1e-6)?, + pw1: linear(w, &format!("{p}.pointwise_conv1"), 256, 1024)?, + pw2: linear(w, &format!("{p}.pointwise_conv2"), 1024, 256)?, + scale: w.get(&format!("{p}.scale"), &[256])?, + }) + }) + .collect::>>()?; + Ok(MemoryEncoder { + down, + final_conv: Conv::load(w, "memory_encoder.mask_downsampler.final_conv", &[256, 256, 1, 1], 0, 1, 1)?, + feature_projection: Conv::load(w, "memory_encoder.feature_projection", &[256, 256, 1, 1], 0, 1, 1)?, + fuser, + projection: Conv::load(w, "memory_encoder.projection", &[64, 256, 1, 1], 0, 1, 1)?, + occlusion: w.get("occlusion_spatial_embedding_parameter", &[1, 64])?, + }) + } + + /// `features`: `[1, 256, g, g]` image features; `mask`: `[1, 1, 16 g, 16 g]` already scaled + /// (`20 * m - 10`). Returns the raw `[1, 64, g, g]` features and position encoding. + pub fn encode(&self, features: &Tensor, mask: &Tensor) -> Result<(Tensor, Tensor)> { + let mut m = mask.clone(); + for (conv, ln) in &self.down { + m = norm_channels_first(ln, &conv.forward(&m)?)?.gelu_erf()?; + } + let m = self.final_conv.forward(&m)?; + let mut x = (self.feature_projection.forward(features)? + m)?; + for b in &self.fuser { + x = b.forward(&x)?; + } + let x = self.projection.forward(&x)?; + let pos = sine_position(32, x.dim(2)?, x.dim(3)?)?; + Ok((x, pos)) + } + + /// The memory kept for a frame: the encoder output, plus the occlusion embedding when the + /// object is absent, rounded through bfloat16 like the reference. + pub fn memory(&self, features: &Tensor, mask: &Tensor, object_present: bool) -> Result { + let (x, pos) = self.encode(features, mask)?; + let x = if object_present { x } else { x.broadcast_add(&self.occlusion.reshape((1, 64, 1, 1))?)? }; + let flat = |t: &Tensor| -> Result { Ok(t.flatten_from(2)?.squeeze(0)?.t()?.contiguous()?) }; + let features = flat(&x)?.to_dtype(candle_core::DType::BF16)?.to_dtype(candle_core::DType::F32)?; + Ok(Memory { features, pos: flat(&pos)? }) + } +} diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index 6536a9eb5..f90d03ab8 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -143,3 +143,33 @@ fn prompt_encoder_and_decoder_match_pytorch() { let e = max_abs(&out.sam_tokens, &r.get("mask_decoder.2").unwrap().squeeze(0).unwrap()); assert!(e < 1e-4, "sam token error {e}"); } + +fn video_fixture_file() -> PathBuf { + std::env::var_os("EFFECTCRAFT_SAM2_VIDEO_REF").map_or_else(|| repo().join("plan/sam2-ref/video.safetensors"), PathBuf::from) +} + +/// The weights and the video reference tensors (`tools/sam2_video_reference.py`), or `None`. +fn load_video() -> Option<(Weights, HashMap)> { + let (wf, rf) = (weights_file(), video_fixture_file()); + if !wf.exists() || !rf.exists() { + eprintln!("skipped: SAM 2 weights or video reference tensors not found ({} / {})", wf.display(), rf.display()); + return None; + } + Some((Weights::open(&wf).unwrap(), candle_core::safetensors::load(&rf, &Device::Cpu).unwrap())) +} + +/// Measured max 3.5e-6 (memory features), 2.6e-5 (memory attention), 2026-10-15. +const MEMORY_TOL: f32 = 1e-4; + +#[test] +fn memory_encoder_matches_pytorch() { + let Some((w, r)) = load_video() else { return }; + let enc = effectcraft_sam2::memory::MemoryEncoder::load(&w).unwrap(); + let (feat, pos) = enc.encode(r.get("mem.enc.vision_features").unwrap(), r.get("mem.enc.masks").unwrap()).unwrap(); + let e = max_abs(&feat, r.get("mem.enc.out_features").unwrap()); + eprintln!("memory features: max abs error {e:e}"); + assert!(e < MEMORY_TOL, "memory feature error {e}"); + let e = max_abs(&pos, r.get("mem.enc.out_pos").unwrap()); + eprintln!("memory position encoding: max abs error {e:e}"); + assert!(e < 1e-5, "memory position error {e}"); +} diff --git a/tools/sam2_video_reference.py b/tools/sam2_video_reference.py new file mode 100644 index 000000000..6e769d955 --- /dev/null +++ b/tools/sam2_video_reference.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python +"""Dump SAM 2 video reference tensors (CPU, float32) for the Rust port's reference tests. + +Usage: python -I tools/sam2_video_reference.py MODEL_DIR OUT.safetensors +Runs `Sam2VideoModel` on a 7-frame synthetic clip (a textured disc moving right, hidden behind a bar on +frames 4 and 5), prompted with one click on frame 2, then tracks forward (2..6) and backward (2..0). +Saves the preprocessed frames, every frame's low-resolution mask and object score, the stored memory +features and object pointers, and the inputs/outputs of the memory encoder and memory attention. +""" +import sys + +import numpy as np +import torch +import torch.nn.functional as F +from safetensors.torch import save_file +from transformers import Sam2VideoModel +from transformers.models.sam2_video.modeling_sam2_video import Sam2VideoInferenceSession + +T, SIZE, PROMPT_FRAME = 7, 256, 2 + + +def clip(): + rng = np.random.default_rng(3) + yy, xx = np.mgrid[0:SIZE, 0:SIZE].astype(np.float32) + bg = 0.3 + 0.2 * np.sin(xx / 9.0) * np.cos(yy / 7.0) + frames = [] + for t in range(T): + cx, cy = 80 + 12 * t, 128 + disc = np.hypot(xx - cx, yy - cy) < 34 + img = np.stack([bg, bg * 0.8 + 0.1, bg * 0.6 + 0.3], -1) + tex = 0.85 + 0.1 * np.sin((xx - cx) / 5) + img[disc] = np.stack([tex, 0.4 + 0.1 * np.cos((yy - cy) / 5), 0.15 + 0 * xx], -1)[disc] + if t in (4, 5): # a bar hides the disc + img[:, int(cx) - 50 : int(cx) + 50] = 0.2 + img += rng.normal(0, 0.01, img.shape).astype(np.float32) + frames.append((img.clip(0, 1) * 255).astype(np.uint8)) + return frames, (80 + 12 * PROMPT_FRAME, 128) + + +def main(model_dir, out): + torch.manual_seed(0) + torch.set_num_threads(4) + model = Sam2VideoModel.from_pretrained(model_dir).eval() + frames, click = clip() + mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + x = torch.from_numpy(np.stack(frames)).permute(0, 3, 1, 2).float() / 255.0 + x = F.interpolate(x, size=(1024, 1024), mode="bilinear", align_corners=False) + pixels = (x - mean) / std + + taps = {"clip.pixel_values": pixels} + first = {"enc": True, "attn": True} + + def enc_hook(_m, args, kwargs, out): + if first["enc"]: + first["enc"] = False + vf = args[0] if args else kwargs["vision_features"] + masks = args[1] if len(args) > 1 else kwargs["masks"] + taps["mem.enc.vision_features"] = vf.detach().float().clone() + taps["mem.enc.masks"] = masks.detach().float().clone() + taps["mem.enc.out_features"] = out[0].detach().float().clone() + taps["mem.enc.out_pos"] = out[1].detach().float().clone() + + def attn_hook(_m, args, kwargs, out): + if first["attn"]: + first["attn"] = False + for k in ("current_vision_features", "memory", "current_vision_position_embeddings", "memory_posision_embeddings"): + taps[f"mem.attn.{k}"] = kwargs[k].detach().float().clone() + taps["mem.attn.num_pointers"] = torch.tensor([float(kwargs["num_object_pointer_tokens"])]) + taps["mem.attn.out"] = out.detach().float().clone() + + model.memory_encoder.register_forward_hook(enc_hook, with_kwargs=True) + model.memory_attention.register_forward_hook(attn_hook, with_kwargs=True) + + sess = Sam2VideoInferenceSession(video=pixels, video_height=SIZE, video_width=SIZE, dtype=torch.float32) + obj_idx = sess.obj_id_to_idx(1) + scale = 1024.0 / SIZE + sess.add_point_inputs( + obj_idx, + PROMPT_FRAME, + { + "point_coords": torch.tensor([[[[click[0] * scale, click[1] * scale]]]], dtype=torch.float32), + "point_labels": torch.tensor([[[1]]], dtype=torch.int32), + }, + ) + sess.obj_with_new_inputs = [1] + + def record(tag, o): + taps[f"video.{tag}.{o.frame_idx}.pred_masks"] = o.pred_masks.detach().float().clone() + taps[f"video.{tag}.{o.frame_idx}.score"] = o.object_score_logits.detach().float().clone() + + with torch.no_grad(): + record("prompt", model(sess, frame_idx=PROMPT_FRAME)) + for o in model.propagate_in_video_iterator(sess, start_frame_idx=PROMPT_FRAME): + record("fwd", o) + for o in model.propagate_in_video_iterator(sess, start_frame_idx=PROMPT_FRAME, reverse=True): + record("rev", o) + + d = sess.output_dict_per_obj[obj_idx] + for kind in ("cond_frame_outputs", "non_cond_frame_outputs"): + for f, o in d[kind].items(): + taps[f"state.{f}.maskmem_features"] = o["maskmem_features"].detach().float().clone() + taps[f"state.{f}.maskmem_pos_enc"] = o["maskmem_pos_enc"].detach().float().clone() + taps[f"state.{f}.object_pointer"] = o["object_pointer"].detach().float().clone() + save_file({k: v.contiguous() for k, v in taps.items()}, out) + print(f"wrote {len(taps)} tensors to {out}") + + +if __name__ == "__main__": + if len(sys.argv) != 3: + sys.exit(__doc__) + main(sys.argv[1], sys.argv[2]) From ea2899d0771be2c6f814a1d269a1d16f52e6151e Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:30:53 +0200 Subject: [PATCH 11/77] S2.22: memory attention with rotary encoding matches PyTorch Self and cross attention with axial 2-D RoPE (keys tiled over memory frames, object-pointer tokens left unrotated). Max abs error 2.6e-5. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/src/lib.rs | 1 + crates/sam2/src/memattn.rs | 173 +++++++++++++++++++++++++++++++++ crates/sam2/tests/reference.rs | 21 ++++ 3 files changed, 195 insertions(+) create mode 100644 crates/sam2/src/memattn.rs diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 18fe5a8ae..53156c7ab 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 memattn; pub mod memory; pub mod model; pub mod neck; diff --git a/crates/sam2/src/memattn.rs b/crates/sam2/src/memattn.rs new file mode 100644 index 000000000..c3bfc0893 --- /dev/null +++ b/crates/sam2/src/memattn.rs @@ -0,0 +1,173 @@ +// Modified work: ported from Hugging Face `transformers` `models/sam2_video/modeling_sam2_video.py` +// (Sam2VideoMemoryAttention, Sam2VideoMemoryAttentionLayer, Sam2VideoRoPEAttention, +// Sam2VideoVisionRotaryEmbedding; Apache-2.0, Copyright The HuggingFace Inc. team) to Rust on +// candle. See NOTICE. +//! Memory attention: the current frame's features attend to themselves and to the memory bank +//! (memory frames plus object pointers) with axial 2-D rotary position encoding. + +use candle_core::{D, Device, Tensor}; +use candle_nn::{Linear, ops::softmax_last_dim}; + +use crate::nn::{LayerNorm, apply, layer_norm, linear}; +use crate::weights::Weights; +use crate::{Error, Result}; + +const DIM: usize = 256; +const LAYERS: usize = 4; +/// Side of the feature grid the rotary encoding is defined on. +pub const GRID: usize = 64; + +/// `cos` and `sin` of the axial rotary encoding for the `GRID x GRID` grid, each `[GRID², 256]`. +fn rope_tables(theta: f32) -> Result<(Tensor, Tensor)> { + let freqs = DIM / 2 / 2; // per axis: head_dim / 4 + let inv: Vec = (0..freqs).map(|i| 1.0 / theta.powf(i as f32 / freqs as f32)).collect(); + let n = GRID * GRID; + let (mut cos, mut sin) = (vec![0f32; n * DIM], vec![0f32; n * DIM]); + for p in 0..n { + let (x, y) = ((p % GRID) as f32, (p / GRID) as f32); + for j in 0..DIM / 2 { + let f = if j < freqs { x * inv.get(j).copied().unwrap_or(0.0) } else { y * inv.get(j - freqs).copied().unwrap_or(0.0) }; + for k in 0..2 { + if let (Some(c), Some(s)) = (cos.get_mut(p * DIM + 2 * j + k), sin.get_mut(p * DIM + 2 * j + k)) { + *c = f.cos(); + *s = f.sin(); + } + } + } + } + Ok((Tensor::from_vec(cos, (n, DIM), &Device::Cpu)?, Tensor::from_vec(sin, (n, DIM), &Device::Cpu)?)) +} + +/// `(x0, x1) -> (-x1, x0)` over adjacent pairs of the last dimension. +fn rotate_pairwise(x: &Tensor) -> Result { + let (n, d) = x.dims2()?; + let pairs = x.reshape((n, d / 2, 2))?; + let (x1, x2) = (pairs.narrow(2, 0, 1)?.squeeze(2)?, pairs.narrow(2, 1, 1)?.squeeze(2)?); + Ok(Tensor::stack(&[&x2.neg()?, &x1], D::Minus1)?.reshape((n, d))?) +} + +fn rope(x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result { + Ok(((x * cos)? + (rotate_pairwise(x)? * sin)?)?) +} + +struct RopeAttention { + q: Linear, + k: Linear, + v: Linear, + o: Linear, +} + +impl RopeAttention { + fn load(w: &Weights, prefix: &str, kv_dim: usize) -> Result { + Ok(RopeAttention { + q: linear(w, &format!("{prefix}.q_proj"), DIM, DIM)?, + k: linear(w, &format!("{prefix}.k_proj"), kv_dim, DIM)?, + v: linear(w, &format!("{prefix}.v_proj"), kv_dim, DIM)?, + o: linear(w, &format!("{prefix}.o_proj"), DIM, DIM)?, + }) + } + + /// `query`: `[N, 256]`, `key`/`value`: `[M, kv]`. The first `M - exclude` keys are rotated + /// (their table is `cos`/`sin` tiled to length); the last `exclude` (object pointers) are not. + fn forward(&self, query: &Tensor, key: &Tensor, value: &Tensor, tables: &(Tensor, Tensor), exclude: usize) -> Result { + let (cos, sin) = tables; + let q = rope(&apply(&self.q, query)?, cos, sin)?; + let k = apply(&self.k, key)?; + let v = apply(&self.v, value)?; + let m = k.dim(0)?; + let rotated = m.checked_sub(exclude).ok_or_else(|| Error::Input("more object pointers than memory tokens".into()))?; + let k = if rotated == 0 { + k + } else { + let n = cos.dim(0)?; + let reps = if rotated == n { + 1 + } else if rotated % n == 0 { + rotated / n + } else { + return Err(Error::Input(format!("memory length {rotated} is not a multiple of {n}"))); + }; + let tile = |t: &Tensor| -> Result { Ok(Tensor::cat(&vec![t; reps], 0)?) }; + let k_rot = rope(&k.narrow(0, 0, rotated)?, &tile(cos)?, &tile(sin)?)?; + if exclude == 0 { k_rot } else { Tensor::cat(&[&k_rot, &k.narrow(0, rotated, exclude)?], 0)? } + }; + let scale = (DIM as f64).powf(-0.5); + let att = softmax_last_dim(&((q * scale)?.matmul(&k.t()?.contiguous()?)?))?; + apply(&self.o, &att.matmul(&v)?) + } +} + +struct Layer { + self_attn: RopeAttention, + cross: RopeAttention, + ln1: LayerNorm, + ln2: LayerNorm, + ln3: LayerNorm, + lin1: Linear, + lin2: Linear, +} + +pub struct MemoryAttention { + layers: Vec, + final_ln: LayerNorm, + tables: (Tensor, Tensor), +} + +impl MemoryAttention { + pub fn load(w: &Weights) -> Result { + let layers = (0..LAYERS) + .map(|i| { + let p = format!("memory_attention.layers.{i}"); + Ok(Layer { + self_attn: RopeAttention::load(w, &format!("{p}.self_attn"), DIM)?, + cross: RopeAttention::load(w, &format!("{p}.cross_attn_image"), 64)?, + ln1: layer_norm(w, &format!("{p}.layer_norm1"), DIM, 1e-5)?, + ln2: layer_norm(w, &format!("{p}.layer_norm2"), DIM, 1e-5)?, + ln3: layer_norm(w, &format!("{p}.layer_norm3"), DIM, 1e-5)?, + lin1: linear(w, &format!("{p}.linear1"), DIM, 2048)?, + lin2: linear(w, &format!("{p}.linear2"), 2048, DIM)?, + }) + }) + .collect::>>()?; + Ok(MemoryAttention { layers, final_ln: layer_norm(w, "memory_attention.layer_norm", DIM, 1e-5)?, tables: rope_tables(10000.0)? }) + } + + /// Condition the current frame on memory. + /// + /// `current` and `current_pos`: `[4096, 256]` (the lowest-resolution neck level and its + /// position encoding, one row per cell); `memory` and `memory_pos`: `[M, 64]` (memory frames + /// followed by `num_pointers` object-pointer tokens). Returns `[4096, 256]`. + pub fn forward(&self, current: &Tensor, current_pos: &Tensor, memory: &Tensor, memory_pos: &Tensor, num_pointers: usize) -> Result { + if current.dims() != [GRID * GRID, DIM] { + return Err(Error::Input(format!("current features must be [{}, {DIM}], got {:?}", GRID * GRID, current.dims()))); + } + let keys_pos = (memory + memory_pos)?; + let mut out = (current + (current_pos * 0.1)?)?; + for l in &self.layers { + let q = l.ln1.forward(&out)?; + out = (&out + l.self_attn.forward(&q, &q, &q, &self.tables, 0)?)?; + let q = l.ln2.forward(&out)?; + out = (&out + l.cross.forward(&q, &keys_pos, memory, &self.tables, num_pointers)?)?; + let q = l.ln3.forward(&out)?; + out = (&out + apply(&l.lin2, &apply(&l.lin1, &q)?.relu()?)?)?; + } + self.final_ln.forward(&out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rotation_is_pairwise_and_tables_start_at_identity() { + let x = Tensor::new(&[[1f32, 2.0, 3.0, 4.0]], &Device::Cpu).unwrap(); + let r = rotate_pairwise(&x).unwrap().to_vec2::().unwrap(); + assert_eq!(r, vec![vec![-2.0, 1.0, -4.0, 3.0]]); + let (cos, sin) = rope_tables(10000.0).unwrap(); + // cell (0, 0): angle 0 everywhere + let c = cos.narrow(0, 0, 1).unwrap().to_vec2::().unwrap(); + let s = sin.narrow(0, 0, 1).unwrap().to_vec2::().unwrap(); + assert!(c[0].iter().all(|v| *v == 1.0) && s[0].iter().all(|v| *v == 0.0)); + } +} diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index f90d03ab8..d6d457871 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -173,3 +173,24 @@ fn memory_encoder_matches_pytorch() { eprintln!("memory position encoding: max abs error {e:e}"); assert!(e < 1e-5, "memory position error {e}"); } + +#[test] +fn memory_attention_matches_pytorch() { + let Some((w, r)) = load_video() else { return }; + let att = effectcraft_sam2::memattn::MemoryAttention::load(&w).unwrap(); + let get = |k: &str| r.get(k).unwrap().squeeze(1).unwrap(); + let pointers = r.get("mem.attn.num_pointers").unwrap().to_vec1::().unwrap()[0] as usize; + let out = att + .forward( + &get("mem.attn.current_vision_features"), + &get("mem.attn.current_vision_position_embeddings"), + &get("mem.attn.memory"), + &get("mem.attn.memory_posision_embeddings"), + pointers, + ) + .unwrap(); + let want = r.get("mem.attn.out").unwrap().squeeze(0).unwrap().squeeze(0).unwrap(); + let e = max_abs(&out, &want); + eprintln!("memory attention: max abs error {e:e} (pointers {pointers})"); + assert!(e < MEMORY_TOL, "memory attention error {e}"); +} From 5b0cc57f87670035eef83439e08146ed959a7b37 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:35:01 +0200 Subject: [PATCH 12/77] S2.23: video tracker session Per-object memory banks, prompts on any frame, forward and backward propagation with streaming results, cancellation, object pointers with temporal encoding, pruning of memories no later frame can use, and checked input handling. Hostile-input and cancel tests. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/src/lib.rs | 1 + crates/sam2/src/model.rs | 10 + crates/sam2/src/prompt.rs | 8 + crates/sam2/src/session.rs | 363 +++++++++++++++++++++++++++++++++ crates/sam2/tests/reference.rs | 53 +++++ 5 files changed, 435 insertions(+) create mode 100644 crates/sam2/src/session.rs diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 53156c7ab..4cb6f6f3e 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -19,6 +19,7 @@ pub mod model; pub mod neck; pub mod nn; pub mod prompt; +pub mod session; pub mod weights; /// Everything that can go wrong in this crate. Never a panic. diff --git a/crates/sam2/src/model.rs b/crates/sam2/src/model.rs index 13481fdb9..c30cbb5b9 100644 --- a/crates/sam2/src/model.rs +++ b/crates/sam2/src/model.rs @@ -71,6 +71,16 @@ impl ImageModel { self.decoder.forward(&image, &self.image_pos, &sparse, &dense, &enc.high, multimask) } + /// Decode against an already memory-conditioned embedding `[1, 256, g, g]` (video tracking). + pub fn decode_embedded(&self, image: &Tensor, high: &(Tensor, Tensor), sparse: &Tensor, dense: &Tensor, multimask: bool) -> Result { + self.decoder.forward(image, &self.image_pos, sparse, dense, high, multimask) + } + + /// The learned embedding added to a frame's features when there is no memory yet. + pub fn no_memory(&self) -> &Tensor { + &self.no_memory + } + pub fn device(&self) -> Device { Device::Cpu } diff --git a/crates/sam2/src/prompt.rs b/crates/sam2/src/prompt.rs index 55dd901a2..51aade1c7 100644 --- a/crates/sam2/src/prompt.rs +++ b/crates/sam2/src/prompt.rs @@ -146,6 +146,14 @@ impl PromptEncoder { Ok(Tensor::from_vec(tokens, (1, n, DIM), &Device::Cpu)?) } + /// Tokens for "no prompt" on a tracked frame: the reference feeds one point labelled -1 and + /// then pads another, so two `not_a_point` tokens, `[1, 2, 256]`. + pub fn sparse_none(&self) -> Result { + let mut v = self.not_a_point.clone(); + v.extend_from_slice(&self.not_a_point); + Ok(Tensor::from_vec(v, (1, 2, DIM), &Device::Cpu)?) + } + /// The dense embedding for "no mask prompt", `[1, 256, GRID, GRID]`. pub fn dense_empty(&self) -> Result { Ok(self.no_mask.reshape((1, DIM, 1, 1))?.broadcast_as((1, DIM, GRID, GRID))?.contiguous()?) diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs new file mode 100644 index 000000000..45d1ffd0a --- /dev/null +++ b/crates/sam2/src/session.rs @@ -0,0 +1,363 @@ +// Modified work: ported from Hugging Face `transformers` `models/sam2_video/modeling_sam2_video.py` +// (Sam2VideoModel and Sam2VideoInferenceSession: memory-conditioned single-frame inference, +// memory bank selection, object pointers, propagation; Apache-2.0, Copyright The HuggingFace +// Inc. team) to Rust on candle. See NOTICE. +//! The video tracker: per-object memory banks, prompts on any frame, propagation forwards and +//! backwards with streaming results and cancellation. + +use std::collections::{BTreeMap, VecDeque}; +use std::path::Path; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use candle_core::{DType, Tensor}; +use candle_nn::{Linear, ops::sigmoid}; + +use crate::decoder::FeedForward; +use crate::memattn::MemoryAttention; +use crate::memory::{Memory, MemoryEncoder}; +use crate::model::ImageModel; +use crate::nn::{apply, linear}; +use crate::prompt::{INPUT_SIZE, Prompt}; +use crate::weights::Weights; +use crate::{Error, Result}; + +/// Mask logit used where the object is absent (the reference's `NO_OBJ_SCORE`). +const NO_OBJECT: f32 = -1024.0; +/// Frames of memory the attention looks back over. +const MEMORY_FRAMES: usize = 7; +/// Object pointers kept per query. +const MAX_POINTERS: usize = 16; +/// Largest clip the session accepts. +const MAX_FRAMES: usize = 1_000_000; +const GRID: usize = 64; + +/// Frames of the clip, already resized to the model input and normalised: `[1, 3, 1024, 1024]`. +pub trait FrameSource: Send { + fn num_frames(&self) -> usize; + fn frame(&self, index: usize) -> Result; +} + +/// Cooperative cancellation shared with the caller. +#[derive(Clone, Default)] +pub struct Cancel(Arc); + +impl Cancel { + pub fn new() -> Cancel { + Cancel::default() + } + pub fn cancel(&self) { + self.0.store(true, Ordering::Relaxed); + } + pub fn is_cancelled(&self) -> bool { + self.0.load(Ordering::Relaxed) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ObjectId(pub u32); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Direction { + Forward, + Backward, +} + +/// One object's mask on one frame. +#[derive(Clone, Debug, PartialEq)] +pub struct FrameMask { + pub frame: usize, + pub object: ObjectId, + /// Mask logits at 256 x 256 (positive is inside the object); `-1024` everywhere when the + /// object is absent. + pub logits: Vec, + /// The model's own estimate of the mask quality (0..1). + pub iou: f32, + /// Object-presence logit; at or below zero the object is not in the frame. + pub object_score: f32, + pub occluded: bool, +} + +/// The weights of everything the video tracker needs. +pub struct VideoModel { + pub image: ImageModel, + memory_encoder: MemoryEncoder, + memory_attention: MemoryAttention, + temporal_pos: Vec, + no_object_pointer: Tensor, + pointer_proj: FeedForward, + pointer_time_proj: Linear, +} + +impl VideoModel { + pub fn load(w: &Weights) -> Result { + Ok(VideoModel { + image: ImageModel::load(w)?, + memory_encoder: MemoryEncoder::load(w)?, + memory_attention: MemoryAttention::load(w)?, + temporal_pos: w.get("memory_temporal_positional_encoding", &[MEMORY_FRAMES, 1, 1, 64])?.flatten_all()?.to_vec1::()?, + no_object_pointer: w.get("no_object_pointer", &[1, 256])?, + pointer_proj: FeedForward::load(w, "object_pointer_proj", 256, 256, 256, 3, false)?, + pointer_time_proj: linear(w, "temporal_positional_encoding_projection_layer", 256, 64)?, + }) + } + + pub fn open(path: &Path) -> Result { + VideoModel::load(&Weights::open(path)?) + } +} + +/// What the tracker keeps for a frame it has processed. +struct FrameOut { + memory: Memory, + /// `[1, 256]` + pointer: Tensor, + /// The mask of a conditioning frame (returned again, unchanged, when propagation passes it). + result: Option, +} + +#[derive(Default)] +struct ObjectState { + cond: BTreeMap, + non_cond: BTreeMap, +} + +struct FrameFeatures { + /// Lowest-resolution level `[1, 256, 64, 64]` and its position encoding. + low: Tensor, + pos: Tensor, + high: (Tensor, Tensor), +} + +pub struct VideoTracker { + model: Arc, + frames: Box, + objects: BTreeMap, + cache: VecDeque<(usize, Arc)>, +} + +fn sine_embedding_1d(positions: &[f32], dim: usize) -> Result { + let half = dim / 2; + let dim_t: Vec = (0..half).map(|i| 10000f32.powf(2.0 * (i / 2) as f32 / half as f32)).collect(); + let mut out = Vec::with_capacity(positions.len() * dim); + for p in positions { + out.extend(dim_t.iter().map(|d| (p / d).sin())); + out.extend(dim_t.iter().map(|d| (p / d).cos())); + } + Ok(Tensor::from_vec(out, (positions.len(), dim), &candle_core::Device::Cpu)?) +} + +fn scalar(t: &Tensor) -> Result { + Ok(t.flatten_all()?.to_vec1::()?.first().copied().unwrap_or(0.0)) +} + +impl VideoTracker { + pub fn new(model: Arc, frames: Box) -> Result { + let n = frames.num_frames(); + if n == 0 || n > MAX_FRAMES { + return Err(Error::Input(format!("a clip needs between 1 and {MAX_FRAMES} frames, got {n}"))); + } + Ok(VideoTracker { model, frames, objects: BTreeMap::new(), cache: VecDeque::new() }) + } + + pub fn num_frames(&self) -> usize { + self.frames.num_frames() + } + + /// Forget every object and cached frame (the clip stays). + pub fn reset(&mut self) { + self.objects.clear(); + self.cache.clear(); + } + + fn features(&mut self, frame: usize) -> Result> { + if let Some((_, f)) = self.cache.iter().find(|(i, _)| *i == frame) { + return Ok(f.clone()); + } + let pixels = self.frames.frame(frame)?; + let enc = self.model.image.encode(&pixels)?; + let (low, pos) = match (enc.features.features.get(2), enc.features.positions.get(2)) { + (Some(l), Some(p)) => (l.clone(), p.clone()), + _ => return Err(Error::Compute("vision encoder returned too few levels".into())), + }; + let f = Arc::new(FrameFeatures { low, pos, high: enc.high }); + self.cache.push_back((frame, f.clone())); + while self.cache.len() > 2 { + self.cache.pop_front(); + } + Ok(f) + } + + /// The memory bank for `frame`: memory frames and object pointers with their position + /// encodings, plus how many trailing tokens are pointers. + fn memory_bank(&self, obj: &ObjectState, frame: usize, reverse: bool) -> Result<(Tensor, Tensor, usize)> { + let (mut feats, mut poss) = (vec![], vec![]); + let temporal = |off: usize| -> Result { + // The reference indexes with `offset - 1`, so the conditioning frames (offset 0) use the last row. + let row = if off == 0 { MEMORY_FRAMES - 1 } else { off - 1 }; + let v = self.model.temporal_pos.get(row * 64..(row + 1) * 64).ok_or_else(|| Error::Weights("temporal position table".into()))?; + Ok(Tensor::from_slice(v, (1, 64), &candle_core::Device::Cpu)?) + }; + for out in obj.cond.values() { + feats.push(out.memory.features.clone()); + poss.push(out.memory.pos.broadcast_add(&temporal(0)?)?); + } + for off in (1..MEMORY_FRAMES).rev() { + let prev = if reverse { frame.checked_add(off) } else { frame.checked_sub(off) }; + if let Some(out) = prev.and_then(|p| obj.non_cond.get(&p)) { + feats.push(out.memory.features.clone()); + poss.push(out.memory.pos.broadcast_add(&temporal(off)?)?); + } + } + // Object pointers: conditioning frames on the already-tracked side, then recent frames. + let n_frames = self.frames.num_frames(); + let max_ptrs = n_frames.min(MAX_POINTERS); + let (mut offsets, mut pointers): (Vec, Vec) = (vec![], vec![]); + let sign = if reverse { -1.0f32 } else { 1.0 }; + for (t, out) in &obj.cond { + if if reverse { *t >= frame } else { *t <= frame } { + offsets.push((frame as f32 - *t as f32) * sign); + pointers.push(out.pointer.clone()); + } + } + for off in 1..max_ptrs { + let r = if reverse { frame.checked_add(off) } else { frame.checked_sub(off) }; + let Some(r) = r.filter(|r| *r < n_frames) else { break }; + if let Some(out) = obj.non_cond.get(&r) { + offsets.push(off as f32); + pointers.push(out.pointer.clone()); + } + } + let mut n_pointer_tokens = 0; + if !pointers.is_empty() { + let max_diff = (max_ptrs.saturating_sub(1)).max(1) as f32; + let norm: Vec = offsets.iter().map(|o| o / max_diff).collect(); + let pe = apply(&self.model.pointer_time_proj, &sine_embedding_1d(&norm, 256)?)?; + let ptrs = Tensor::cat(&pointers, 0)?; // [P, 256] + let p = ptrs.dim(0)?; + let tokens = ptrs.reshape((p * 4, 64))?; + let pe = pe.unsqueeze(1)?.broadcast_as((p, 4, 64))?.reshape((p * 4, 64))?; + n_pointer_tokens = p * 4; + feats.push(tokens); + poss.push(pe.contiguous()?); + } + if feats.is_empty() { + return Err(Error::Input("the object has no conditioning frame; add a prompt first".into())); + } + Ok((Tensor::cat(&feats, 0)?, Tensor::cat(&poss, 0)?, n_pointer_tokens)) + } + + /// Run the model on one frame for one object and store the result. + fn step(&mut self, id: ObjectId, frame: usize, prompt: Option<&Prompt>, reverse: bool, conditioning: bool) -> Result { + let feats = self.features(frame)?; + let model = self.model.clone(); + let tokens = |t: &Tensor| -> Result { Ok(t.flatten_from(2)?.squeeze(0)?.t()?.contiguous()?) }; + let conditioned = if conditioning { + feats.low.broadcast_add(model.image.no_memory())? + } else { + let obj = self.objects.get(&id).ok_or_else(|| Error::Input("unknown object".into()))?; + let (mem, mem_pos, n_ptr) = self.memory_bank(obj, frame, reverse)?; + let out = model.memory_attention.forward(&tokens(&feats.low)?, &tokens(&feats.pos)?, &mem, &mem_pos, n_ptr)?; + out.t()?.contiguous()?.reshape((1, 256, GRID, GRID))? + }; + let num_points = prompt.map_or(0, |p| p.points.len() + 2 * usize::from(p.bbox.is_some())); + let multimask = num_points <= 1; + let sparse = match prompt { + Some(p) => model.image.prompt.sparse(p)?, + None => model.image.prompt.sparse_none()?, + }; + let dense = model.image.prompt.dense_empty()?; + let out = model.image.decode_embedded(&conditioned, &feats.high, &sparse, &dense, multimask)?; + + let score = scalar(&out.object_score)?; + let present = score > 0.0; + let ious = out.iou.flatten_all()?.to_vec1::()?; + let best = if multimask { ious.iter().enumerate().fold((0usize, f32::NEG_INFINITY), |b, (i, v)| if *v > b.1 { (i, *v) } else { b }).0 } else { 0 }; + let low = out.masks.narrow(1, best, 1)?; + let low = if present { low } else { Tensor::full(NO_OBJECT, low.dims(), low.device())? }; + let token = if multimask { out.sam_tokens.narrow(1, best, 1)? } else { out.sam_tokens.narrow(1, 0, 1)? }.squeeze(1)?; + let pointer = if present { model.pointer_proj.forward(&token)? } else { model.no_object_pointer.clone() }; + + let high = low.upsample_bilinear2d(INPUT_SIZE, INPUT_SIZE, false)?; + let for_memory = if prompt.is_some() { high.gt(0f32)?.to_dtype(DType::F32)? } else { sigmoid(&high)? }; + let for_memory = ((for_memory * 20.0)? - 10.0)?; + let memory = model.memory_encoder.memory(&feats.low, &for_memory, present)?; + + let iou = ious.get(best).copied().unwrap_or(0.0); + let result = FrameMask { frame, object: id, logits: low.flatten_all()?.to_vec1::()?, iou, object_score: score, occluded: !present }; + let state = self.objects.entry(id).or_default(); + let entry = FrameOut { memory, pointer, result: conditioning.then(|| result.clone()) }; + if conditioning { + state.non_cond.remove(&frame); + state.cond.insert(frame, entry); + } else { + state.non_cond.insert(frame, entry); + Self::prune(state, frame); + } + Ok(result) + } + + /// Drop memories that no later frame in either direction can still look at. + fn prune(state: &mut ObjectState, frame: usize) { + let keep = |f: usize, cond: &BTreeMap| f.abs_diff(frame) <= MAX_POINTERS || cond.keys().any(|c| f.abs_diff(*c) <= MEMORY_FRAMES); + let doomed: Vec = state.non_cond.keys().copied().filter(|f| !keep(*f, &state.cond)).collect(); + for f in doomed { + state.non_cond.remove(&f); + } + } + + fn check_frame(&self, frame: usize) -> Result<()> { + if frame >= self.frames.num_frames() { + return Err(Error::Input(format!("frame {frame} is outside the clip of {} frames", self.frames.num_frames()))); + } + Ok(()) + } + + /// Prompt an object on a frame (clicks and/or a box). The frame becomes one of the object's + /// conditioning frames when it has not been tracked yet; the mask comes back at once. + pub fn prompt(&mut self, frame: usize, object: ObjectId, prompt: &Prompt) -> Result { + self.check_frame(frame)?; + prompt.validate()?; + if prompt.points.is_empty() && prompt.bbox.is_none() { + return Err(Error::Input("a prompt needs at least one point or a box".into())); + } + let tracked = self.objects.get(&object).is_some_and(|o| o.non_cond.contains_key(&frame)); + if tracked { + // A correction on a tracked frame is conditioned on the memory like any tracked frame. + let reverse = false; + return self.step(object, frame, Some(prompt), reverse, false); + } + self.objects.entry(object).or_default(); + self.step(object, frame, Some(prompt), false, true) + } + + /// Track every object from `start` through `max_frames` frames (the rest of the clip when + /// `None`) in `direction`, calling `out` for each object on each frame as it is finished. + pub fn propagate(&mut self, start: usize, direction: Direction, max_frames: Option, cancel: &Cancel, out: &mut dyn FnMut(FrameMask)) -> Result<()> { + self.check_frame(start)?; + let n = self.frames.num_frames(); + let reverse = direction == Direction::Backward; + let span = max_frames.unwrap_or(n).min(n); + let ids: Vec = self.objects.iter().filter(|(_, o)| !o.cond.is_empty()).map(|(id, _)| *id).collect(); + if ids.is_empty() { + return Err(Error::Input("no object has been prompted".into())); + } + for k in 0..=span { + let frame = if reverse { start.checked_sub(k) } else { Some(start + k).filter(|f| *f < n) }; + let Some(frame) = frame else { break }; + for id in &ids { + if cancel.is_cancelled() { + return Err(Error::Cancelled); + } + let cached = self.objects.get(id).is_some_and(|o| o.cond.contains_key(&frame)); + let stored = self.objects.get(id).and_then(|o| o.cond.get(&frame)).and_then(|o| o.result.clone()); + let mask = match stored { + Some(m) if cached => m, + _ => self.step(*id, frame, None, reverse, false)?, + }; + out(mask); + } + } + Ok(()) + } +} diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index d6d457871..7e480af10 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -194,3 +194,56 @@ fn memory_attention_matches_pytorch() { eprintln!("memory attention: max abs error {e:e} (pointers {pointers})"); assert!(e < MEMORY_TOL, "memory attention error {e}"); } + +/// Frames of the video fixture, served one at a time. +struct Clip(Tensor); + +impl effectcraft_sam2::session::FrameSource for Clip { + fn num_frames(&self) -> usize { + self.0.dim(0).unwrap() + } + fn frame(&self, i: usize) -> effectcraft_sam2::Result { + Ok(self.0.narrow(0, i, 1)?.contiguous()?) + } +} + +#[test] +fn tracker_rejects_hostile_input_and_honours_cancel() { + use effectcraft_sam2::Error; + use effectcraft_sam2::prompt::Prompt; + use effectcraft_sam2::session::{Cancel, Direction, ObjectId, VideoModel, VideoTracker}; + let Some((w, r)) = load_video() else { return }; + let model = std::sync::Arc::new(VideoModel::load(&w).unwrap()); + let clip = r.get("clip.pixel_values").unwrap().narrow(0, 0, 3).unwrap(); + let mut t = VideoTracker::new(model.clone(), Box::new(Clip(clip.clone()))).unwrap(); + let click = Prompt { points: vec![([416.0, 512.0], true)], bbox: None }; + let obj = ObjectId(7); + // outside the clip, empty and non-finite prompts, nothing prompted yet + assert!(matches!(t.prompt(3, obj, &click), Err(Error::Input(_)))); + assert!(matches!(t.prompt(0, obj, &Prompt::default()), Err(Error::Input(_)))); + assert!(matches!(t.prompt(0, obj, &Prompt { points: vec![([f32::NAN, 1.0], true)], bbox: None }), Err(Error::Input(_)))); + let mut sink = |_m| {}; + assert!(matches!(t.propagate(0, Direction::Forward, None, &Cancel::new(), &mut sink), Err(Error::Input(_)))); + assert!(matches!(t.propagate(9, Direction::Forward, None, &Cancel::new(), &mut sink), Err(Error::Input(_)))); + // an empty clip is refused; a cancelled run stops with Cancelled and keeps the tracker usable + assert!(VideoTracker::new(model.clone(), Box::new(Clip(clip.narrow(0, 0, 0).unwrap()))).is_err()); + t.prompt(0, obj, &click).unwrap(); + let cancel = Cancel::new(); + cancel.cancel(); + assert!(matches!(t.propagate(0, Direction::Forward, None, &cancel, &mut sink), Err(Error::Cancelled))); + let mut seen = vec![]; + t.propagate(0, Direction::Forward, Some(1), &Cancel::new(), &mut |m| seen.push(m.frame)).unwrap(); + assert_eq!(seen, vec![0, 1]); + // a frame source that returns the wrong size is an error, not a crash + struct Bad; + impl effectcraft_sam2::session::FrameSource for Bad { + fn num_frames(&self) -> usize { + 2 + } + fn frame(&self, _i: usize) -> effectcraft_sam2::Result { + Ok(Tensor::zeros((1, 3, 64, 64), candle_core::DType::F32, &Device::Cpu)?) + } + } + let mut bad = VideoTracker::new(model, Box::new(Bad)).unwrap(); + assert!(matches!(bad.prompt(0, obj, &click), Err(Error::Input(_)))); +} From d1bff2baf6ec093a39be96b333003d7d0560ce4a Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:35:07 +0200 Subject: [PATCH 13/77] S2.24: video tracking matches Sam2VideoModel on a synthetic clip Click on frame 2 of a 7-frame clip with an occlusion, tracked forwards and backwards: IoU 1.0000 on all 9 frame results, worst mask-logit error 0.0038, object scores equal to 3 decimals. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/tests/reference.rs | 54 ++++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index 7e480af10..70af32cd0 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -207,6 +207,60 @@ impl effectcraft_sam2::session::FrameSource for Clip { } } +fn iou(a: &[f32], b: &[f32]) -> f64 { + let (mut i, mut u) = (0usize, 0usize); + for (x, y) in a.iter().zip(b) { + let (x, y) = (*x > 0.0, *y > 0.0); + i += usize::from(x && y); + u += usize::from(x || y); + } + if u == 0 { 1.0 } else { i as f64 / u as f64 } +} + +/// Through the clip the memory bank feeds back on itself (bfloat16 memories, rotary attention), so +/// the bound on logits is looser than for a single stage; the masks themselves must agree. +/// Measured max 0.0038 over 9 frames forwards and backwards, 2026-10-15. +const VIDEO_LOGIT_TOL: f32 = 0.02; +const VIDEO_MIN_IOU: f64 = 0.995; + +#[test] +fn video_tracking_matches_sam2_video_model() { + use effectcraft_sam2::prompt::Prompt; + use effectcraft_sam2::session::{Cancel, Direction, ObjectId, VideoModel, VideoTracker}; + let Some((w, r)) = load_video() else { return }; + let model = std::sync::Arc::new(VideoModel::load(&w).unwrap()); + let clip = r.get("clip.pixel_values").unwrap().clone(); + let mut tracker = VideoTracker::new(model, Box::new(Clip(clip))).unwrap(); + let obj = ObjectId(1); + let first = tracker.prompt(2, obj, &Prompt { points: vec![([416.0, 512.0], true)], bbox: None }).unwrap(); + + let check = |tag: &str, m: &effectcraft_sam2::session::FrameMask| { + let want = r.get(&format!("video.{tag}.{}.pred_masks", m.frame)).unwrap().flatten_all().unwrap().to_vec1::().unwrap(); + let want_score = r.get(&format!("video.{tag}.{}.score", m.frame)).unwrap().flatten_all().unwrap().to_vec1::().unwrap()[0]; + let err = m.logits.iter().zip(&want).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max); + let j = iou(&m.logits, &want); + eprintln!("{tag} frame {}: logit error {err:.4}, IoU {j:.4}, score {:.3} (want {want_score:.3})", m.frame, m.object_score); + assert!(j >= VIDEO_MIN_IOU, "{tag} frame {} IoU {j}", m.frame); + assert!(err < VIDEO_LOGIT_TOL, "{tag} frame {} logit error {err}", m.frame); + assert!((m.object_score - want_score).abs() < 0.05, "{tag} frame {} score {} vs {want_score}", m.frame, m.object_score); + assert_eq!(m.occluded, want_score <= 0.0); + }; + check("prompt", &first); + let cancel = Cancel::new(); + let mut fwd = vec![]; + tracker.propagate(2, Direction::Forward, None, &cancel, &mut |m| fwd.push(m)).unwrap(); + assert_eq!(fwd.iter().map(|m| m.frame).collect::>(), vec![2, 3, 4, 5, 6]); + for m in &fwd { + check("fwd", m); + } + let mut rev = vec![]; + tracker.propagate(2, Direction::Backward, None, &cancel, &mut |m| rev.push(m)).unwrap(); + assert_eq!(rev.iter().map(|m| m.frame).collect::>(), vec![2, 1, 0]); + for m in &rev { + check("rev", m); + } +} + #[test] fn tracker_rejects_hostile_input_and_honours_cancel() { use effectcraft_sam2::Error; From e5842c0c923621d5b1a2d06408624ef380fc1923 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:37:17 +0200 Subject: [PATCH 14/77] S2.30: SAM 2.1 Tiny in the model registry Apache-2.0 entry pinned to a Hugging Face commit with its SHA-256 and size, listed in Settings > Roto Brush and installable with the existing verified install flow. No weights are bundled. Co-Authored-By: Claude Sonnet 5.5 --- crates/engine/src/tests/models.rs | 4 +++- crates/segment/README.md | 1 + crates/segment/src/lib.rs | 23 ++++++++++++++++++++++- docs/agents.md | 2 +- 4 files changed, 27 insertions(+), 3 deletions(-) diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index 21d316e75..786770ec5 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -36,7 +36,9 @@ fn models_are_listed_chosen_and_verified() { for (prefix, model, task) in [("roto", "mobilesam", Task::Mask), ("face", "mediapipe-face", Task::Face)] { let list = s.execute_checked(&format!("{prefix}.models"), json!({})).unwrap(); let ids: Vec<&str> = list["models"].as_array().unwrap().iter().map(|m| m["id"].as_str().unwrap()).collect(); - assert_eq!(ids, ["classical", model]); + // Roto Brush also offers SAM 2.1 Tiny (a video tracker; see crates/sam2). + let expect: Vec<&str> = if task == Task::Mask { vec!["classical", model, "sam2.1-hiera-tiny"] } else { vec!["classical", model] }; + assert_eq!(ids, expect); let m = &list["models"][1]; assert_eq!((m["licence"].as_str(), m["installed"].as_bool()), (Some("Apache-2.0"), Some(false))); assert!(m["url"].as_str().unwrap().starts_with("https://") && !m["authors"].as_str().unwrap().is_empty()); diff --git a/crates/segment/README.md b/crates/segment/README.md index 79b1d996b..4962f8ef8 100644 --- a/crates/segment/README.md +++ b/crates/segment/README.md @@ -36,6 +36,7 @@ kernels at run time). No ML framework, no unsafe code. | id | Task | Model | Authors | Licence | File | Size | |---|---|---|---|---|---|---| | `mobilesam` | Roto Brush | [MobileSAM](https://github.com/ChaoningZhang/MobileSAM): Segment Anything with a 5M-parameter TinyViT encoder | Chaoning Zhang et al., building on Meta AI's Segment Anything and Microsoft's TinyViT | Apache-2.0 | `mobile_sam.pt` from the official repository | 40.7 MB | +| `sam2.1-hiera-tiny` | Roto Brush | [SAM 2.1 Tiny](https://github.com/facebookresearch/sam2): Hiera-Tiny video object tracker with streaming memory | Meta AI (Nikhila Ravi et al.) | Apache-2.0 | `model.safetensors` of `facebook/sam2.1-hiera-tiny` (Hugging Face, pinned to a commit); loaded by `effectcraft-sam2` | 155.9 MB | | `mediapipe-face` | Face tracking | [MediaPipe Face Landmarker](https://developers.google.com/edge/mediapipe/solutions/vision/face_landmarker): BlazeFace detector and Face Mesh V2 | Google LLC | Apache-2.0 ([model card](https://storage.googleapis.com/mediapipe-assets/Model%20Card%20MediaPipe%20Face%20Mesh%20V2.pdf)) | `face_landmarker.task` (float16, version 1) from Google's model storage | 3.8 MB | How they were checked: diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index adeb44881..dbfe07c1f 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -96,6 +96,24 @@ pub const MOBILE_SAM: ModelInfo = ModelInfo { sha256: "6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f", }; +/// SAM 2.1 Tiny (Meta AI, Apache-2.0): a video object tracker with streaming memory. The weights are +/// the Hugging Face `transformers` checkpoint, pinned to a commit so the SHA-256 stays valid; they +/// are loaded by `effectcraft-sam2` (native only), not by [`load`]. +pub const SAM2_TINY: ModelInfo = ModelInfo { + id: "sam2.1-hiera-tiny", + task: Task::Mask, + name: "SAM 2.1 Tiny", + description: "Meta's Segment Anything Model 2 (Hiera-Tiny, 39M parameters). Follows the object through the clip with a streaming memory: it keeps tracking through occlusion and when the object leaves and re-enters the frame. Slower than MobileSAM on a CPU.", + authors: "Meta AI (Nikhila Ravi et al., SAM 2), Hugging Face transformers checkpoint", + licence: "Apache-2.0", + licence_url: "https://github.com/facebookresearch/sam2/blob/main/LICENSE", + homepage: "https://github.com/facebookresearch/sam2", + url: "https://huggingface.co/facebook/sam2.1-hiera-tiny/resolve/de431c4043854a71d8101e17995dfe596bf101a5/model.safetensors", + file_name: "sam2.1-hiera-tiny.safetensors", + size: 155_908_064, + sha256: "48c14467e5cf9e51870511feb72c89688e82dd74523142c0538b663e193ac2a7", +}; + pub const FACE_LANDMARKER: ModelInfo = ModelInfo { id: "mediapipe-face", task: Task::Face, @@ -112,7 +130,7 @@ pub const FACE_LANDMARKER: ModelInfo = ModelInfo { }; /// The trained models, besides the built-in classical engines. -pub const MODELS: &[ModelInfo] = &[MOBILE_SAM, FACE_LANDMARKER]; +pub const MODELS: &[ModelInfo] = &[MOBILE_SAM, SAM2_TINY, FACE_LANDMARKER]; /// The id of the built-in classical engines (graph cut segmenter, shape-model face tracker): no /// weights, always available. @@ -150,6 +168,7 @@ pub fn load(id: &str, bytes: &[u8]) -> Result { Ok(Loaded::Mask(Arc::new(Sam { net: mobilesam::MobileSam::from_weights(&w)? }))) } "mediapipe-face" => Ok(Loaded::Face(Arc::new(mediapipe::FaceLandmarker::from_task(bytes)?))), + "sam2.1-hiera-tiny" => Err(format!("{} is loaded by the desktop application (effectcraft-sam2)", info.name)), _ => Err(format!("no loader for `{id}`")), } } @@ -263,6 +282,8 @@ mod tests { assert!(m.url.starts_with("https://") && m.licence_url.starts_with("https://") && m.homepage.starts_with("https://")); assert!(!m.authors.is_empty() && notice(m).contains(m.authors)); } + assert_eq!(models(Task::Mask).map(|m| m.id).collect::>(), ["mobilesam", "sam2.1-hiera-tiny"]); + assert!(SAM2_TINY.url.contains("/resolve/") && SAM2_TINY.file_name.ends_with(".safetensors"), "pinned to a commit, not a branch"); assert_eq!(models(Task::Face).map(|m| m.id).collect::>(), ["mediapipe-face"]); assert!(load("mediapipe-face", b"nope").err().is_some_and(|e| e.contains("does not match"))); // A file that isn't the published one is refused before anything is parsed. diff --git a/docs/agents.md b/docs/agents.md index d801fcec6..864e7cf59 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -506,7 +506,7 @@ Trained model (Roto Brush 2.0 / 3.0): licence, size, URL, SHA-256 and whether it is installed, selected and active. - `roto.model.download {id}` fetches and verifies the official weights in the background; `roto.model.install {path, id?}` installs a file you already have. -- `roto.model.select {id: classical|mobilesam}` chooses which one Roto Brush uses. Loading runs +- `roto.model.select {id: classical|mobilesam|sam2.1-hiera-tiny}` chooses which one Roto Brush uses. Loading runs in the background; `"wait": true` (also on download and install) returns once it is in use. - `roto.model.remove {id}` deletes the installed weights. - The `version` property (`effects/#1/version`: 0 = 1.0, 1 = 2.0, 2 = 3.0) picks between the From 92ac189a49f10d8981894259934f61111f6094b4 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:44:45 +0200 Subject: [PATCH 15/77] 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()); +} From a62d59da4bf7b5489d4e86a0652fffc0b7619e7d Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:47:59 +0200 Subject: [PATCH 16/77] S2.40: SAM 2 as the Roto Brush object tracker ObjectTracker interface in the segment crate (prompts on some frames, the clip tracked both ways, each frame streamed with confidence and an occluded flag), implemented by SAM 2. MaskModel::tracker() is None for other models. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/src/lib.rs | 1 + crates/sam2/src/mask.rs | 4 + crates/sam2/src/session.rs | 10 +-- crates/sam2/src/track.rs | 131 +++++++++++++++++++++++++++++++++ crates/sam2/tests/reference.rs | 60 +++++++++++++++ crates/segment/src/lib.rs | 48 ++++++++++++ 6 files changed, 249 insertions(+), 5 deletions(-) create mode 100644 crates/sam2/src/track.rs diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 0d7700a0e..6a0eb7fbf 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -21,6 +21,7 @@ pub mod neck; pub mod nn; pub mod prompt; pub mod session; +pub mod track; pub mod weights; /// Everything that can go wrong in this crate. Never a panic. diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index 8bb309da0..a0eac9cd6 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -105,6 +105,10 @@ impl MaskModel for Sam2Mask { 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()) } + + fn tracker(&self) -> Option<&dyn effectcraft_segment::ObjectTracker> { + Some(self) + } } /// Verify the installed file against the registry and build the model. diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 45d1ffd0a..3440ff379 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -33,7 +33,7 @@ const MAX_FRAMES: usize = 1_000_000; const GRID: usize = 64; /// Frames of the clip, already resized to the model input and normalised: `[1, 3, 1024, 1024]`. -pub trait FrameSource: Send { +pub trait FrameSource { fn num_frames(&self) -> usize; fn frame(&self, index: usize) -> Result; } @@ -129,9 +129,9 @@ struct FrameFeatures { high: (Tensor, Tensor), } -pub struct VideoTracker { +pub struct VideoTracker<'a> { model: Arc, - frames: Box, + frames: Box, objects: BTreeMap, cache: VecDeque<(usize, Arc)>, } @@ -151,8 +151,8 @@ fn scalar(t: &Tensor) -> Result { Ok(t.flatten_all()?.to_vec1::()?.first().copied().unwrap_or(0.0)) } -impl VideoTracker { - pub fn new(model: Arc, frames: Box) -> Result { +impl<'a> VideoTracker<'a> { + pub fn new(model: Arc, frames: Box) -> Result> { let n = frames.num_frames(); if n == 0 || n > MAX_FRAMES { return Err(Error::Input(format!("a clip needs between 1 and {MAX_FRAMES} frames, got {n}"))); diff --git a/crates/sam2/src/track.rs b/crates/sam2/src/track.rs new file mode 100644 index 000000000..9320fbbe4 --- /dev/null +++ b/crates/sam2/src/track.rs @@ -0,0 +1,131 @@ +// Modified work: adapts the ported SAM 2 video tracker to EffectCraft's `ObjectTracker` +// interface (Apache-2.0, see NOTICE). +//! SAM 2 as the Roto Brush's object tracker: prompts on some frames, the whole clip tracked +//! forwards and backwards, each frame reported as it is finished. + +use std::cell::RefCell; +use std::collections::HashMap; +use std::rc::Rc; +use std::sync::Arc; + +use candle_core::Tensor; +use effectcraft_segment::{ClipFrame, ObjectTracker, TrackRequest, TrackedFrame}; + +use crate::mask::{Sam2Mask, preprocess}; +use crate::prompt::{INPUT_SIZE, Prompt}; +use crate::session::{Cancel, Direction, FrameMask, FrameSource, ObjectId, VideoTracker}; +use crate::{Error, Result}; + +/// Most frames one tracking request may cover. +const MAX_CLIP: usize = 100_000; + +type Frames<'a> = &'a mut dyn FnMut(i64) -> Option>; + +/// Feeds the host's frames to the tracker, remembering each frame's size for the output. +struct Source<'a> { + first: i64, + count: usize, + fetch: RefCell>, + sizes: Rc>>, +} + +impl FrameSource for Source<'_> { + fn num_frames(&self) -> usize { + self.count + } + + fn frame(&self, index: usize) -> Result { + let f = (self.fetch.borrow_mut())(self.first + index as i64) + .ok_or_else(|| Error::Input(format!("frame {} is not available", self.first + index as i64)))?; + self.sizes.borrow_mut().insert(index, (f.w, f.h)); + preprocess(&f.rgb, f.w, f.h) + } +} + +fn to_frame(m: &FrameMask, size: (usize, usize)) -> Result { + let (w, h) = size; + let logits = Tensor::from_slice(&m.logits, (1, 1, 256, 256), &candle_core::Device::Cpu)?; + let prob = candle_nn::ops::sigmoid(&logits.upsample_bilinear2d(h, w, false)?)?.flatten_all()?.to_vec1::()?; + let presence = 1.0 / (1.0 + (-m.object_score).exp()); + Ok(TrackedFrame { prob, w, h, confidence: (presence * m.iou.clamp(0.0, 1.0)).clamp(0.0, 1.0), occluded: m.occluded }) +} + +impl Sam2Mask { + fn run_track(&self, req: &TrackRequest, frames: Frames<'_>, cancel: &dyn Fn() -> bool, out: &mut dyn FnMut(i64, TrackedFrame)) -> Result<()> { + if req.count == 0 || req.count > MAX_CLIP { + return Err(Error::Input(format!("a tracking request needs 1 to {MAX_CLIP} frames, got {}", req.count))); + } + if req.prompts.is_empty() { + return Err(Error::Input("tracking needs a prompt on at least one frame".into())); + } + let sizes = Rc::new(RefCell::new(HashMap::new())); + let source = Source { first: req.first, count: req.count, fetch: RefCell::new(frames), sizes: sizes.clone() }; + // Prompt coordinates are in the frame's pixels and the model works at 1024 x 1024, so the + // prompted frames' sizes are needed first. + let mut prompt_frames = vec![]; + for (frame, p) in &req.prompts { + let idx = usize::try_from(frame - req.first) + .ok() + .filter(|i| *i < req.count) + .ok_or_else(|| Error::Input(format!("prompt frame {frame} is outside the clip")))?; + let f = (source.fetch.borrow_mut())(*frame).ok_or_else(|| Error::Input(format!("frame {frame} is not available")))?; + sizes.borrow_mut().insert(idx, (f.w, f.h)); + prompt_frames.push((idx, p)); + } + let mut tracker = VideoTracker::new(self.video_model().clone(), Box::new(source))?; + let object = ObjectId(1); + let mut prompted = vec![]; + for (idx, p) in prompt_frames { + if cancel() { + return Err(Error::Cancelled); + } + let (w, h) = sizes.borrow().get(&idx).copied().ok_or_else(|| Error::Compute("frame size unknown".into()))?; + let (sx, sy) = (INPUT_SIZE as f32 / w as f32, INPUT_SIZE as f32 / h as f32); + let prompt = Prompt { + points: p.points.iter().map(|(c, fg)| ([c[0] * sx, c[1] * sy], *fg)).collect(), + bbox: p.bbox.map(|b| [b[0] * sx, b[1] * sy, b[2] * sx, b[3] * sy]), + }; + tracker.prompt(idx, object, &prompt)?; + prompted.push(idx); + } + let first = prompted.iter().copied().min().unwrap_or(0); + let flag = Cancel::new(); + let mut failure: Option = None; + let mut emit = |m: FrameMask| { + if cancel() { + flag.cancel(); + } + match to_frame(&m, sizes.borrow().get(&m.frame).copied().unwrap_or((0, 0))) { + Ok(t) => out(req.first + m.frame as i64, t), + Err(e) => { + failure.get_or_insert(e); + flag.cancel(); + } + } + }; + tracker.propagate(first, Direction::Forward, None, &flag, &mut emit)?; + if first > 0 { + // frame `first` was reported by the forward pass; the backward pass starts there again + let mut skip_first = true; + tracker.propagate(first, Direction::Backward, None, &flag, &mut |m| { + if std::mem::take(&mut skip_first) { + return; + } + emit(m); + })?; + } + failure.map_or(Ok(()), Err) + } +} + +impl ObjectTracker for Sam2Mask { + fn track( + &self, + req: &TrackRequest, + frames: &mut dyn FnMut(i64) -> Option>, + cancel: &dyn Fn() -> bool, + out: &mut dyn FnMut(i64, TrackedFrame), + ) -> effectcraft_segment::Result<()> { + self.run_track(req, frames, cancel, out).map_err(|e| e.to_string()) + } +} diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index 63c1a06ed..ce1fd14d8 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -337,3 +337,63 @@ fn sam2_works_as_a_mask_model() { assert!(model.segment(&rgb, fw, fh, &Prompt::default()).is_err()); assert!(model.segment(&rgb, fw + 1, fh, &prompt).is_err()); } + +#[test] +fn sam2_tracks_a_clip_through_the_object_tracker_interface() { + use effectcraft_segment::{ClipFrame, MaskModel, Prompt, TrackRequest}; + let Some((w, _)) = load() else { return }; + let (fw, fh, n) = (128usize, 96usize, 6usize); + let centre = |t: usize| (40.0 + 8.0 * t as f32, 48.0f32); + let inside = |t: usize, x: usize, y: usize| { + let c = centre(t); + ((x as f32 - c.0).powi(2) + (y as f32 - c.1).powi(2)).sqrt() < 26.0 + }; + let frames: Vec> = (0..n) + .map(|t| { + let mut rgb = vec![[0f32; 3]; fw * fh]; + 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(t, x, y) { [0.85 + 0.1 * ((x as f32 - centre(t).0) / 5.0).sin(), 0.4, 0.15] } else { [bg, bg * 0.8 + 0.1, bg * 0.6 + 0.3] }; + } + } + std::sync::Arc::new(ClipFrame { rgb, w: fw, h: fh }) + }) + .collect(); + let model = effectcraft_sam2::mask::Sam2Mask::new(std::sync::Arc::new(effectcraft_sam2::session::VideoModel::load(&w).unwrap())); + assert!(model.tracker().is_some()); + let c = centre(1); + let req = TrackRequest { first: 100, count: n, prompts: vec![(101, Prompt { points: vec![([c.0, c.1], true)], bbox: None, mask: None })] }; + let mut seen = std::collections::BTreeMap::new(); + let mut get = |f: i64| frames.get((f - 100) as usize).cloned(); + model + .tracker() + .unwrap() + .track(&req, &mut get, &|| false, &mut |f, t| { + assert!(seen.insert(f, t).is_none(), "frame {f} reported twice"); + }) + .unwrap(); + assert_eq!(seen.keys().copied().collect::>(), (100..100 + n as i64).collect::>()); + for (f, t) in &seen { + assert_eq!((t.w, t.h, t.prob.len()), (fw, fh, fw * fh)); + let k = (f - 100) as usize; + let (mut i, mut u) = (0usize, 0usize); + for y in 0..fh { + for x in 0..fw { + let (a, b) = (t.prob[y * fw + x] > 0.5, inside(k, x, y)); + i += usize::from(a && b); + u += usize::from(a || b); + } + } + let j = i as f64 / u.max(1) as f64; + eprintln!("frame {f}: IoU {j:.3}, confidence {:.2}, occluded {}", t.confidence, t.occluded); + assert!(j > 0.85 && !t.occluded && t.confidence > 0.3, "frame {f}: IoU {j}"); + } + // Cancelling before the work starts, and bad requests, are errors. + assert!(model.tracker().unwrap().track(&req, &mut get, &|| true, &mut |_, _| {}).is_err()); + let none = TrackRequest { first: 100, count: n, prompts: vec![] }; + assert!(model.tracker().unwrap().track(&none, &mut get, &|| false, &mut |_, _| {}).is_err()); + let outside = TrackRequest { first: 100, count: n, prompts: vec![(500, Prompt { points: vec![([1.0, 1.0], true)], bbox: None, mask: None })] }; + assert!(model.tracker().unwrap().track(&outside, &mut get, &|| false, &mut |_, _| {}).is_err()); +} diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index dbfe07c1f..37b3a4ca0 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -45,6 +45,54 @@ pub trait MaskModel: Send + Sync { fn info(&self) -> &'static ModelInfo; /// Foreground probability (0–1) of each pixel of a frame (straight RGB 0–1, row-major `w×h`). fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &Prompt) -> Result>; + /// A video tracker built on the same weights, if the model has one. + fn tracker(&self) -> Option<&dyn ObjectTracker> { + None + } +} + +/// One decoded frame handed to a tracker (straight RGB 0–1, row-major `w×h`). +#[derive(Clone, Debug, Default)] +pub struct ClipFrame { + pub rgb: Vec<[f32; 3]>, + pub w: usize, + pub h: usize, +} + +/// What to track: the clip's frames `first..first + count` (the host's frame numbers) and the +/// prompts of the one object, each on its frame. +#[derive(Clone, Debug)] +pub struct TrackRequest { + pub first: i64, + pub count: usize, + pub prompts: Vec<(i64, Prompt)>, +} + +/// A tracker's answer for one frame. +#[derive(Clone, Debug, Default)] +pub struct TrackedFrame { + /// Foreground probability of every pixel (`w×h`); all zero when the object is absent. + pub prob: Vec, + pub w: usize, + pub h: usize, + /// How sure the tracker is (0–1), for the timeline. + pub confidence: f32, + /// The object is not in this frame (occluded or out of view). + pub occluded: bool, +} + +/// A video tracker with memory (SAM 2): given prompts on some frames, it follows the object +/// through the clip in both directions. Optional: a [`MaskModel`] may offer one. +pub trait ObjectTracker: Send + Sync { + /// Track the object. `frames(i)` supplies frame `i` (`None`: unavailable); `cancel()` is polled + /// between frames; `out(i, result)` is called as soon as each frame is finished. + fn track( + &self, + req: &TrackRequest, + frames: &mut dyn FnMut(i64) -> Option>, + cancel: &dyn Fn() -> bool, + out: &mut dyn FnMut(i64, TrackedFrame), + ) -> Result<()>; } /// What a model is for. From c1f5ff9f7dec765521e75a41dba8c9d599bcbc2c Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 02:09:10 +0200 Subject: [PATCH 17/77] S2.41: Roto Brush follows the object with SAM 2 When the chosen model has a tracker, roto.propagate runs one tracking pass over the span: prompts from the strokes (a handful of points, which SAM 2 trusts), the clip tracked both ways, each frame's probability map refined by the same graph cut with the strokes as hard constraints. Results are cached under their own keys with a per-frame confidence and an occluded flag; roto.status reports them (source: tracker|classical). Without a tracker nothing changes. Real weights: IoU 0.98 on all 21 frames, confidence 0.99. Co-Authored-By: Claude Sonnet 5.5 --- crates/effects/src/roto.rs | 114 +++++++++++++++++++++++++-- crates/engine/src/analysis/models.rs | 5 ++ crates/engine/src/analysis/roto.rs | 29 +++++++ crates/engine/src/cmd/roto_cmds.rs | 13 +++ crates/engine/src/tests/roto.rs | 38 +++++++++ crates/track/src/roto.rs | 41 ++++++++++ crates/track/src/roto/tests.rs | 33 ++++++++ 7 files changed, 268 insertions(+), 5 deletions(-) diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 2467a719d..96d6b467c 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -29,7 +29,7 @@ use std::sync::{Arc, Mutex, OnceLock, RwLock, Weak}; use effectcraft_keyframe::Value; use effectcraft_project::ParamUi; use effectcraft_raster::Image; -use effectcraft_segment::MaskModel; +use effectcraft_segment::{ClipFrame, MaskModel, TrackRequest}; use effectcraft_track::roto::refine::{self, MatteParams}; use effectcraft_track::roto::{self as rb, RotoData, SegOpts, rle}; use rayon::prelude::*; @@ -298,9 +298,39 @@ fn from_store(key: u64, scale: f64) -> Option> { Some(seg) } +/// Tracker results live under the chain key mixed with this, so a frame the tracker followed and +/// one the classic engine computed never share a cache slot. +const TRACK_SALT: u64 = 0x7472_6163_6b65_7231; + +fn tracked_key(key: u64) -> u64 { + key.rotate_left(17) ^ TRACK_SALT +} + +/// How far to trust a frame the tracker followed. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct FrameTrust { + /// 0–1, from the tracker's own object score and mask quality. + pub confidence: f32, + /// The object is not in this frame (an empty matte). + pub occluded: bool, +} + +fn trust_cache() -> &'static Mutex> { + static T: OnceLock>> = OnceLock::new(); + T.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// The tracker's confidence in the frame with chain key `key`, if the tracker followed it. +pub fn trust(key: u64) -> Option { + trust_cache().lock().unwrap_or_else(|e| e.into_inner()).get(&tracked_key(key)).copied() +} + /// Drop the cached segmentations of these chain keys (every scale). pub fn forget(keys: impl IntoIterator) { - let keys: std::collections::HashSet = keys.into_iter().collect(); + let mut trusts = trust_cache().lock().unwrap_or_else(|e| e.into_inner()); + let keys: std::collections::HashSet = keys.into_iter().flat_map(|k| [k, tracked_key(k)]).collect(); + trusts.retain(|k, _| !keys.contains(k)); + drop(trusts); let mut c = cache().lock().unwrap_or_else(|e| e.into_inner()); let gone: Vec<(u64, u64)> = c.map.keys().filter(|k| keys.contains(&k.0)).copied().collect(); for k in gone { @@ -313,6 +343,7 @@ pub fn forget(keys: impl IntoIterator) { /// Drop every cached segmentation (Edit ▸ Purge). pub fn purge() { + trust_cache().lock().unwrap_or_else(|e| e.into_inner()).clear(); let mut c = cache().lock().unwrap_or_else(|e| e.into_inner()); c.map.clear(); c.order.clear(); @@ -333,6 +364,16 @@ pub fn layer_grid(buf: &Buf, layer_size: [f64; 2]) -> Image { img } +/// What a tracking pass found. +#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct TrackSummary { + pub frames: usize, + /// Frames where the object was absent. + pub occluded: usize, + pub min_confidence: f32, +} + /// Everything needed to compute an instance's segmentations. pub struct Chain { pub data: Arc, @@ -357,11 +398,16 @@ impl Chain { let sc = self.opts.scale; let lookup = |g: i64| -> Option> { let k = *self.keys.get(&g)?; + // A frame the tracker followed wins over the classic engine's (memory only: its + // confidence is not persisted). + if let Some(s) = cached(tracked_key(k), sc) { + return Some(s); + } if let Some(s) = cached(k, sc).or_else(|| from_store(k, sc)) { return Some(s); } if sc != 1.0 { - let full = cached(k, 1.0).or_else(|| from_store(k, 1.0))?; + let full = cached(tracked_key(k), 1.0).or_else(|| cached(k, 1.0)).or_else(|| from_store(k, 1.0))?; let s = Arc::new(downscale(&full, sc)); store(k, sc, s.clone()); return Some(s); @@ -407,7 +453,8 @@ impl Chain { if pimg.width != img.width || pimg.height != img.height { return None; } - rb::propagate_with(&pimg, p, &img, &strokes, &self.opts, self.refine_radius, self.model.as_deref()) + // A tracker is too slow to prompt on every frame (its pass is `track`); the classic flow carries the rest. + rb::propagate_with(&pimg, p, &img, &strokes, &self.opts, self.refine_radius, if self.has_tracker() { None } else { self.model.as_deref() }) } }; let seg = Arc::new(seg); @@ -424,7 +471,64 @@ impl Chain { /// Whether frame `f` is computed (at full resolution or at `scale`). pub fn is_cached(&self, f: i64, scale: f64) -> bool { - self.keys.get(&f).is_some_and(|k| cached(*k, 1.0).is_some() || cached(*k, scale).is_some()) + self.keys.get(&f).is_some_and(|k| [*k, tracked_key(*k)].iter().any(|k| cached(*k, 1.0).is_some() || cached(*k, scale).is_some())) + } + + /// Whether the chosen model can follow an object through a clip (SAM 2). + pub fn has_tracker(&self) -> bool { + self.model.as_ref().is_some_and(|m| m.tracker().is_some()) + } + + /// Follow the object through the whole span with the model's tracker, from the strokes on + /// their frames, and cache every frame's segmentation (under the tracker's keys, with its + /// confidence). The strokes stay hard constraints and the edges are refined at full + /// resolution; frames the tracker can't give are left to [`Chain::ensure`]'s classic engine. + /// `progress` gets the number of frames finished. `Ok(None)`: no tracker, nothing done. + pub fn track( + &self, + frame: &mut dyn FnMut(i64) -> Option>, + cancel: &dyn Fn() -> bool, + progress: &mut dyn FnMut(usize), + ) -> Result, String> { + let Some(tracker) = self.model.as_ref().and_then(|m| m.tracker()) else { return Ok(None) }; + let (Some(_), true) = (self.data.base, self.data.span[1] >= self.data.span[0]) else { return Err("no base frame".into()) }; + let prompts: Vec<(i64, effectcraft_segment::Prompt)> = self + .data + .stroke_frames() + .into_iter() + .filter(|f| self.data.in_span(*f)) + .filter_map(|f| { + let strokes = self.data.strokes_at(f); + let p = rb::tracker_prompt(&strokes, self.opts.scale); + p.points.iter().any(|(_, fg)| *fg).then_some((f, p)) + }) + .collect(); + if prompts.is_empty() { + return Err("no foreground strokes to follow".into()); + } + let req = TrackRequest { first: self.data.span[0], count: (self.data.span[1] - self.data.span[0] + 1) as usize, prompts }; + let frame = std::cell::RefCell::new(frame); + let mut summary = TrackSummary::default(); + tracker.track( + &req, + &mut |g| { + let img = (frame.borrow_mut())(g)?; + Some(Arc::new(ClipFrame { rgb: rb::rgb_of(&img), w: img.width as usize, h: img.height as usize })) + }, + cancel, + &mut |g, t| { + let (Some(key), Some(img)) = (self.keys.get(&g).copied(), (frame.borrow_mut())(g)) else { return }; + let strokes = self.data.strokes_at(g); + let seg = rb::segment_tracked(&img, &strokes, t.prob, t.occluded, &self.opts, self.refine_radius); + store(tracked_key(key), self.opts.scale, Arc::new(seg)); + trust_cache().lock().unwrap_or_else(|e| e.into_inner()).insert(tracked_key(key), FrameTrust { confidence: t.confidence, occluded: t.occluded }); + summary.frames += 1; + summary.occluded += usize::from(t.occluded); + summary.min_confidence = if summary.frames == 1 { t.confidence } else { summary.min_confidence.min(t.confidence) }; + progress(summary.frames); + }, + )?; + Ok(Some(summary)) } } diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index 83d0d3c51..dbe8a1d5f 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -188,6 +188,11 @@ impl Session { }); } + /// The Roto Brush model in use (`None`: the classic engine). + pub fn models_roto_active(&self) -> Option { + self.models.status(Task::Mask).active + } + /// Finish what the model workers did: a toast, and loading a model that just arrived. pub fn poll_models(&mut self) { for task in [Task::Mask, Task::Face] { diff --git a/crates/engine/src/analysis/roto.rs b/crates/engine/src/analysis/roto.rs index 813561969..12ce21260 100644 --- a/crates/engine/src/analysis/roto.rs +++ b/crates/engine/src/analysis/roto.rs @@ -112,6 +112,10 @@ pub struct RotoProgress { pub error: Option, /// The job works on an Object Matte (else a Roto Brush & Refine Edge). pub object_matte: bool, + /// What the model's tracker found, when one followed the object (SAM 2). + pub tracked: Option, + /// Why the tracker was not used, when a tracking model was chosen but the classic engine ran. + pub tracker_note: Option, } impl RotoProgress { @@ -436,6 +440,31 @@ async fn run_work(w: Work, shared: &RotoShared) { let model = fx::installed_model(); let key = if w.object_matte { om::frozen_key(&w.params, w.size, w.fps, model.as_deref()) } else { fx::frozen_key(&w.params, w.size, w.fps) }; let mut frozen = FrozenCache { key, ..Default::default() }; + // A tracking model (SAM 2) follows the object through the whole span in one pass; the loop + // below then finds every frame cached. Without one (or if it fails) nothing changes: the + // classic engine runs frame by frame. + if w.task == RotoTask::Propagate && !w.object_matte && chain.has_tracker() { + let mut get = |g: i64| src.get(g); + let mut report = |n: usize| { + let mut s = lock(&shared.state); + s.done = n as u64; + s.elapsed = t0.elapsed().as_secs_f64(); + s.fps = if s.elapsed > 0.0 { s.done as f64 / s.elapsed } else { 0.0 }; + crate::offload::report(s.done, s.total); + }; + let result = chain.track(&mut get, &cancel, &mut report); + if cancel() { + let mut s = lock(&shared.state); + s.cancelled = true; + s.finished = true; + return; + } + let mut s = lock(&shared.state); + match result { + Ok(summary) => s.tracked = summary, + Err(e) => s.tracker_note = Some(format!("tracker unavailable, classic engine used: {e}")), + } + } for (k, f) in w.frames.iter().enumerate() { if cancel() { let mut s = lock(&shared.state); diff --git a/crates/engine/src/cmd/roto_cmds.rs b/crates/engine/src/cmd/roto_cmds.rs index a6d389ca1..c37adb0c2 100644 --- a/crates/engine/src/cmd/roto_cmds.rs +++ b/crates/engine/src/cmd/roto_cmds.rs @@ -463,6 +463,19 @@ fn status(s: &mut Session, p: &Value) -> Result { } } } + // Where each frame's matte came from, and how far to trust it: frames a tracker followed + // report its confidence and whether the object was absent; the rest came from the classic engine. + { + let size = rt::layer_size(&s.project, &comp, &layer); + let chain = fx::Chain::new(¶ms_static(&g), size, comp.frame_rate.as_f64(), 1.0); + let tracked: serde_json::Map = chain + .keys + .iter() + .filter_map(|(f, k)| fx::trust(*k).map(|t| (f.to_string(), json!({"confidence": t.confidence, "occluded": t.occluded})))) + .collect(); + out["tracker"] = json!({"model": s.models_roto_active(), "frames": tracked}); + out["source"] = json!(if tracked.contains_key(&frame.to_string()) { "tracker" } else { "classical" }); + } let computed = s.roto_computed(cid, lid, uid); out["computed"] = json!(computed.iter().filter(|(_, c)| **c).map(|(f, _)| *f).collect::>()); Ok(out) diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index 64ee30d5e..895a4222a 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -668,3 +668,41 @@ fn object_matte_mattes_persist_in_the_disk_cache() { assert!(om::ensure(&moved, size, 25.0, 0, None, &mut |_| None, &|| false).is_none()); let _ = std::fs::remove_dir_all(&dir); } + +/// SAM 2 tracks the clip end to end (needs the real weights: `EFFECTCRAFT_SAM2=path/to/model.safetensors`). +/// Every frame comes from the tracker with a confidence, the mattes are right, and nothing changes +/// once the model is deselected. +#[cfg(not(target_arch = "wasm32"))] +#[test] +fn sam2_tracks_the_clip_through_roto_propagate() { + let Ok(weights) = std::env::var("EFFECTCRAFT_SAM2") else { return }; + let dir = std::env::temp_dir().join(format!("ec-roto-sam2-{}", std::process::id())); + let (mut s, clip, _) = painted(center, "disk-sam2.mov"); + s.models_dir = Some(dir.join("models")); + std::fs::create_dir_all(dir.join("models")).unwrap(); + s.execute_checked("roto.model.install", json!({"path": weights})).unwrap(); + s.execute_checked("roto.model.select", json!({"id": "sam2.1-hiera-tiny", "wait": true})).unwrap(); + assert_eq!(s.models.status(effectcraft_segment::Task::Mask).active.as_deref(), Some("sam2.1-hiera-tiny")); + let r = s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})).unwrap(); + assert_eq!(r["frames"], 21, "{r}"); + let mut worst: f64 = 1.0; + for f in 0..=20 { + let j = iou_at(&mut s, clip, f, center(f)); + let st = s.execute("roto.status", json!({"layer": clip.0, "frame": f})).unwrap(); + eprintln!("frame {f}: IoU {j:.3}, {:?}", st["tracker"]["frames"][f.to_string()]); + worst = worst.min(j); + } + assert!(worst > 0.9, "worst IoU over the span with SAM 2: {worst}"); + let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); + assert_eq!(st["source"], "tracker", "{st}"); + let frames = st["tracker"]["frames"].as_object().unwrap(); + assert_eq!(frames.len(), 21); + assert!(frames.values().all(|t| t["occluded"] == false && t["confidence"].as_f64().unwrap() > 0.3), "{frames:?}"); + assert_eq!(st["tracker"]["model"], "sam2.1-hiera-tiny"); + // Back to the classic engine: same strokes, its own segmentation, no tracker data. + s.execute_checked("roto.model.select", json!({"id": "classical", "wait": true})).unwrap(); + let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); + assert_eq!(st["source"], "classical", "{st}"); + assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 0); + let _ = std::fs::remove_dir_all(dir); +} diff --git a/crates/track/src/roto.rs b/crates/track/src/roto.rs index abcf28aa0..ce52a79a4 100644 --- a/crates/track/src/roto.rs +++ b/crates/track/src/roto.rs @@ -903,6 +903,47 @@ fn points_along(s: &Stroke, scale: f64, max: usize) -> Vec<[f32; 2]> { out } +/// A tracker's prompt from a frame's strokes (points along the foreground and background strokes, +/// in frame pixels at `scale`). +/// +/// SAM 2 was trained on a handful of clicks: with sixteen or more points it doubts its own mask +/// (measured: confidence 0.25 with 16 points, 0.00 with 32, 0.99 with 1 to 7), so this picks at +/// most [`TRACKER_FG_POINTS`] foreground and [`TRACKER_BG_POINTS`] background points, spread +/// along the strokes. +pub fn tracker_prompt(strokes: &[&Stroke], scale: f64) -> Prompt { + let mut points = vec![]; + for (kind, fg, max) in [(StrokeKind::Fg, true, TRACKER_FG_POINTS), (StrokeKind::Bg, false, TRACKER_BG_POINTS)] { + let all: Vec<[f32; 2]> = strokes.iter().filter(|s| s.kind == kind).flat_map(|s| points_along(s, scale, 3)).collect(); + let n = all.len().min(max); + // evenly spaced picks, keeping the first and the last + points.extend((0..n).filter_map(|i| all.get(if n > 1 { i * (all.len() - 1) / (n - 1) } else { all.len() / 2 }).map(|p| (*p, fg)))); + } + // a box the user drew on the frame prompts the tracker too + let bbox = strokes.iter().filter(|s| s.kind == StrokeKind::Box).find_map(|s| box_of(s, scale)); + Prompt { points, bbox, ..Default::default() } +} + +/// Most foreground / background points a tracker prompt carries. +pub const TRACKER_FG_POINTS: usize = 3; +pub const TRACKER_BG_POINTS: usize = 3; + +/// A frame's segmentation from a tracker's probability map: the map is the shape prior of the +/// same graph cut as everywhere else, so the strokes stay hard constraints and the edges snap +/// to colour. An absent object gives an empty matte (strokes on the frame still win); a map of +/// the wrong size is ignored. +pub fn segment_tracked(img: &Image, strokes: &[&Stroke], prob: Vec, occluded: bool, opts: &SegOpts, refine_radius: f64) -> FrameSeg { + let (w, h) = (img.width as usize, img.height as usize); + if prob.len() != w * h || w == 0 || h == 0 { + return segment(img, strokes, None, opts, refine_radius); + } + let has_fg = strokes.iter().any(|s| s.kind == StrokeKind::Fg); + if (occluded || !prob.iter().any(|p| *p >= 0.5)) && !has_fg { + return FrameSeg::empty(w, h); + } + let prior = Prior { soft: prob, fg: Gmm::default(), bg: Gmm::default(), refine: vec![0; w * h], weight: MODEL_PRIOR }; + segment(img, strokes, Some(&prior), opts, refine_radius) +} + /// A model's prompt from the strokes: points along the foreground and background strokes, but /// not where a later stroke of the other kind painted over them (`hard`, from [`resolve`]), and /// the box of a [`StrokeKind::Box`] stroke (matte pixels). A stroke painted over at every point diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index 8a5cab98c..c08306060 100644 --- a/crates/track/src/roto/tests.rs +++ b/crates/track/src/roto/tests.rs @@ -579,3 +579,36 @@ fn a_correction_keeps_the_model_prompted_with_what_is_left() { let hard = resolve(&refs, StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h); assert!(stroke_prompt(&refs, &hard, 1.0, w, h).points.iter().all(|p| !p.1)); } + +#[test] +fn tracker_prompts_carry_a_handful_of_points() { + let st = base_strokes(0); + let refs: Vec<&Stroke> = st.iter().collect(); + let p = tracker_prompt(&refs, 1.0); + let (fg, bg) = (p.points.iter().filter(|(_, f)| *f).count(), p.points.iter().filter(|(_, f)| !*f).count()); + assert!((1..=TRACKER_FG_POINTS).contains(&fg) && (1..=TRACKER_BG_POINTS).contains(&bg), "{fg} fg, {bg} bg"); + // a long, finely sampled stroke still gives at most the limit, and an empty one gives none + let long = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 2.0, points: (0..400).map(|i| [i as f64 * 0.4, 50.0]).collect() }; + assert!(tracker_prompt(&[&long], 1.0).points.len() <= TRACKER_FG_POINTS); + assert!(tracker_prompt(&[], 1.0).points.is_empty()); +} + +#[test] +fn tracked_segmentation_handles_absent_objects_and_bad_maps() { + let (img, gt) = frame_at(center(0)); + let st = base_strokes(0); + let refs: Vec<&Stroke> = st.iter().collect(); + let opts = SegOpts::default(); + let n = (W * H) as usize; + // a good map (the true disk) reproduces it, with the strokes as constraints + let soft: Vec = gt.iter().map(|v| f32::from(*v)).collect(); + let seg = segment_tracked(&img, &refs, soft, false, &opts, 0.0); + assert!(iou(&seg.matte, >) > 0.95); + // an absent object (no strokes on the frame) is an empty matte, whatever the map says + let none = segment_tracked(&img, &[], vec![1.0; n], true, &opts, 0.0); + assert_eq!(none.area(), 0); + assert_eq!(segment_tracked(&img, &[], vec![0.0; n], false, &opts, 0.0).area(), 0); + // a map of the wrong size is ignored: the strokes alone segment the frame + let wrong = segment_tracked(&img, &refs, vec![1.0; 3], false, &opts, 0.0); + assert!(iou(&wrong.matte, >) > 0.95); +} From 4655425d0113568e23507d054445973a05ca5466 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 02:22:28 +0200 Subject: [PATCH 18/77] S2.50: bench --roto adds the SAM 2 tracker column run_tracker follows the object with the model's video tracker on the same synthetic clips (--sam2 FILE or the installed weights). 12 frames, dev build: re-entry clip mean IoU 0.632 classical vs 0.975 SAM 2 (0.000 vs 0.967 after the object returns); the other clips 0.98 vs 0.99. Co-Authored-By: Claude Sonnet 5.5 --- Cargo.lock | 1 + apps/effectcraft-cli/Cargo.toml | 1 + apps/effectcraft-cli/src/main.rs | 18 +++++++++++++++--- crates/track/src/roto/bench.rs | 30 ++++++++++++++++++++++++++++-- 4 files changed, 45 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7d951deb7..423abb85a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1455,6 +1455,7 @@ dependencies = [ "effectcraft-gpu", "effectcraft-host", "effectcraft-render", + "effectcraft-sam2", "effectcraft-segment", "effectcraft-time", "effectcraft-track", diff --git a/apps/effectcraft-cli/Cargo.toml b/apps/effectcraft-cli/Cargo.toml index 7d8fe96ff..abfc014eb 100644 --- a/apps/effectcraft-cli/Cargo.toml +++ b/apps/effectcraft-cli/Cargo.toml @@ -22,6 +22,7 @@ effectcraft-gpu = { workspace = true } effectcraft-time = { workspace = true } effectcraft-track = { workspace = true } effectcraft-segment = { workspace = true } +effectcraft-sam2 = { workspace = true } serde_json = { workspace = true } image = { workspace = true } cpu-time = "1" diff --git a/apps/effectcraft-cli/src/main.rs b/apps/effectcraft-cli/src/main.rs index 3de4449b6..b60918304 100644 --- a/apps/effectcraft-cli/src/main.rs +++ b/apps/effectcraft-cli/src/main.rs @@ -27,8 +27,8 @@ //! (--gpu: CPU vs GPU ms/frame for every comp at Full and Half) //! effectcraft-cli bench --ops [--small] [--layers N] [--comps N] [--footage N] everyday-operation timings //! on a large generated project (open, save, auto-save, undo/redo, timeline, Project panel) -//! effectcraft-cli bench --roto [--frames N] [--json] Roto Brush trackers on synthetic clips with known mattes -//! (classical engine, MobileSAM when installed; IoU, flicker, re-appearance, seconds per frame) +//! effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--json] Roto Brush trackers on synthetic clips with known mattes +//! (classical engine, MobileSAM and SAM 2.1 Tiny when installed; IoU, flicker, re-appearance, seconds per frame) //! effectcraft-cli script FILE.jsx [F.ecproj] | --eval CODE run an After Effects-style script //! effectcraft-cli mcp [--autosave | --bridge PORT] MCP server on stdio //! @@ -126,7 +126,8 @@ const USAGE: &str = "usage: effectcraft-cli Result<(), Failure> { if mobilesam.is_none() { note!("bench --roto: MobileSAM is not installed; comparing the classical engine only"); } + // `--sam2 FILE` uses the weights at FILE; otherwise the installed ones. + let sam2_file = args.opt("--sam2").map(std::path::PathBuf::from).or_else(|| session.model_path(&effectcraft_segment::SAM2_TINY)); + let sam2: Option> = + sam2_file.and_then(|p| std::fs::read(p).ok()).and_then(|bytes| effectcraft_sam2::mask::load(&bytes).ok()).map(|m| m as std::sync::Arc); + if sam2.is_none() { + note!("bench --roto: SAM 2.1 Tiny is not installed; no SAM 2 column"); + } let mut runs = vec![]; for sc in bench::scenarios(frames) { runs.push(bench::run_flow(&sc, "classical", None)); if let Some(m) = &mobilesam { runs.push(bench::run_flow(&sc, "mobilesam", Some(m.as_ref()))); } + if let Some(run) = sam2.as_ref().and_then(|m| bench::run_tracker(&sc, "sam2-tiny", m.as_ref())) { + runs.push(run); + } } if args.flag("--json") { say!("{}", serde_json::to_string_pretty(&runs).map_err(|e| Failure::Error(e.to_string()))?); diff --git a/crates/track/src/roto/bench.rs b/crates/track/src/roto/bench.rs index d52836979..810bf34e5 100644 --- a/crates/track/src/roto/bench.rs +++ b/crates/track/src/roto/bench.rs @@ -5,11 +5,13 @@ //! later SAM 2) is added when the caller has one. `effectcraft-cli bench --roto` prints the table. use effectcraft_raster::Image; -use effectcraft_segment::MaskModel; +use std::sync::Arc; + +use effectcraft_segment::{ClipFrame, MaskModel, TrackRequest}; use serde::Serialize; use web_time::Instant; -use super::{FrameSeg, SegOpts, Stroke, StrokeKind, iou, propagate_with, segment_with}; +use super::{FrameSeg, SegOpts, Stroke, StrokeKind, iou, propagate_with, rgb_of, segment_tracked, segment_with, tracker_prompt}; pub const W: u32 = 200; pub const H: u32 = 150; @@ -210,6 +212,30 @@ pub fn run_flow(sc: &Scenario, tracker: &str, model: Option<&dyn MaskModel>) -> metrics(sc, tracker, &mattes, secs) } +/// Follow the object through the whole clip with the model's video tracker, as the Roto Brush's +/// tracking pass does (prompt from the strokes, probability maps refined by the same graph cut). +/// `None` when the model has no tracker or the tracker fails. +pub fn run_tracker(sc: &Scenario, tracker: &str, model: &dyn MaskModel) -> Option { + let t = model.tracker()?; + let opts = SegOpts::default(); + let refs: Vec<&Stroke> = sc.strokes.iter().collect(); + let clip: Vec> = sc.frames.iter().map(|i| Arc::new(ClipFrame { rgb: rgb_of(i), w: i.width as usize, h: i.height as usize })).collect(); + let req = TrackRequest { first: 0, count: sc.frames.len(), prompts: vec![(sc.base as i64, tracker_prompt(&refs, 1.0))] }; + let start = Instant::now(); + let mut mattes: Vec> = vec![vec![0u8; (W * H) as usize]; sc.frames.len()]; + let mut get = |f: i64| usize::try_from(f).ok().and_then(|i| clip.get(i).cloned()); + let r = t.track(&req, &mut get, &|| false, &mut |f, tf| { + let (Ok(i), true) = (usize::try_from(f), tf.prob.len() == (W * H) as usize) else { return }; + let Some(img) = sc.frames.get(i) else { return }; + let strokes: Vec<&Stroke> = if i == sc.base { refs.clone() } else { vec![] }; + if let Some(slot) = mattes.get_mut(i) { + *slot = segment_tracked(img, &strokes, tf.prob, tf.occluded, &opts, 0.0).matte; + } + }); + r.ok()?; + Some(metrics(sc, tracker, &mattes, start.elapsed().as_secs_f64())) +} + /// A fixed-width table of the runs, grouped by scenario. pub fn table(runs: &[Run]) -> String { let mut out = format!("{:<10} {:<12} {:>8} {:>8} {:>8} {:>10} {:>9}\n", "clip", "tracker", "meanIoU", "minIoU", "flicker", "after ret.", "s/frame"); From a4c0a2b9c7d48bfc011aaecf220a3c2e1c61c033 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 02:25:14 +0200 Subject: [PATCH 19/77] S2.42: a failing tracker falls back to the classic engine Test with a tracker that always errors: roto.propagate completes, every frame comes from the classic engine (status source: classical), IoU unchanged. Co-Authored-By: Claude Sonnet 5.5 --- crates/engine/src/tests/roto.rs | 44 +++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index 895a4222a..b78bc2df7 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -706,3 +706,47 @@ fn sam2_tracks_the_clip_through_roto_propagate() { assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 0); let _ = std::fs::remove_dir_all(dir); } + +/// A model whose tracker always fails: Roto Brush must carry on with the classic engine and say so. +#[test] +fn a_failing_tracker_falls_back_to_the_classic_engine() { + use effectcraft_segment::{ClipFrame, MaskModel, ModelInfo, ObjectTracker, Prompt, TrackRequest, TrackedFrame}; + use std::sync::Arc; + struct Broken; + impl MaskModel for Broken { + fn info(&self) -> &'static ModelInfo { + &effectcraft_segment::SAM2_TINY + } + fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &Prompt) -> effectcraft_segment::Result> { + Err("no weights".into()) + } + fn tracker(&self) -> Option<&dyn ObjectTracker> { + Some(self) + } + } + impl ObjectTracker for Broken { + fn track( + &self, + _: &TrackRequest, + _: &mut dyn FnMut(i64) -> Option>, + _: &dyn Fn() -> bool, + _: &mut dyn FnMut(i64, TrackedFrame), + ) -> effectcraft_segment::Result<()> { + Err("the tracker broke".into()) + } + } + let (mut s, clip, _) = painted(center, "disk-broken-tracker.mov"); + effectcraft_effects::roto::set_model(Some(Arc::new(Broken))); + let r = s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})); + effectcraft_effects::roto::set_model(None); + let r = r.unwrap(); + assert_eq!(r["frames"], 21, "{r}"); + let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 10})).unwrap(); + assert_eq!(st["source"], "classical", "{st}"); + assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 0); + let mut worst: f64 = 1.0; + for f in 0..=20 { + worst = worst.min(iou_at(&mut s, clip, f, center(f))); + } + assert!(worst > 0.9, "the classic engine should have carried the clip: {worst}"); +} From 371e67f51eadb2e70d7cdbf39c0e8ac2066b1ed3 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 02:25:28 +0200 Subject: [PATCH 20/77] S2.46: document Roto Brush with SAM 2 How to install and use it, how the port was verified, speed and the benchmark. Co-Authored-By: Claude Sonnet 5.5 --- docs/agents.md | 2 +- docs/roto-brush-sam2.md | 57 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) create mode 100644 docs/roto-brush-sam2.md diff --git a/docs/agents.md b/docs/agents.md index 864e7cf59..e0c0055ec 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -506,7 +506,7 @@ Trained model (Roto Brush 2.0 / 3.0): licence, size, URL, SHA-256 and whether it is installed, selected and active. - `roto.model.download {id}` fetches and verifies the official weights in the background; `roto.model.install {path, id?}` installs a file you already have. -- `roto.model.select {id: classical|mobilesam|sam2.1-hiera-tiny}` chooses which one Roto Brush uses. Loading runs +- `roto.model.select {id: classical|mobilesam|sam2.1-hiera-tiny}` chooses which one Roto Brush uses (SAM 2 also tracks the object through the clip, see [roto-brush-sam2.md](roto-brush-sam2.md); `roto.status` then reports `source` and per-frame `tracker.frames` with `confidence` and `occluded`). Loading runs in the background; `"wait": true` (also on download and install) returns once it is in use. - `roto.model.remove {id}` deletes the installed weights. - The `version` property (`effects/#1/version`: 0 = 1.0, 1 = 2.0, 2 = 3.0) picks between the diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md new file mode 100644 index 000000000..b2c356e74 --- /dev/null +++ b/docs/roto-brush-sam2.md @@ -0,0 +1,57 @@ +# Roto Brush with SAM 2 + +Roto Brush can follow an object through a clip with Meta's **SAM 2.1 Tiny** (Hiera-Tiny, 39M parameters, +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 + 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"}`. +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 + object leaves the frame and returns. Each frame's probability map is then refined by the same graph cut as every other + Roto Brush frame: your strokes stay hard constraints and edges snap to colour. +4. `roto.status {frame}` reports `source` (`tracker` or `classical`) and, for tracked frames, `tracker.frames[f]`: + `confidence` (0–1, the model's object score times its predicted mask quality) and `occluded` (the object is not + in the frame: an empty matte). + +If the model is missing, fails, or is deselected, the classic engine runs exactly as before (`roto.status` says +`classical`, and the job's `trackerNote` says why). + +## How it works +`crates/sam2` is a Rust port (on [candle](https://github.com/huggingface/candle), native only) of the Hugging Face +`transformers` SAM 2 implementation, Apache-2.0 (see its `NOTICE`). Every stage is checked against PyTorch by +`crates/sam2/tests/reference.rs` (they skip without the weights; reference tensors come from +`tools/sam2_reference.py` and `tools/sam2_video_reference.py`): + +| Stage | Max abs error vs PyTorch | +|---|---| +| Hiera backbone, 12 blocks | 2.6e-5 | +| Neck features / position encodings | 9.7e-6 / 6e-8 | +| Prompt embeddings | 1.2e-7 | +| Mask decoder (logits up to 11.6) | 1.1e-4 | +| Memory encoder | 3.5e-6 | +| Memory attention (rotary) | 2.6e-5 | +| 7-frame clip, forwards and backwards (9 frame results) | logits 0.0038, IoU 1.0000 | + +SAM 2 is prompted with at most three foreground and three background points per frame: it was trained on a handful +of clicks, and measured confidence collapses with many (0.99 with 1–7 points, 0.25 with 16, 0.00 with 32). + +## Speed (this PC: 16 threads, CPU only) +A 1024 × 1024 frame takes about **3.4 s** to encode (release build); a click on an encoded frame about 120 ms. Tracking a +clip costs one encode per frame, so a 300-frame clip takes on the order of 17 minutes on a CPU. Candle has no AMD/Vulkan +backend; GPU work is Phase 6 of the plan (the early Hiera blocks are memory-bound and have headroom on the CPU too). + +## Benchmark +`effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors]` runs synthetic clips with known mattes. 12 frames, +debug build, mean IoU (after the object returns): + +| Clip | Classical | SAM 2.1 Tiny | +|---|---|---| +| move | 0.982 | 0.993 | +| light | 0.981 | 0.992 | +| occlusion | 0.978 | 0.988 | +| reenter | 0.632 (0.000) | 0.975 (0.967) | + +MobileSAM's column appears when it is installed. From b72637bb55e38d05d2f996a127df39cd56f6276a Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 07:16:50 +0200 Subject: [PATCH 21/77] S2.50: bench --roto compares classical, MobileSAM and SAM 2 --mobilesam FILE; results recorded in docs/roto-brush-sam2.md (re-entry clip: 0.632 classical, 0.499 MobileSAM, 0.975 SAM 2). Co-Authored-By: Claude Sonnet 5.5 --- apps/effectcraft-cli/src/main.rs | 17 +++++++++-------- docs/roto-brush-sam2.md | 22 +++++++++++++--------- 2 files changed, 22 insertions(+), 17 deletions(-) diff --git a/apps/effectcraft-cli/src/main.rs b/apps/effectcraft-cli/src/main.rs index b60918304..27710a4a2 100644 --- a/apps/effectcraft-cli/src/main.rs +++ b/apps/effectcraft-cli/src/main.rs @@ -27,7 +27,7 @@ //! (--gpu: CPU vs GPU ms/frame for every comp at Full and Half) //! effectcraft-cli bench --ops [--small] [--layers N] [--comps N] [--footage N] everyday-operation timings //! on a large generated project (open, save, auto-save, undo/redo, timeline, Project panel) -//! effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--json] Roto Brush trackers on synthetic clips with known mattes +//! effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt] [--json] Roto Brush trackers on synthetic clips with known mattes //! (classical engine, MobileSAM and SAM 2.1 Tiny when installed; IoU, flicker, re-appearance, seconds per frame) //! effectcraft-cli script FILE.jsx [F.ecproj] | --eval CODE run an After Effects-style script //! effectcraft-cli mcp [--autosave | --bridge PORT] MCP server on stdio @@ -126,7 +126,7 @@ const USAGE: &str = "usage: effectcraft-cli Result<(), Failure> { use effectcraft_track::roto::bench; let frames = args.num("--frames")?.unwrap_or(24.0).clamp(12.0, 240.0) as usize; let session = effectcraft_host::session(); + // `--mobilesam FILE` uses the weights at FILE; otherwise the installed ones. + let mobilesam_file = args.opt("--mobilesam").map(std::path::PathBuf::from).or_else(|| session.model_path(&MOBILE_SAM)); let mobilesam: Option> = - session.model_path(&MOBILE_SAM).and_then(|p| std::fs::read(p).ok()).and_then(|bytes| effectcraft_segment::load(MOBILE_SAM.id, &bytes).ok()).and_then( - |m| match m { - Loaded::Mask(m) => Some(m), - _ => None, - }, - ); + mobilesam_file.and_then(|p| std::fs::read(p).ok()).and_then(|bytes| effectcraft_segment::load(MOBILE_SAM.id, &bytes).ok()).and_then(|m| match m { + Loaded::Mask(m) => Some(m), + _ => None, + }); if mobilesam.is_none() { note!("bench --roto: MobileSAM is not installed; comparing the classical engine only"); } diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index b2c356e74..efa7a5eb8 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -44,14 +44,18 @@ clip costs one encode per frame, so a 300-frame clip takes on the order of 17 mi backend; GPU work is Phase 6 of the plan (the early Hiera blocks are memory-bound and have headroom on the CPU too). ## Benchmark -`effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors]` runs synthetic clips with known mattes. 12 frames, -debug build, mean IoU (after the object returns): +`effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt]` runs synthetic clips +with known mattes through all three. 12 frames, debug build, this PC; mean IoU (mean IoU after the object returns), seconds +per frame: -| Clip | Classical | SAM 2.1 Tiny | -|---|---|---| -| move | 0.982 | 0.993 | -| light | 0.981 | 0.992 | -| occlusion | 0.978 | 0.988 | -| reenter | 0.632 (0.000) | 0.975 (0.967) | +| Clip | Classical | MobileSAM | SAM 2.1 Tiny | +|---|---|---|---| +| move | 0.982 | 0.992 | 0.993 | +| light | 0.981 | 0.992 | 0.992 | +| occlusion | 0.978 | 0.881 | 0.988 | +| reenter | 0.632 (0.000) | 0.499 (0.725) | 0.975 (0.967) | +| s/frame | 0.01 | 1.1 | 8.2 | -MobileSAM's column appears when it is installed. +MobileSAM prompts every frame independently from the flow-warped matte, so it follows the object while the flow does and +loses it when it is hidden; SAM 2's streaming memory is what lets it re-acquire the object. The synthetic clips are easy +(one textured disc), so the "move" and "light" rows mostly show all three working; the re-entry row is the difference. From 60ad10abebfe012f539ee3d4f821be25f2977bcf Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 07:46:42 +0200 Subject: [PATCH 22/77] S2.42: without a tracker the chain is exactly the classic engine Bit-for-bit test: every frame of the Roto Brush chain equals a direct segment_with / propagate_with call with no model. Co-Authored-By: Claude Sonnet 5.5 --- crates/engine/src/tests/roto.rs | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index b78bc2df7..b8cef6677 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -750,3 +750,33 @@ fn a_failing_tracker_falls_back_to_the_classic_engine() { } assert!(worst > 0.9, "the classic engine should have carried the clip: {worst}"); } + +/// Without a tracking model the Roto Brush chain is exactly the classic engine: every frame +/// equals a direct call of `segment_with` / `propagate_with` (no model), bit for bit. +#[test] +fn without_a_tracker_the_chain_is_exactly_the_classic_engine() { + use effectcraft_track::roto as rb; + use std::sync::Arc; + let (s, clip, uid) = painted(center, "disk-bitexact.mov"); + forget(&s, clip, uid); + let comp = s.active_comp().unwrap(); + let l = comp.layer(clip).unwrap(); + let g = l.props.find_group(uid).unwrap(); + let size = crate::roto::layer_size(&s.project, comp, l); + let chain = effectcraft_effects::roto::Chain::new(&crate::roto::params_static(g), size, comp.frame_rate.as_f64(), 1.0); + assert!(!chain.has_tracker()); + let img = |f: i64| Arc::new(draw(center(f as u32))); + let mut prev: Option<(Arc, rb::FrameSeg)> = None; + for f in 0..=5i64 { + let got = chain.ensure(f, &mut |g| Some(img(g)), &|| false).unwrap(); + let cur = img(f); + let strokes = chain.data.strokes_at(f); + let want = match &prev { + None => rb::segment_with(&cur, &strokes, None, &chain.opts, chain.refine_radius, None), + Some((pi, ps)) => rb::propagate_with(pi, ps, &cur, &strokes, &chain.opts, chain.refine_radius, None), + }; + assert_eq!(*got, want, "frame {f} differs from the classic engine"); + prev = Some((cur, want)); + } + forget(&s, clip, uid); +} From a7417b578df68996aa98a84345647c141551f9bd Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 15 Oct 2026 11:41:33 +0000 Subject: [PATCH 23/77] Roto Brush: don't run the classic engine silently under a selected model With SAM 2 (or MobileSAM) chosen, a propagate that started before the model had loaded ran the whole job on the classic engine and reported nothing. The chosen model is now in use before a job that segments with it starts (Roto Brush 2.0 / 3.0 and Object Matte): - roto.propagate and the other Roto jobs wait for a load or download under way when called with wait: true (bounded to 10 minutes). Without wait they refuse while the model is still loading, and start the load in the background. A model that is not installed, or failed to load, is refused either way. - roto.status reports model {selected, installed, active, loading, error} and a notice while the chosen model is not the one in use. source is "model" for a frame the model segmented on its own, instead of "classical". - Auto-propagate keeps edited instances pending while the model loads, and shows errors as toasts instead of dropping them. - Docs: what propagate does now; tracked frames live in the session that ran propagate and are not saved in the project file, so a separate render process does not show them; roto.model.select lists sam2.1-hiera-tiny. Tests: a selected model that is not installed refuses to propagate (no weights needed); SAM 2 selected without wait is loaded before propagate (EFFECTCRAFT_SAM2). The SAM 2 tests and the failing-tracker test take the Roto test lock, since they share the process-wide model slot. Claude-Session: https://claude.ai/code/session_0136kZtXrYo5VUWEFxWduwUD --- crates/effects/src/roto.rs | 8 ++- crates/engine/src/analysis/models.rs | 87 +++++++++++++++++++++++++++- crates/engine/src/analysis/roto.rs | 21 ++++++- crates/engine/src/cmd/roto_cmds.rs | 14 ++++- crates/engine/src/tests/roto.rs | 45 ++++++++++++++ docs/agents.md | 3 + docs/roto-brush-sam2.md | 17 ++++-- 7 files changed, 183 insertions(+), 12 deletions(-) diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 96d6b467c..ffc2be376 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -191,9 +191,15 @@ pub fn set_model(m: Option>) { } } +/// Whether an instance with `params` segments with the model Settings chose: Roto Brush 2.0 and +/// 3.0 do, 1.0 is the classic engine. +pub fn uses_model(params: &Params) -> bool { + e(params, "version") != 0 +} + /// The model an instance segments with: its version's choice of the installed one. pub fn model_for(params: &Params) -> Option> { - if e(params, "version") == 0 { + if !uses_model(params) { return None; } installed_model() diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index dbe8a1d5f..11f856d5f 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -14,10 +14,12 @@ use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; +use std::time::Duration; use effectcraft_segment::face::FaceModel; use effectcraft_segment::{self as seg, Loaded, ModelInfo, Task}; use serde_json::{Value, json}; +use web_time::Instant; use crate::commands::{CommandSpec, always, b_p, bad, str_p}; use crate::{EngineError, Result, Session, cmd, query}; @@ -126,6 +128,10 @@ fn kind(task: Task) -> &'static Kind { } } +/// How long a job that needs the Roto Brush model waits for it to load (or download) before it +/// gives up. A worker that never reports back must not hang the command. +const ROTO_MODEL_WAIT: Duration = Duration::from_secs(600); + impl Session { /// The folder installed models live in: the host's, else `models` next to the settings. pub fn models_dir(&self) -> Option { @@ -193,6 +199,85 @@ impl Session { self.models.status(Task::Mask).active } + /// The Roto Brush model Settings chose, when it is one (not the classic engine). + fn roto_model_info(&self) -> Option<&'static ModelInfo> { + seg::info(self.chosen_model(Task::Mask)).filter(|m| m.task == Task::Mask) + } + + /// Whether the chosen Roto Brush model is still loading or downloading. + pub fn roto_model_loading(&self) -> bool { + self.models.status(Task::Mask).busy.is_some() && self.chosen_model(Task::Mask) != seg::CLASSICAL + } + + /// The Roto Brush model Settings chose and its state (`roto.status`). + pub fn roto_model(&self) -> Value { + let st = self.models.status(Task::Mask); + json!({ + "selected": self.chosen_model(Task::Mask), + "installed": self.roto_model_info().is_none_or(|m| self.model_path(m).is_some()), + "active": st.active, + "loading": st.busy.is_some(), + "error": st.error, + }) + } + + /// What `roto.status` says when the chosen model is not the one in use for this layer (`None`: + /// nothing to say). `uses_model`: the layer segments with a model at all (Roto Brush 1.0 does not). + pub fn roto_model_notice(&self, uses_model: bool) -> Option { + let info = self.roto_model_info()?; + if !uses_model { + return Some(format!("{} is selected, but this layer is Roto Brush 1.0: it segments with the classic engine", info.name)); + } + let st = self.models.status(Task::Mask); + if st.active.as_deref() == Some(info.id) { + None + } else if self.model_path(info).is_none() { + Some(format!("{} is selected but not installed: propagate refuses until it is installed or Classic is chosen", info.name)) + } else if let Some(e) = st.error { + Some(format!("{} is selected but failed to load ({e}): propagate refuses until it is selected again", info.name)) + } else { + Some(format!("{} is selected but not loaded yet: propagate loads it first with wait: true, and refuses without it", info.name)) + } + } + + /// Make sure the Roto Brush model Settings chose is in use before a job that segments with it + /// starts, so the job never runs on the classic engine without saying so. With `wait`, a load + /// under way is waited for and one not started is loaded here; without it, a model that is not + /// ready is refused (and starts loading). A model that is not installed, or failed to load, is + /// refused either way. The classic engine needs nothing. + pub fn ensure_roto_model(&mut self, wait: bool) -> std::result::Result<(), String> { + let Some(info) = self.roto_model_info() else { return Ok(()) }; + // A download or load under way (from Settings, or `select` without `wait`) finishes first. + // Nothing is in flight in the browser build, so this never sleeps there. + if wait && self.models.status(Task::Mask).busy.is_some() { + let started = Instant::now(); + while self.models.status(Task::Mask).busy.is_some() { + if started.elapsed() >= ROTO_MODEL_WAIT { + return Err(format!("{} is still loading after {} minutes", info.name, ROTO_MODEL_WAIT.as_secs() / 60)); + } + std::thread::sleep(Duration::from_millis(5)); + } + } + if self.model_path(info).is_none() { + if self.models.status(Task::Mask).busy.is_some() { + return Err(format!("{} is still downloading: try again when it is ready, or pass wait: true", info.name)); + } + return Err(format!("{} is selected for Roto Brush but not installed: install or download it, or choose {}", info.name, ROTO.classic_name)); + } + let st = self.models.status(Task::Mask); + if st.active.as_deref() != Some(info.id) && st.busy.is_none() && st.error.is_none() { + self.apply_model(Task::Mask, wait); + } + let st = self.models.status(Task::Mask); + if let Some(e) = st.error { + return Err(format!("{} failed to load: {e}", info.name)); + } + if st.active.as_deref() == Some(info.id) { + return Ok(()); + } + Err(format!("{} is still loading: try again when it is ready, or pass wait: true", info.name)) + } + /// Finish what the model workers did: a toast, and loading a model that just arrived. pub fn poll_models(&mut self) { for task in [Task::Mask, Task::Face] { @@ -417,7 +502,7 @@ pub fn specs() -> Vec { "Use Roto Brush Model", [], None, - "{id: classical|mobilesam, wait?} (Roto Brush 2.0 / 3.0 use it; wait: loaded before returning)", + "{id: classical|mobilesam|sam2.1-hiera-tiny, wait?} (Roto Brush 2.0 / 3.0 use it; wait: loaded before returning)", always, |s, p| select(s, p, Task::Mask) ), diff --git a/crates/engine/src/analysis/roto.rs b/crates/engine/src/analysis/roto.rs index 12ce21260..357e77b13 100644 --- a/crates/engine/src/analysis/roto.rs +++ b/crates/engine/src/analysis/roto.rs @@ -621,9 +621,21 @@ impl Session { } } + /// Whether `effect` on `layer` segments with the model Settings chose: Object Matte always, + /// Roto Brush unless it is version 1.0 (the classic engine). + pub(crate) fn uses_roto_model(&self, comp: ItemId, layer: LayerId, effect: Uid) -> bool { + let g = self.project.comp(comp).and_then(|c| c.layer(layer)).and_then(|l| l.effects()).and_then(|f| f.groups().find(|g| g.uid == effect)); + g.is_some_and(|g| is_object_matte(g) || fx::uses_model(¶ms_static(g))) + } + /// Start propagating (or freezing) Roto Brush `effect` on `layer`. `wait`: block until done /// (wasm32 always does). Returns the number of frames. pub fn start_roto(&mut self, comp: ItemId, layer: LayerId, effect: Uid, task: RotoTask, dir: Direction, wait: bool) -> Result { + // The model must be in use before the job starts: otherwise it would run on the classic + // engine without saying so. + if self.uses_roto_model(comp, layer, effect) { + self.ensure_roto_model(wait)?; + } if self.is_roto_running() { if !wait { return Err("a Roto Brush job is already running".into()); @@ -786,11 +798,14 @@ impl Session { pub fn poll_roto(&mut self, auto: bool) -> bool { self.poll_models(); let changed = self.poll_roto_job(); - if auto && self.roto_job.is_none() && self.offloaded_roto().is_none() && self.state.roto.auto_propagate { + // Edited instances stay pending while the model they will segment with is still loading. + let loading = self.roto_model_loading(); + if auto && self.roto_job.is_none() && self.offloaded_roto().is_none() && self.state.roto.auto_propagate && !loading { while let Some((c, l, e)) = self.roto_pending.first().copied() { self.roto_pending.remove(0); - if self.start_roto(c, l, e, RotoTask::Propagate, Direction::Both, false).is_ok() { - return true; + match self.start_roto(c, l, e, RotoTask::Propagate, Direction::Both, false) { + Ok(_) => return true, + Err(message) => self.events.push(Event::Toast { message, error: true }), } } } diff --git a/crates/engine/src/cmd/roto_cmds.rs b/crates/engine/src/cmd/roto_cmds.rs index c37adb0c2..c7f190fde 100644 --- a/crates/engine/src/cmd/roto_cmds.rs +++ b/crates/engine/src/cmd/roto_cmds.rs @@ -431,6 +431,9 @@ fn status(s: &mut Session, p: &Value) -> Result { out["strokeFrames"] = json!(d.stroke_frames()); out["frozen"] = json!(frozen_of(&g)); out["pending"] = json!(s.roto_pending.contains(&(cid, lid, uid))); + // The model Settings chose and whether this layer segments with it (propagate refuses or waits otherwise). + out["model"] = s.roto_model(); + out["notice"] = json!(s.roto_model_notice(s.uses_roto_model(cid, lid, uid))); out["limits"] = json!(rt::frame_limits(&comp, &layer)); let frame = frame_p(s, p, cid, &layer); out["frame"] = json!(frame); @@ -464,7 +467,8 @@ fn status(s: &mut Session, p: &Value) -> Result { } } // Where each frame's matte came from, and how far to trust it: frames a tracker followed - // report its confidence and whether the object was absent; the rest came from the classic engine. + // report its confidence and whether the object was absent; frames the model segmented on + // their own report "model"; the rest came from the classic engine. { let size = rt::layer_size(&s.project, &comp, &layer); let chain = fx::Chain::new(¶ms_static(&g), size, comp.frame_rate.as_f64(), 1.0); @@ -474,7 +478,13 @@ fn status(s: &mut Session, p: &Value) -> Result { .filter_map(|(f, k)| fx::trust(*k).map(|t| (f.to_string(), json!({"confidence": t.confidence, "occluded": t.occluded})))) .collect(); out["tracker"] = json!({"model": s.models_roto_active(), "frames": tracked}); - out["source"] = json!(if tracked.contains_key(&frame.to_string()) { "tracker" } else { "classical" }); + out["source"] = json!(if tracked.contains_key(&frame.to_string()) { + "tracker" + } else if chain.model.is_some() { + "model" + } else { + "classical" + }); } let computed = s.roto_computed(cid, lid, uid); out["computed"] = json!(computed.iter().filter(|(_, c)| **c).map(|(f, _)| *f).collect::>()); diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index b8cef6677..9ef6a848f 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -677,6 +677,7 @@ fn object_matte_mattes_persist_in_the_disk_cache() { fn sam2_tracks_the_clip_through_roto_propagate() { let Ok(weights) = std::env::var("EFFECTCRAFT_SAM2") else { return }; let dir = std::env::temp_dir().join(format!("ec-roto-sam2-{}", std::process::id())); + hold_roto_cache_test_lock(); let (mut s, clip, _) = painted(center, "disk-sam2.mov"); s.models_dir = Some(dir.join("models")); std::fs::create_dir_all(dir.join("models")).unwrap(); @@ -707,6 +708,49 @@ fn sam2_tracks_the_clip_through_roto_propagate() { let _ = std::fs::remove_dir_all(dir); } +/// Choosing a model that is not installed must not fall back to the classic engine silently: +/// propagate refuses and `roto.status` says why. Choosing the classic engine runs it again. +#[test] +fn a_selected_model_that_is_not_installed_refuses_to_propagate() { + hold_roto_cache_test_lock(); + let (mut s, clip, _) = painted(center, "disk-not-installed.mov"); + s.models_dir = Some(std::env::temp_dir().join(format!("ec-roto-none-{}", std::process::id()))); + s.execute_checked("roto.model.select", json!({"id": "sam2.1-hiera-tiny"})).unwrap(); + let st = s.execute("roto.status", json!({"layer": clip.0})).unwrap(); + assert_eq!(st["model"]["selected"], "sam2.1-hiera-tiny", "{st}"); + assert_eq!(st["model"]["installed"], false, "{st}"); + assert!(st["notice"].as_str().unwrap().contains("not installed"), "{st}"); + let e = s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})).unwrap_err().to_string(); + assert!(e.contains("not installed"), "{e}"); + assert_eq!(s.execute("roto.status", json!({"layer": clip.0, "frame": 10})).unwrap()["source"], "classical"); + s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); + let r = s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})).unwrap(); + assert_eq!(r["frames"], 21, "{r}"); + assert!(s.execute("roto.status", json!({"layer": clip.0})).unwrap()["notice"].is_null()); +} + +/// Selecting SAM 2 without `wait` and propagating at once: propagate loads the model first, so the +/// tracker (not the classic engine) does the work. Needs EFFECTCRAFT_SAM2 like the test above. +#[cfg(not(target_arch = "wasm32"))] +#[test] +fn sam2_selected_without_wait_is_loaded_before_propagate() { + let Ok(weights) = std::env::var("EFFECTCRAFT_SAM2") else { return }; + hold_roto_cache_test_lock(); + let dir = std::env::temp_dir().join(format!("ec-roto-sam2-wait-{}", std::process::id())); + let (mut s, clip, _) = painted(center, "disk-sam2-wait.mov"); + s.models_dir = Some(dir.join("models")); + std::fs::create_dir_all(dir.join("models")).unwrap(); + s.execute_checked("roto.model.install", json!({"path": weights})).unwrap(); + s.execute_checked("roto.model.select", json!({"id": "sam2.1-hiera-tiny"})).unwrap(); + let r = s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})).unwrap(); + assert_eq!(r["frames"], 21, "{r}"); + let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); + assert_eq!(st["source"], "tracker", "{st}"); + assert_eq!(st["tracker"]["model"], "sam2.1-hiera-tiny"); + assert!(st["notice"].is_null(), "{st}"); + let _ = std::fs::remove_dir_all(dir); +} + /// A model whose tracker always fails: Roto Brush must carry on with the classic engine and say so. #[test] fn a_failing_tracker_falls_back_to_the_classic_engine() { @@ -735,6 +779,7 @@ fn a_failing_tracker_falls_back_to_the_classic_engine() { Err("the tracker broke".into()) } } + hold_roto_cache_test_lock(); let (mut s, clip, _) = painted(center, "disk-broken-tracker.mov"); effectcraft_effects::roto::set_model(Some(Arc::new(Broken))); let r = s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})); diff --git a/docs/agents.md b/docs/agents.md index e0c0055ec..54ba21afa 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -508,6 +508,9 @@ Trained model (Roto Brush 2.0 / 3.0): `roto.model.install {path, id?}` installs a file you already have. - `roto.model.select {id: classical|mobilesam|sam2.1-hiera-tiny}` chooses which one Roto Brush uses (SAM 2 also tracks the object through the clip, see [roto-brush-sam2.md](roto-brush-sam2.md); `roto.status` then reports `source` and per-frame `tracker.frames` with `confidence` and `occluded`). Loading runs in the background; `"wait": true` (also on download and install) returns once it is in use. +- `roto.propagate` never falls back to the classic engine silently: a chosen model that is not installed, or failed + to load, is an error, and without `"wait": true` so is a model still loading. `roto.status` reports `model` + (`selected`, `installed`, `active`, `loading`, `error`) and a `notice` while the chosen model is not the one in use. - `roto.model.remove {id}` deletes the installed weights. - The `version` property (`effects/#1/version`: 0 = 1.0, 1 = 2.0, 2 = 3.0) picks between the classic engine (1.0) and the chosen model (2.0, 3.0). diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index efa7a5eb8..f5946a98e 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -12,12 +12,19 @@ Apache-2.0). It is optional: nothing changes until you install the model and cho whole span in both directions, using SAM 2's streaming memory, so it keeps tracking through occlusion and when the object leaves the frame and returns. Each frame's probability map is then refined by the same graph cut as every other Roto Brush frame: your strokes stay hard constraints and edges snap to colour. -4. `roto.status {frame}` reports `source` (`tracker` or `classical`) and, for tracked frames, `tracker.frames[f]`: - `confidence` (0–1, the model's object score times its predicted mask quality) and `occluded` (the object is not - in the frame: an empty matte). +4. `roto.status {frame}` reports `source` (`tracker`, `model` for a frame the model segmented on its own, or + `classical`) and, for tracked frames, `tracker.frames[f]`: `confidence` (0–1, the model's object score times its + predicted mask quality) and `occluded` (the object is not in the frame: an empty matte). `model` says what Settings + chose and whether it is installed, loaded and loading; `notice` says when it is not the one in use for the layer. -If the model is missing, fails, or is deselected, the classic engine runs exactly as before (`roto.status` says -`classical`, and the job's `trackerNote` says why). +Propagate never falls back to the classic engine without saying so. With SAM 2 chosen but not installed, or failing to +load, it refuses with the reason. While it is still loading, `wait: true` (agents) loads it first; without `wait` it +refuses, and the load goes on in the background. Deselecting SAM 2 (choosing Classic) runs Propagate exactly as before, +and Roto Brush 1.0 never uses a model. + +Tracked frames live in the session that ran Propagate and are not saved in the project file. A separate render process +(`effectcraft-cli render-frame`) does not have them, so render in the same session: in the app, or in one `run` that +includes `roto.propagate` and then `comp.saveFrameAs`. ## How it works `crates/sam2` is a Rust port (on [candle](https://github.com/huggingface/candle), native only) of the Hugging Face From 1eee687ebc7062a60bc0c269338e86d57c9e5267 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 15 Oct 2026 13:38:06 +0000 Subject: [PATCH 24/77] Roto Brush: save the tracked frames in the project The tracked frames lived only in memory, so a saved project lost them: a render in another process (effectcraft-cli render-frame) segmented the frames again with the classic engine. A propagate with a tracker now saves the frames it followed with the instance, in a hidden parameter beside Freeze's (TrackCache: each frame's matte, Refine Edge band, colour models and confidence, keyed by the inputs it was tracked from). - They apply while the strokes and span they came from are unchanged, whichever model Settings choose, so a project renders the same when it reopens, without the model loaded or chosen. - A stroke or span change makes them stale until the next Propagate. A Propagate replaces them; with the classic engine it leaves none. - roto.status reports tracker.saved (the model they came from) and a notice when Settings choose another model. - The save at job completion is not an undo step (the frames are a cache of the strokes) and is skipped if the strokes changed while tracking. Decoding checks the saved size, so a corrupt project is ignored rather than overflowing. Object Matte saves nothing. Projects saved before this change open unchanged (the parameter is filled in on open). Tests: a saved project reopens with the model unloaded and renders the saved frames through Chain::ensure; choosing the classic engine keeps them until a classic Propagate replaces them; a stroke change makes them stale. The SAM 2 tests and the failing-tracker test take the Roto test lock. Claude-Session: https://claude.ai/code/session_0136kZtXrYo5VUWEFxWduwUD --- crates/effects/src/object_matte.rs | 2 +- crates/effects/src/roto.rs | 187 ++++++++++++++++++++++++++- crates/engine/src/analysis/models.rs | 12 ++ crates/engine/src/analysis/roto.rs | 28 +++- crates/engine/src/cmd/roto_cmds.rs | 11 +- crates/engine/src/tests/roto.rs | 117 ++++++++++++++++- docs/agents.md | 3 + docs/roto-brush-sam2.md | 9 +- 8 files changed, 353 insertions(+), 16 deletions(-) diff --git a/crates/effects/src/object_matte.rs b/crates/effects/src/object_matte.rs index 6804ea2c7..4b5e1fd42 100644 --- a/crates/effects/src/object_matte.rs +++ b/crates/effects/src/object_matte.rs @@ -227,7 +227,7 @@ pub fn job_frames(params: &Params, limits: [i64; 2]) -> Vec { pub fn chain(params: &Params, layer_size: [f64; 2], fps: f64, f: i64, scale: f64, model: Option>) -> Option { let data = roto_data(params, layer_size, f, model.is_some())?; let keys = data.chain_keys(seed(params, layer_size, fps, model.as_deref())); - Some(Chain { data: Arc::new(data), keys, opts: seg_opts(scale), refine_radius: 0.0, model }) + Some(Chain { data: Arc::new(data), keys, opts: seg_opts(scale), refine_radius: 0.0, model, saved: None }) } /// Frame `f`'s segmentation at full resolution, computing what is missing (the background diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index ffc2be376..9596b6d50 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -31,7 +31,7 @@ use effectcraft_project::ParamUi; use effectcraft_raster::Image; use effectcraft_segment::{ClipFrame, MaskModel, TrackRequest}; use effectcraft_track::roto::refine::{self, MatteParams}; -use effectcraft_track::roto::{self as rb, RotoData, SegOpts, rle}; +use effectcraft_track::roto::{self as rb, Gmm, RotoData, SegOpts, rle}; use rayon::prelude::*; use serde::{Deserialize, Serialize}; @@ -46,6 +46,8 @@ pub const ID: &str = "ec.matte.rotobrush"; pub const STROKES: &str = "strokes"; pub const INPUT_KEY: &str = "inputKey"; pub const FROZEN: &str = "frozen"; +/// The saved tracking (see [`TrackCache`]). +pub const TRACKED: &str = "tracked"; /// Default segmentation span each side of the base frame (frames). pub const DEFAULT_SPAN: i64 = 20; @@ -95,6 +97,7 @@ pub fn specs() -> Vec { p(STROKES, "Strokes", Value::Str(String::new()), ParamUi::Hidden), p(INPUT_KEY, "Input Key", Value::Str(String::new()), ParamUi::Hidden), p(FROZEN, "Frozen", Value::Str(String::new()), ParamUi::Hidden), + p(TRACKED, "Tracked", Value::Str(String::new()), ParamUi::Hidden), ], render, gpu: false, @@ -214,6 +217,11 @@ pub fn installed_model() -> Option> { /// Seed of the chain keys: what the frames are (the engine-maintained Input Key), the layer /// size, the frame rate, the segmentation settings and the model. pub fn seed(params: &Params, layer_size: [f64; 2], fps: f64) -> u64 { + seed_with(params, layer_size, fps, model_for(params).map_or(effectcraft_segment::CLASSICAL, |m| m.info().id)) +} + +/// [`seed`] with the segmenting model named by its id (saved tracking names the one that made it). +fn seed_with(params: &Params, layer_size: [f64; 2], fps: f64, model: &str) -> u64 { let mut h = 0xcbf2_9ce4_8422_2325; fnv(&mut h, s(params, INPUT_KEY).as_bytes()); for v in [ @@ -227,7 +235,7 @@ pub fn seed(params: &Params, layer_size: [f64; 2], fps: f64) -> u64 { fnv(&mut h, &v.to_bits().to_le_bytes()); } fnv(&mut h, &[e(params, "quality") as u8, e(params, "version") as u8]); - fnv(&mut h, model_for(params).map_or(effectcraft_segment::CLASSICAL, |m| m.info().id).as_bytes()); + fnv(&mut h, model.as_bytes()); h } @@ -388,13 +396,57 @@ pub struct Chain { pub refine_radius: f64, /// The trained model (Version 2.0 / 3.0 with one installed). pub model: Option>, + /// The saved tracking, while it applies (see [`TrackCache`]). + pub(crate) saved: Option, } impl Chain { pub fn new(params: &Params, layer_size: [f64; 2], fps: f64, scale: f64) -> Chain { let data = data(params); let keys = data.chain_keys(seed(params, layer_size, fps)); - Chain { refine_radius: data.refine_radius(), data, keys, opts: seg_opts(params, scale), model: model_for(params) } + let saved = saved_for(params, &data, layer_size, fps); + Chain { refine_radius: data.refine_radius(), data, keys, opts: seg_opts(params, scale), model: model_for(params), saved } + } + + /// Frame `f`'s segmentation at `scale` from the saved tracking or the caches; computes nothing. + pub fn cached_seg(&self, f: i64, scale: f64) -> Option> { + self.saved_seg(f, scale).or_else(|| self.keys.get(&f).and_then(|k| cached(*k, scale))) + } + + /// The tracker's trust in frame `f`: from the saved tracking, else from this session's tracking. + pub fn trust_at(&self, f: i64) -> Option { + if let Some(fr) = self.saved.as_ref().and_then(|t| t.cache.frames.get(&f)) { + return Some(FrameTrust { confidence: fr.c, occluded: fr.o }); + } + trust(*self.keys.get(&f)?) + } + + /// The model the saved tracking was made with, while it applies. + pub fn saved_model(&self) -> Option<&str> { + self.saved.as_ref().map(|t| t.cache.model.as_str()) + } + + /// Frame `f` of the saved tracking at `scale` (decoded once, then cached). + fn saved_seg(&self, f: i64, scale: f64) -> Option> { + let t = self.saved.as_ref()?; + let k = tracked_key(*t.keys.get(&f)?); + if let Some(s) = cached(k, scale) { + return Some(s); + } + let full = match cached(k, 1.0) { + Some(s) => s, + None => { + let s = Arc::new(decode_tracked(&t.cache, f)?); + store(k, 1.0, s.clone()); + s + } + }; + if scale == 1.0 { + return Some(full); + } + let s = Arc::new(downscale(&full, scale)); + store(k, scale, s.clone()); + Some(s) } /// Frame `f`'s segmentation at this chain's scale: cached, or computed from the nearest @@ -403,9 +455,12 @@ impl Chain { pub fn ensure(&self, f: i64, frame: &mut dyn FnMut(i64) -> Option>, cancel: &dyn Fn() -> bool) -> Option> { let sc = self.opts.scale; let lookup = |g: i64| -> Option> { + // A frame the tracker followed wins over the classic engine's: the saved tracking, else + // this session's. + if let Some(s) = self.saved_seg(g, sc) { + return Some(s); + } let k = *self.keys.get(&g)?; - // A frame the tracker followed wins over the classic engine's (memory only: its - // confidence is not persisted). if let Some(s) = cached(tracked_key(k), sc) { return Some(s); } @@ -654,6 +709,128 @@ fn frozen(params: &Params) -> Option> { Some(d) } +// ------------------------------------------------------------------ saved tracking + +/// The tracker's frames of an instance, saved with the project (the hidden [`TRACKED`] parameter, +/// as [`FrozenCache`] is). They apply while `key` matches the instance's inputs ([`tracks_key`]), +/// whichever model Settings choose: a stroke or span change makes them stale, and the next Propagate +/// replaces them (with the classic engine, that leaves none). +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct TrackCache { + /// What the frames were tracked from ([`tracks_key`]). + pub key: String, + /// The tracker's model id. + pub model: String, + /// The mattes' width and height (the layer's size). + pub size: [usize; 2], + pub frames: BTreeMap, +} + +/// One tracked frame: its matte and Refine Edge band (RLE, base64), its colour models and the +/// tracker's trust in it. +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(default)] +pub struct TrackFrame { + /// Matte (RLE, base64). + pub m: String, + /// Refine Edge band (RLE, base64). + pub r: String, + pub fg: Gmm, + pub bg: Gmm, + /// 0–1: the tracker's confidence. + pub c: f32, + /// The object is not in the frame. + pub o: bool, +} + +impl TrackCache { + pub fn from_json(s: &str) -> Option { + if s.trim().is_empty() { + return None; + } + serde_json::from_str(s).ok() + } + pub fn to_json(&self) -> String { + serde_json::to_string(self).unwrap_or_default() + } +} + +/// What saved tracking depends on: the chain keys of its frames (the inputs and the model). +pub fn tracks_key(keys: &BTreeMap) -> String { + let mut h = 0xcbf2_9ce4_8422_2325; + for (f, k) in keys { + fnv(&mut h, &f.to_le_bytes()); + fnv(&mut h, &k.to_le_bytes()); + } + format!("{h:016x}") +} + +/// The saved tracking of `params`, parsed once per content. +fn saved_tracks(params: &Params) -> Option> { + static C: OnceLock>>> = OnceLock::new(); + let j = s(params, TRACKED); + if j.is_empty() { + return None; + } + let mut h = 0xcbf2_9ce4_8422_2325; + fnv(&mut h, j.as_bytes()); + let c = C.get_or_init(Default::default); + if let Some(d) = c.lock().unwrap_or_else(|e| e.into_inner()).get(&h) { + return Some(d.clone()); + } + let d = Arc::new(TrackCache::from_json(j)?); + let mut m = c.lock().unwrap_or_else(|e| e.into_inner()); + if m.len() > 8 { + m.clear(); + } + m.insert(h, d.clone()); + Some(d) +} + +/// A saved tracking that applies to a chain: its frames, keyed under the model that made them. +pub(crate) struct SavedTracking { + cache: Arc, + keys: BTreeMap, +} + +/// The saved tracking of `params` if it applies: the inputs are the instance's now. It does not depend +/// on which model Settings choose or has loaded, so a project renders the same anywhere. +fn saved_for(params: &Params, data: &RotoData, layer_size: [f64; 2], fps: f64) -> Option { + let cache = saved_tracks(params)?; + let keys = data.chain_keys(seed_with(params, layer_size, fps, &cache.model)); + (tracks_key(&keys) == cache.key).then_some(SavedTracking { cache, keys }) +} + +/// Decode a tracked frame of saved tracking (full resolution). +fn decode_tracked(cache: &TrackCache, f: i64) -> Option { + let fr = cache.frames.get(&f)?; + // The size comes from the project file: checked, so a corrupt one is ignored rather than overflowing. + let [w, h] = cache.size; + let n = w.checked_mul(h)?; + let matte = rle::decode(&fr.m)?; + let refine = rle::decode(&fr.r)?; + (matte.len() == n && refine.len() == n).then(|| FrameSeg { w, h, matte, refine, fg: fr.fg.clone(), bg: fr.bg.clone() }) +} + +/// The tracking a chain has just made, to save with its instance: every frame of its span that the +/// tracker followed, at full resolution, with its trust. `None` when it followed none. +pub fn tracks_of(ch: &Chain) -> Option { + let model = ch.model.as_ref()?.info().id; + let mut size = None; + let mut frames = BTreeMap::new(); + for (f, k) in &ch.keys { + let (Some(seg), Some(t)) = (cached(tracked_key(*k), 1.0), trust(*k)) else { continue }; + size.get_or_insert([seg.w, seg.h]); + frames.insert( + *f, + TrackFrame { m: rle::encode(&seg.matte), r: rle::encode(&seg.refine), fg: seg.fg.clone(), bg: seg.bg.clone(), c: t.confidence, o: t.occluded }, + ); + } + let cache = TrackCache { key: tracks_key(&ch.keys), model: model.to_string(), size: size?, frames }; + (!cache.frames.is_empty()).then_some(cache) +} + /// Encode one frame for [`FrozenCache`]. pub fn freeze_frame(alpha: &[f32], band: &[u8]) -> FrozenFrame { let a: Vec = alpha.iter().map(|v| (v.clamp(0.0, 1.0) * 255.0).round() as u8).collect(); diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index 11f856d5f..bad9fd5ff 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -240,6 +240,18 @@ impl Session { } } + /// What `roto.status` says when the saved tracking applies but Settings choose another model (`None`: + /// they agree, or nothing is saved). + pub fn saved_tracking_notice(&self, saved: Option<&str>) -> Option { + let saved = saved?; + let chosen = self.chosen_model(Task::Mask); + if saved == chosen { + return None; + } + let name = |id: &str| seg::info(id).map_or(ROTO.classic_name, |m| m.name); + Some(format!("The saved tracking was made with {}; {} is chosen: Propagate again to replace it", name(saved), name(chosen))) + } + /// Make sure the Roto Brush model Settings chose is in use before a job that segments with it /// starts, so the job never runs on the classic engine without saying so. With `wait`, a load /// under way is waited for and one not started is loaded here; without it, a model that is not diff --git a/crates/engine/src/analysis/roto.rs b/crates/engine/src/analysis/roto.rs index 357e77b13..286b95095 100644 --- a/crates/engine/src/analysis/roto.rs +++ b/crates/engine/src/analysis/roto.rs @@ -22,7 +22,7 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use effectcraft_effects::object_matte as om; -use effectcraft_effects::roto::{self as fx, FROZEN, FrozenCache, INPUT_KEY}; +use effectcraft_effects::roto::{self as fx, FROZEN, FrozenCache, INPUT_KEY, TRACKED}; use effectcraft_effects::{Params, flatten_params}; use effectcraft_keyframe::Value; use effectcraft_project::{GroupKind, ItemId, ItemKind, Layer, LayerId, Project, PropGroup, Uid}; @@ -135,6 +135,8 @@ pub struct RotoShared { pub cancel: AtomicBool, pub state: Mutex, frozen: Mutex>, + /// The tracker's frames, as the saved tracking (JSON) to store with the instance. + saved: Mutex>, } pub struct RotoJob { @@ -253,6 +255,7 @@ pub(crate) fn sync(before: &Project, after: &mut Project) -> Vec<(ItemId, LayerI set_str(g, INPUT_KEY, key); if had { set_str(g, FROZEN, String::new()); + set_str(g, TRACKED, String::new()); } } } @@ -459,11 +462,19 @@ async fn run_work(w: Work, shared: &RotoShared) { s.finished = true; return; } + // The frames it followed replace the saved tracking (see `poll_roto_job`): none if it failed. + *lock(&shared.saved) = Some(match &result { + Ok(_) => fx::tracks_of(&chain).map(|t| t.to_json()).unwrap_or_default(), + Err(_) => String::new(), + }); let mut s = lock(&shared.state); match result { Ok(summary) => s.tracked = summary, Err(e) => s.tracker_note = Some(format!("tracker unavailable, classic engine used: {e}")), } + } else if w.task == RotoTask::Propagate && !w.object_matte { + // With the classic engine, Propagate leaves no saved tracking behind. + *lock(&shared.saved) = Some(String::new()); } for (k, f) in w.frames.iter().enumerate() { if cancel() { @@ -738,6 +749,17 @@ impl Session { Ok(n) } + /// Save the tracker's frames in the instance, unless its strokes changed while it tracked. No undo + /// step: the frames are a cache of the strokes and the span (a stale copy is ignored). + fn save_tracks(&mut self, cid: ItemId, lid: LayerId, uid: Uid, strokes: &str, json: String) { + let p = Arc::make_mut(&mut self.project); + if let Some(g) = p.comp_mut(cid).and_then(|c| c.layer_mut(lid)).and_then(|l| l.props.find_group_mut(uid)) + && guard_of(g) == strokes + { + set_str(g, TRACKED, json); + } + } + /// Finish a done job: Freeze writes the frozen mattes (one undo step). fn poll_roto_job(&mut self) -> bool { let Some(job) = &self.roto_job else { return false }; @@ -764,6 +786,10 @@ impl Session { RotoTask::Propagate => { let verb = if job.object_matte { "tracked" } else { "propagated" }; self.events.push(Event::Toast { message: format!("{name}: {verb} {} frame(s) in {:.1} s", st.done, st.elapsed), error: false }); + // The frames the tracker followed are saved with the instance. + if let Some(json) = lock(&job.shared.saved).take() { + self.save_tracks(job.comp, job.layer, job.effect, &job.strokes, json); + } // Rendered frames depend on the cache: redraw. self.bump(); } diff --git a/crates/engine/src/cmd/roto_cmds.rs b/crates/engine/src/cmd/roto_cmds.rs index c7f190fde..0931236a2 100644 --- a/crates/engine/src/cmd/roto_cmds.rs +++ b/crates/engine/src/cmd/roto_cmds.rs @@ -433,7 +433,6 @@ fn status(s: &mut Session, p: &Value) -> Result { out["pending"] = json!(s.roto_pending.contains(&(cid, lid, uid))); // The model Settings chose and whether this layer segments with it (propagate refuses or waits otherwise). out["model"] = s.roto_model(); - out["notice"] = json!(s.roto_model_notice(s.uses_roto_model(cid, lid, uid))); out["limits"] = json!(rt::frame_limits(&comp, &layer)); let frame = frame_p(s, p, cid, &layer); out["frame"] = json!(frame); @@ -444,7 +443,7 @@ fn status(s: &mut Session, p: &Value) -> Result { let fps = comp.frame_rate.as_f64(); let ps = params_static(&g); let chain = fx::Chain::new(&ps, size, fps, 1.0); - let seg = match chain.keys.get(&frame).and_then(|k| fx::cached(*k, 1.0)) { + let seg = match chain.cached_seg(frame, 1.0) { Some(sg) => Some(sg), None if compute => { let index = layer.effects().and_then(|f| f.groups().position(|x| x.uid == uid)).unwrap_or(0); @@ -474,10 +473,12 @@ fn status(s: &mut Session, p: &Value) -> Result { let chain = fx::Chain::new(¶ms_static(&g), size, comp.frame_rate.as_f64(), 1.0); let tracked: serde_json::Map = chain .keys - .iter() - .filter_map(|(f, k)| fx::trust(*k).map(|t| (f.to_string(), json!({"confidence": t.confidence, "occluded": t.occluded})))) + .keys() + .filter_map(|f| chain.trust_at(*f).map(|t| (f.to_string(), json!({"confidence": t.confidence, "occluded": t.occluded})))) .collect(); - out["tracker"] = json!({"model": s.models_roto_active(), "frames": tracked}); + out["tracker"] = json!({"model": s.models_roto_active(), "saved": chain.saved_model(), "frames": tracked}); + let notice = s.roto_model_notice(s.uses_roto_model(cid, lid, uid)).or_else(|| s.saved_tracking_notice(chain.saved_model())); + out["notice"] = json!(notice); out["source"] = json!(if tracked.contains_key(&frame.to_string()) { "tracker" } else if chain.model.is_some() { diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index 9ef6a848f..b18cd386d 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -700,9 +700,14 @@ fn sam2_tracks_the_clip_through_roto_propagate() { assert_eq!(frames.len(), 21); assert!(frames.values().all(|t| t["occluded"] == false && t["confidence"].as_f64().unwrap() > 0.3), "{frames:?}"); assert_eq!(st["tracker"]["model"], "sam2.1-hiera-tiny"); - // Back to the classic engine: same strokes, its own segmentation, no tracker data. + // Settings choose the classic engine: the saved tracking stays until a Propagate replaces it. s.execute_checked("roto.model.select", json!({"id": "classical", "wait": true})).unwrap(); let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); + assert_eq!(st["source"], "tracker", "{st}"); + assert!(st["notice"].as_str().unwrap().contains("Propagate again"), "{st}"); + // A classic Propagate: its own segmentation, no tracker data. + s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})).unwrap(); + let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); assert_eq!(st["source"], "classical", "{st}"); assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 0); let _ = std::fs::remove_dir_all(dir); @@ -751,6 +756,116 @@ fn sam2_selected_without_wait_is_loaded_before_propagate() { let _ = std::fs::remove_dir_all(dir); } +/// A tracker that follows the disc exactly: its true mask on every frame of the request. A stand-in +/// for SAM 2, so the saved tracking can be tested without the weights. +struct Exact; + +impl effectcraft_segment::MaskModel for Exact { + fn info(&self) -> &'static effectcraft_segment::ModelInfo { + &effectcraft_segment::SAM2_TINY + } + fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &effectcraft_segment::Prompt) -> effectcraft_segment::Result> { + Err("no weights".into()) + } + fn tracker(&self) -> Option<&dyn effectcraft_segment::ObjectTracker> { + Some(self) + } +} + +impl effectcraft_segment::ObjectTracker for Exact { + fn track( + &self, + req: &effectcraft_segment::TrackRequest, + frame: &mut dyn FnMut(i64) -> Option>, + _: &dyn Fn() -> bool, + out: &mut dyn FnMut(i64, effectcraft_segment::TrackedFrame), + ) -> effectcraft_segment::Result<()> { + for g in req.first..req.first + req.count as i64 { + let Some(img) = frame(g) else { return Err(format!("no frame {g}")) }; + let c = center(g as u32); + let prob = (0..img.w * img.h) + .map(|i| ((((i % img.w) as f64 + 0.5 - c[0]).hypot((i / img.w) as f64 + 0.5 - c[1])) <= R) as u8 as f32) + .collect::>(); + out(g, effectcraft_segment::TrackedFrame { prob, w: img.w, h: img.h, confidence: 0.99, occluded: false }); + } + Ok(()) + } +} + +/// Saved tracking is part of the project: written with the instance, it applies when the project +/// opens, with or without the model loaded, and Settings choosing the classic engine leaves it in place +/// until a Propagate replaces it. +#[test] +fn saved_tracking_is_kept_in_the_project_and_applies_when_it_reopens() { + let sam2 = "sam2.1-hiera-tiny"; + let (s, clip, uid, path) = saved_project("saved"); + // Reopened with SAM 2 chosen but not loaded (none is installed here): the saved frames apply. + forget(&s, clip, uid); + let mut s2 = Session::default(); + s2.execute_checked("file.open", json!({"path": path})).unwrap(); + s2.execute_checked("roto.model.select", json!({"id": sam2})).unwrap(); + let st = s2.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); + assert_eq!(st["source"], "tracker", "{st}"); + assert_eq!(st["tracker"]["saved"], sam2, "{st}"); + assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 21, "{st}"); + let j = iou_at(&mut s2, clip, 12, center(12)); + assert!(j > 0.9, "IoU of a saved frame: {j}"); + // The render path (Chain::ensure) reads the same frame. + let comp = s2.active_comp().unwrap(); + let l = comp.layer(clip).unwrap(); + let g = l.props.find_group(uid).unwrap(); + let size = crate::roto::layer_size(&s2.project, comp, l); + let ch = effectcraft_effects::roto::Chain::new(&crate::roto::params_static(g), size, comp.frame_rate.as_f64(), 1.0); + let seg = ch.ensure(12, &mut |_| None, &|| false).expect("the render path reads the saved frame"); + let truth = rle::decode(>(center(12))).unwrap(); + let inter = seg.matte.iter().zip(&truth).filter(|(a, b)| **a != 0 && **b != 0).count(); + let uni = seg.matte.iter().zip(&truth).filter(|(a, b)| **a != 0 || **b != 0).count(); + assert!(inter as f64 / uni as f64 > 0.9, "the render path's frame: IoU {}", inter as f64 / uni as f64); + // Settings choose the classic engine: the saved frames still apply, and the notice says so. + s2.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); + let st = s2.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); + assert_eq!(st["source"], "tracker", "{st}"); + assert!(st["notice"].as_str().unwrap().contains("Propagate again"), "{st}"); + // A classic Propagate replaces the saved tracking with nothing. + s2.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})).unwrap(); + let st = s2.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); + assert_eq!(st["source"], "classical", "{st}"); + assert!(st["tracker"]["saved"].is_null() && st["notice"].is_null(), "{st}"); + let _ = std::fs::remove_dir_all(std::path::Path::new(&path).parent().unwrap()); +} + +/// A tracked clip saved to a project file: the session that tracked it, its layer and instance, and the +/// file. +fn saved_project(name: &str) -> (Session, LayerId, u64, String) { + let (mut s, clip, uid) = painted(center, &format!("disk-{name}.mov")); + effectcraft_effects::roto::set_model(Some(std::sync::Arc::new(Exact))); + let r = s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})); + effectcraft_effects::roto::set_model(None); + assert_eq!(r.unwrap()["frames"], 21); + let dir = std::env::temp_dir().join(format!("ec-roto-{name}-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("tracked.ecproj").to_string_lossy().to_string(); + s.execute_checked("file.saveAs", json!({"path": path})).unwrap(); + let saved = std::fs::read_to_string(&path).unwrap(); + assert!(saved.contains("sam2.1-hiera-tiny"), "the project file carries the tracking and the model that made it"); + (s, clip, uid, path) +} + +/// A stroke change makes the saved tracking stale: the instance is classical until the next Propagate. +#[test] +fn a_stroke_change_makes_the_saved_tracking_stale() { + let (_s, clip, _uid, path) = saved_project("stale"); + let mut s2 = Session::default(); + s2.execute_checked("file.open", json!({"path": path})).unwrap(); + s2.execute_checked("roto.model.select", json!({"id": "sam2.1-hiera-tiny"})).unwrap(); + assert_eq!(s2.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap()["source"], "tracker"); + s2.execute_checked("roto.stroke", json!({"layer": clip.0, "kind": "bg", "points": line([10.0, 10.0], [60.0, 12.0]), "radius": 4})).unwrap(); + let st = s2.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); + assert_eq!(st["source"], "classical", "{st}"); + assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 0, "{st}"); + let _ = std::fs::remove_dir_all(std::path::Path::new(&path).parent().unwrap()); +} + /// A model whose tracker always fails: Roto Brush must carry on with the classic engine and say so. #[test] fn a_failing_tracker_falls_back_to_the_classic_engine() { diff --git a/docs/agents.md b/docs/agents.md index 54ba21afa..ddd730557 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -511,6 +511,9 @@ Trained model (Roto Brush 2.0 / 3.0): - `roto.propagate` never falls back to the classic engine silently: a chosen model that is not installed, or failed to load, is an error, and without `"wait": true` so is a model still loading. `roto.status` reports `model` (`selected`, `installed`, `active`, `loading`, `error`) and a `notice` while the chosen model is not the one in use. + Tracked frames are saved in the project with the instance (`tracker.saved` names their model): they apply while the + strokes and span are unchanged, whichever model Settings choose (`notice` says when that differs), and a new + Propagate replaces them. - `roto.model.remove {id}` deletes the installed weights. - The `version` property (`effects/#1/version`: 0 = 1.0, 1 = 2.0, 2 = 3.0) picks between the classic engine (1.0) and the chosen model (2.0, 3.0). diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index f5946a98e..d5216e6d4 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -16,15 +16,18 @@ Apache-2.0). It is optional: nothing changes until you install the model and cho `classical`) and, for tracked frames, `tracker.frames[f]`: `confidence` (0–1, the model's object score times its predicted mask quality) and `occluded` (the object is not in the frame: an empty matte). `model` says what Settings chose and whether it is installed, loaded and loading; `notice` says when it is not the one in use for the layer. + `tracker.saved` names the model the saved frames came from. Propagate never falls back to the classic engine without saying so. With SAM 2 chosen but not installed, or failing to load, it refuses with the reason. While it is still loading, `wait: true` (agents) loads it first; without `wait` it refuses, and the load goes on in the background. Deselecting SAM 2 (choosing Classic) runs Propagate exactly as before, and Roto Brush 1.0 never uses a model. -Tracked frames live in the session that ran Propagate and are not saved in the project file. A separate render process -(`effectcraft-cli render-frame`) does not have them, so render in the same session: in the app, or in one `run` that -includes `roto.propagate` and then `comp.saveFrameAs`. +Tracked frames are saved in the project with the instance, as Freeze's are, so a saved project renders the same when it +reopens, in the app or with `effectcraft-cli`, without the model being loaded or chosen there. They apply while the +strokes and span they came from are unchanged. Changing a stroke or the span makes them stale until the next Propagate. +Choosing another model, or the classic engine, leaves them in place and `roto.status` says so; a classic Propagate +replaces them with classic mattes. Object Matte saves nothing. ## How it works `crates/sam2` is a Rust port (on [candle](https://github.com/huggingface/candle), native only) of the Hugging Face From 1fb51600d5c93a74c8742b18ff741731f51fbcc3 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:13:38 +0200 Subject: [PATCH 25/77] S2.65: prompt SAM 2 with a box from the classic matte on stroke frames On real footage the stroke frame came out partial (clicks make SAM 2 pick a part of the object). Measured on a multi-part figure: clicks only IoU 0.41, with a box around the figure IoU 1.00. Chain::track now adds the box of the classic engine's matte of the strokes to each prompt; a box stroke wins. Co-Authored-By: Claude Sonnet 5.5 --- crates/effects/src/roto.rs | 24 ++++++----- crates/sam2/tests/figure.rs | 77 ++++++++++++++++++++++++++++++++++ crates/track/src/roto.rs | 23 ++++++++++ crates/track/src/roto/tests.rs | 23 ++++++++++ 4 files changed, 136 insertions(+), 11 deletions(-) create mode 100644 crates/sam2/tests/figure.rs diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 9596b6d50..f4f5e6bd0 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -553,17 +553,19 @@ impl Chain { ) -> Result, String> { let Some(tracker) = self.model.as_ref().and_then(|m| m.tracker()) else { return Ok(None) }; let (Some(_), true) = (self.data.base, self.data.span[1] >= self.data.span[0]) else { return Err("no base frame".into()) }; - let prompts: Vec<(i64, effectcraft_segment::Prompt)> = self - .data - .stroke_frames() - .into_iter() - .filter(|f| self.data.in_span(*f)) - .filter_map(|f| { - let strokes = self.data.strokes_at(f); - let p = rb::tracker_prompt(&strokes, self.opts.scale); - p.points.iter().any(|(_, fg)| *fg).then_some((f, p)) - }) - .collect(); + // A prompt is a few clicks from the strokes and a box around what the classic engine makes of + // them on that frame (clicks alone pick a part of the object, not all of it). + let mut prompts: Vec<(i64, effectcraft_segment::Prompt)> = vec![]; + for f in self.data.stroke_frames().into_iter().filter(|f| self.data.in_span(*f)) { + let strokes = self.data.strokes_at(f); + let p = match frame(f) { + Some(img) => rb::tracker_prompt_boxed(&strokes, self.opts.scale, &rb::segment(&img, &strokes, None, &self.opts, self.refine_radius)), + None => rb::tracker_prompt(&strokes, self.opts.scale), + }; + if p.points.iter().any(|(_, fg)| *fg) { + prompts.push((f, p)); + } + } if prompts.is_empty() { return Err("no foreground strokes to follow".into()); } diff --git a/crates/sam2/tests/figure.rs b/crates/sam2/tests/figure.rs new file mode 100644 index 000000000..2ca4d0b13 --- /dev/null +++ b/crates/sam2/tests/figure.rs @@ -0,0 +1,77 @@ +use effectcraft_sam2::model::ImageModel; +use effectcraft_sam2::prompt::Prompt; +use effectcraft_sam2::weights::Weights; + +fn figure(w: usize, h: usize) -> (Vec<[f32; 3]>, Vec) { + let mut rgb = vec![[0f32; 3]; w * h]; + let mut gt = vec![false; w * h]; + let part = |x: f32, y: f32| -> Option<[f32; 3]> { + let (cx, _cy) = (128.0, 0.0); + let r = |x0: f32, y0: f32, x1: f32, y1: f32| x >= x0 && x < x1 && y >= y0 && y < y1; + if (x - cx).powi(2) + (y - 40.0).powi(2) < 18.0f32.powi(2) { + return Some([0.8, 0.6, 0.5]); + } // head + if r(104.0, 60.0, 152.0, 130.0) { + return Some([0.2, 0.3, 0.7 + 0.1 * (x / 7.0).sin()]); + } // torso + if r(76.0, 62.0, 104.0, 74.0) || r(152.0, 62.0, 180.0, 74.0) { + return Some([0.8, 0.6, 0.5]); + } // arms + if r(108.0, 130.0, 124.0, 210.0) || r(132.0, 130.0, 148.0, 210.0) { + return Some([0.15, 0.15, 0.2]); + } // legs + if r(102.0, 210.0, 126.0, 224.0) || r(130.0, 210.0, 154.0, 224.0) { + return Some([0.95, 0.95, 0.95]); + } // white shoes + None + }; + for y in 0..h { + for x in 0..w { + let bg = 0.45 + 0.1 * (x as f32 / 11.0).sin() * (y as f32 / 9.0).cos(); + let p = part(x as f32, y as f32); + rgb[y * w + x] = p.unwrap_or([bg, bg * 0.9, bg * 0.8]); + gt[y * w + x] = p.is_some(); + } + } + (rgb, gt) +} + +fn iou(p: &[f32], gt: &[bool]) -> f64 { + let (mut i, mut u) = (0usize, 0usize); + for (a, b) in p.iter().zip(gt) { + let a = *a > 0.0; + i += usize::from(a && *b); + u += usize::from(a || *b); + } + i as f64 / u.max(1) as f64 +} + +/// A multi-part figure (torso, head, legs, arms, white shoes in different colours): clicks alone make +/// SAM 2 pick the torso and head, a box around the figure makes it cover all of it. This is why Roto +/// Brush prompts the tracker with a box from the classic engine's matte (found on real footage, where +/// the stroke frame came out partial). +#[test] +fn a_box_makes_sam2_cover_every_part_of_the_object() { + let wf = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../sam2-tiny/model.safetensors"); + let wf = std::env::var_os("EFFECTCRAFT_SAM2_DIR").map_or(wf, |d| std::path::PathBuf::from(d).join("model.safetensors")); + if !wf.exists() { + eprintln!("skipped: no SAM 2 weights"); + return; + } + let model = ImageModel::load(&Weights::open(&wf).unwrap()).unwrap(); + let (w, h) = (256usize, 256usize); + let (rgb, gt) = figure(w, h); + let enc = model.encode(&effectcraft_sam2::mask::preprocess(&rgb, w, h).unwrap()).unwrap(); + let s = 4.0f32; + let clicks = vec![([128.0 * s, 75.0 * s], true), ([128.0 * s, 100.0 * s], true), ([128.0 * s, 125.0 * s], true), ([20.0 * s, 20.0 * s], false)]; + let mask = |p: &Prompt| -> (f64, f32) { + let out = model.segment(&enc, p, false).unwrap(); + let m = out.masks.narrow(1, 0, 1).unwrap().upsample_bilinear2d(h, w, false).unwrap().flatten_all().unwrap().to_vec1::().unwrap(); + (iou(&m, >), out.iou.flatten_all().unwrap().to_vec1::().unwrap()[0]) + }; + let (points_only, _) = mask(&Prompt { points: clicks.clone(), bbox: None }); + let (boxed, predicted) = mask(&Prompt { points: clicks, bbox: Some([70.0 * s, 18.0 * s, 186.0 * s, 228.0 * s]) }); + eprintln!("clicks only: IoU {points_only:.2}; with the box: IoU {boxed:.2} (predicted quality {predicted:.2})"); + assert!(points_only < 0.6, "the clicks-only failure this test documents is gone ({points_only}): revisit the box prompt"); + assert!(boxed > 0.95 && predicted > 0.7, "boxed IoU {boxed}, predicted {predicted}"); +} diff --git a/crates/track/src/roto.rs b/crates/track/src/roto.rs index ce52a79a4..000f7b55f 100644 --- a/crates/track/src/roto.rs +++ b/crates/track/src/roto.rs @@ -923,6 +923,29 @@ pub fn tracker_prompt(strokes: &[&Stroke], scale: f64) -> Prompt { Prompt { points, bbox, ..Default::default() } } +/// [`tracker_prompt`] plus a box around the classic engine's matte of the same strokes. Clicks alone +/// make SAM 2 pick the part they land on (a torso, not the legs and shoes: measured IoU 0.41 on a +/// multi-part figure, against 1.00 with the box); the box says how far the object reaches. A box +/// stroke the user drew wins over the matte's. +pub fn tracker_prompt_boxed(strokes: &[&Stroke], scale: f64, classic: &FrameSeg) -> Prompt { + let mut p = tracker_prompt(strokes, scale); + if p.bbox.is_none() { + let (w, h) = (classic.w, classic.h); + let (mut x0, mut y0, mut x1, mut y1) = (usize::MAX, usize::MAX, 0usize, 0usize); + for (i, v) in classic.matte.iter().enumerate() { + if *v != 0 && w > 0 { + let (x, y) = (i % w, i / w); + (x0, y0, x1, y1) = (x0.min(x), y0.min(y), x1.max(x), y1.max(y)); + } + } + if x0 <= x1 { + let (mx, my) = (((x1 - x0) as f32 * 0.05).max(4.0), ((y1 - y0) as f32 * 0.05).max(4.0)); + p.bbox = Some([(x0 as f32 - mx).max(0.0), (y0 as f32 - my).max(0.0), (x1 as f32 + 1.0 + mx).min(w as f32), (y1 as f32 + 1.0 + my).min(h as f32)]); + } + } + p +} + /// Most foreground / background points a tracker prompt carries. pub const TRACKER_FG_POINTS: usize = 3; pub const TRACKER_BG_POINTS: usize = 3; diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index c08306060..b7d8c7809 100644 --- a/crates/track/src/roto/tests.rs +++ b/crates/track/src/roto/tests.rs @@ -612,3 +612,26 @@ fn tracked_segmentation_handles_absent_objects_and_bad_maps() { let wrong = segment_tracked(&img, &refs, vec![1.0; 3], false, &opts, 0.0); assert!(iou(&wrong.matte, >) > 0.95); } + +#[test] +fn tracker_prompts_carry_the_box_of_the_classic_matte() { + let (img, gt) = frame_at(center(0)); + let st = base_strokes(0); + let refs: Vec<&Stroke> = st.iter().collect(); + let classic = segment(&img, &refs, None, &SegOpts::default(), 0.0); + let p = tracker_prompt_boxed(&refs, 1.0, &classic); + let b = p.bbox.expect("a box"); + let c = center(0); + // the box reaches the disk's edges, with a margin + assert!(b[0] <= (c[0] - R) as f32 && b[0] > (c[0] - R) as f32 - 12.0, "{b:?}"); + assert!(b[2] >= (c[0] + R) as f32 && b[2] < (c[0] + R) as f32 + 12.0, "{b:?}"); + assert!(b[1] <= (c[1] - R) as f32 && b[3] >= (c[1] + R) as f32, "{b:?}"); + assert!(iou(&classic.matte, >) > 0.95); + // no matte, no box; a box the user drew wins + let none = tracker_prompt_boxed(&refs, 1.0, &FrameSeg::empty(W as usize, H as usize)); + assert_eq!(none.bbox, None); + let bx = Stroke { kind: StrokeKind::Box, frame: 0, radius: 0.0, points: vec![[10.0, 10.0], [50.0, 60.0]] }; + let mut with_box: Vec<&Stroke> = refs.clone(); + with_box.push(&bx); + assert_eq!(tracker_prompt_boxed(&with_box, 1.0, &classic).bbox, Some([10.0, 10.0, 50.0, 60.0])); +} From 7eb91ca44ca0c0a088997f98ad033b977f3e9db4 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:19:35 +0200 Subject: [PATCH 26/77] S2.43: track several objects in one SAM 2 pass TrackRequest takes a list of objects (each with its prompts on its frames) and the tracker reports (object, frame, result); frame encodings are shared, objects may start on different frames. VideoTracker::propagate_objects runs a subset. Limits: 32 objects. Test: two discs, one prompted on frame 0 and one on frame 2. Co-Authored-By: Claude Sonnet 5.5 --- crates/effects/src/roto.rs | 4 +- crates/engine/src/tests/roto.rs | 6 +-- crates/sam2/src/session.rs | 15 ++++++- crates/sam2/src/track.rs | 71 +++++++++++++++++++------------- crates/sam2/tests/reference.rs | 73 +++++++++++++++++++++++++++++---- crates/segment/src/lib.rs | 22 +++++++--- crates/track/src/roto/bench.rs | 4 +- 7 files changed, 145 insertions(+), 50 deletions(-) diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index f4f5e6bd0..869e2c132 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -569,7 +569,7 @@ impl Chain { if prompts.is_empty() { return Err("no foreground strokes to follow".into()); } - let req = TrackRequest { first: self.data.span[0], count: (self.data.span[1] - self.data.span[0] + 1) as usize, prompts }; + let req = TrackRequest::single(self.data.span[0], (self.data.span[1] - self.data.span[0] + 1) as usize, prompts); let frame = std::cell::RefCell::new(frame); let mut summary = TrackSummary::default(); tracker.track( @@ -579,7 +579,7 @@ impl Chain { Some(Arc::new(ClipFrame { rgb: rb::rgb_of(&img), w: img.width as usize, h: img.height as usize })) }, cancel, - &mut |g, t| { + &mut |_, g, t| { let (Some(key), Some(img)) = (self.keys.get(&g).copied(), (frame.borrow_mut())(g)) else { return }; let strokes = self.data.strokes_at(g); let seg = rb::segment_tracked(&img, &strokes, t.prob, t.occluded, &self.opts, self.refine_radius); diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index b18cd386d..4ba557781 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -778,7 +778,7 @@ impl effectcraft_segment::ObjectTracker for Exact { req: &effectcraft_segment::TrackRequest, frame: &mut dyn FnMut(i64) -> Option>, _: &dyn Fn() -> bool, - out: &mut dyn FnMut(i64, effectcraft_segment::TrackedFrame), + out: &mut dyn FnMut(usize, i64, effectcraft_segment::TrackedFrame), ) -> effectcraft_segment::Result<()> { for g in req.first..req.first + req.count as i64 { let Some(img) = frame(g) else { return Err(format!("no frame {g}")) }; @@ -786,7 +786,7 @@ impl effectcraft_segment::ObjectTracker for Exact { let prob = (0..img.w * img.h) .map(|i| ((((i % img.w) as f64 + 0.5 - c[0]).hypot((i / img.w) as f64 + 0.5 - c[1])) <= R) as u8 as f32) .collect::>(); - out(g, effectcraft_segment::TrackedFrame { prob, w: img.w, h: img.h, confidence: 0.99, occluded: false }); + out(0, g, effectcraft_segment::TrackedFrame { prob, w: img.w, h: img.h, confidence: 0.99, occluded: false }); } Ok(()) } @@ -889,7 +889,7 @@ fn a_failing_tracker_falls_back_to_the_classic_engine() { _: &TrackRequest, _: &mut dyn FnMut(i64) -> Option>, _: &dyn Fn() -> bool, - _: &mut dyn FnMut(i64, TrackedFrame), + _: &mut dyn FnMut(usize, i64, TrackedFrame), ) -> effectcraft_segment::Result<()> { Err("the tracker broke".into()) } diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 3440ff379..4b20b5c64 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -334,11 +334,24 @@ impl<'a> VideoTracker<'a> { /// Track every object from `start` through `max_frames` frames (the rest of the clip when /// `None`) in `direction`, calling `out` for each object on each frame as it is finished. pub fn propagate(&mut self, start: usize, direction: Direction, max_frames: Option, cancel: &Cancel, out: &mut dyn FnMut(FrameMask)) -> Result<()> { + self.propagate_objects(None, start, direction, max_frames, cancel, out) + } + + /// [`VideoTracker::propagate`] for the listed objects only (`None`: every prompted object). + pub fn propagate_objects( + &mut self, + only: Option<&[ObjectId]>, + start: usize, + direction: Direction, + max_frames: Option, + cancel: &Cancel, + out: &mut dyn FnMut(FrameMask), + ) -> Result<()> { self.check_frame(start)?; let n = self.frames.num_frames(); let reverse = direction == Direction::Backward; let span = max_frames.unwrap_or(n).min(n); - let ids: Vec = self.objects.iter().filter(|(_, o)| !o.cond.is_empty()).map(|(id, _)| *id).collect(); + let ids: Vec = self.objects.iter().filter(|(id, o)| !o.cond.is_empty() && only.is_none_or(|l| l.contains(id))).map(|(id, _)| *id).collect(); if ids.is_empty() { return Err(Error::Input("no object has been prompted".into())); } diff --git a/crates/sam2/src/track.rs b/crates/sam2/src/track.rs index 9320fbbe4..a46f7f6c6 100644 --- a/crates/sam2/src/track.rs +++ b/crates/sam2/src/track.rs @@ -50,32 +50,41 @@ fn to_frame(m: &FrameMask, size: (usize, usize)) -> Result { Ok(TrackedFrame { prob, w, h, confidence: (presence * m.iou.clamp(0.0, 1.0)).clamp(0.0, 1.0), occluded: m.occluded }) } +/// Most objects one request may track. +const MAX_OBJECTS: usize = 32; + impl Sam2Mask { - fn run_track(&self, req: &TrackRequest, frames: Frames<'_>, cancel: &dyn Fn() -> bool, out: &mut dyn FnMut(i64, TrackedFrame)) -> Result<()> { + fn run_track(&self, req: &TrackRequest, frames: Frames<'_>, cancel: &dyn Fn() -> bool, out: &mut dyn FnMut(usize, i64, TrackedFrame)) -> Result<()> { if req.count == 0 || req.count > MAX_CLIP { return Err(Error::Input(format!("a tracking request needs 1 to {MAX_CLIP} frames, got {}", req.count))); } - if req.prompts.is_empty() { - return Err(Error::Input("tracking needs a prompt on at least one frame".into())); + if req.objects.is_empty() || req.objects.len() > MAX_OBJECTS || req.objects.iter().any(|o| o.is_empty()) { + return Err(Error::Input(format!("tracking needs 1 to {MAX_OBJECTS} objects, each with a prompt on at least one frame"))); } let sizes = Rc::new(RefCell::new(HashMap::new())); let source = Source { first: req.first, count: req.count, fetch: RefCell::new(frames), sizes: sizes.clone() }; // Prompt coordinates are in the frame's pixels and the model works at 1024 x 1024, so the // prompted frames' sizes are needed first. let mut prompt_frames = vec![]; - for (frame, p) in &req.prompts { - let idx = usize::try_from(frame - req.first) - .ok() - .filter(|i| *i < req.count) - .ok_or_else(|| Error::Input(format!("prompt frame {frame} is outside the clip")))?; - let f = (source.fetch.borrow_mut())(*frame).ok_or_else(|| Error::Input(format!("frame {frame} is not available")))?; - sizes.borrow_mut().insert(idx, (f.w, f.h)); - prompt_frames.push((idx, p)); + for (oi, prompts) in req.objects.iter().enumerate() { + for (frame, p) in prompts { + let idx = usize::try_from(frame - req.first) + .ok() + .filter(|i| *i < req.count) + .ok_or_else(|| Error::Input(format!("prompt frame {frame} is outside the clip")))?; + if !sizes.borrow().contains_key(&idx) { + let f = (source.fetch.borrow_mut())(*frame).ok_or_else(|| Error::Input(format!("frame {frame} is not available")))?; + sizes.borrow_mut().insert(idx, (f.w, f.h)); + } + prompt_frames.push((oi, idx, p)); + } } let mut tracker = VideoTracker::new(self.video_model().clone(), Box::new(source))?; - let object = ObjectId(1); - let mut prompted = vec![]; - for (idx, p) in prompt_frames { + let object_id = |oi: usize| ObjectId(oi as u32 + 1); + // each object's first prompted frame: its pass starts there, forwards and backwards + let mut starts: std::collections::BTreeMap> = Default::default(); + let mut first_of: HashMap = HashMap::new(); + for (oi, idx, p) in prompt_frames { if cancel() { return Err(Error::Cancelled); } @@ -85,34 +94,38 @@ impl Sam2Mask { points: p.points.iter().map(|(c, fg)| ([c[0] * sx, c[1] * sy], *fg)).collect(), bbox: p.bbox.map(|b| [b[0] * sx, b[1] * sy, b[2] * sx, b[3] * sy]), }; - tracker.prompt(idx, object, &prompt)?; - prompted.push(idx); + tracker.prompt(idx, object_id(oi), &prompt)?; + let e = first_of.entry(oi).or_insert(idx); + *e = (*e).min(idx); + } + for (oi, idx) in first_of { + starts.entry(idx).or_default().push(object_id(oi)); } - let first = prompted.iter().copied().min().unwrap_or(0); let flag = Cancel::new(); let mut failure: Option = None; let mut emit = |m: FrameMask| { if cancel() { flag.cancel(); } + let oi = (m.object.0 as usize).saturating_sub(1); match to_frame(&m, sizes.borrow().get(&m.frame).copied().unwrap_or((0, 0))) { - Ok(t) => out(req.first + m.frame as i64, t), + Ok(t) => out(oi, req.first + m.frame as i64, t), Err(e) => { failure.get_or_insert(e); flag.cancel(); } } }; - tracker.propagate(first, Direction::Forward, None, &flag, &mut emit)?; - if first > 0 { - // frame `first` was reported by the forward pass; the backward pass starts there again - let mut skip_first = true; - tracker.propagate(first, Direction::Backward, None, &flag, &mut |m| { - if std::mem::take(&mut skip_first) { - return; - } - emit(m); - })?; + for (start, ids) in starts { + tracker.propagate_objects(Some(&ids), start, Direction::Forward, None, &flag, &mut emit)?; + if start > 0 { + // the start frame was reported by the forward pass; the backward pass passes it again + tracker.propagate_objects(Some(&ids), start, Direction::Backward, None, &flag, &mut |m| { + if m.frame != start { + emit(m); + } + })?; + } } failure.map_or(Ok(()), Err) } @@ -124,7 +137,7 @@ impl ObjectTracker for Sam2Mask { req: &TrackRequest, frames: &mut dyn FnMut(i64) -> Option>, cancel: &dyn Fn() -> bool, - out: &mut dyn FnMut(i64, TrackedFrame), + out: &mut dyn FnMut(usize, i64, TrackedFrame), ) -> effectcraft_segment::Result<()> { self.run_track(req, frames, cancel, out).map_err(|e| e.to_string()) } diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index ce1fd14d8..da7dd4706 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -364,13 +364,14 @@ fn sam2_tracks_a_clip_through_the_object_tracker_interface() { let model = effectcraft_sam2::mask::Sam2Mask::new(std::sync::Arc::new(effectcraft_sam2::session::VideoModel::load(&w).unwrap())); assert!(model.tracker().is_some()); let c = centre(1); - let req = TrackRequest { first: 100, count: n, prompts: vec![(101, Prompt { points: vec![([c.0, c.1], true)], bbox: None, mask: None })] }; + let req = TrackRequest::single(100, n, vec![(101, Prompt { points: vec![([c.0, c.1], true)], bbox: None, mask: None })]); let mut seen = std::collections::BTreeMap::new(); let mut get = |f: i64| frames.get((f - 100) as usize).cloned(); model .tracker() .unwrap() - .track(&req, &mut get, &|| false, &mut |f, t| { + .track(&req, &mut get, &|| false, &mut |o, f, t| { + assert_eq!(o, 0); assert!(seen.insert(f, t).is_none(), "frame {f} reported twice"); }) .unwrap(); @@ -391,9 +392,67 @@ fn sam2_tracks_a_clip_through_the_object_tracker_interface() { assert!(j > 0.85 && !t.occluded && t.confidence > 0.3, "frame {f}: IoU {j}"); } // Cancelling before the work starts, and bad requests, are errors. - assert!(model.tracker().unwrap().track(&req, &mut get, &|| true, &mut |_, _| {}).is_err()); - let none = TrackRequest { first: 100, count: n, prompts: vec![] }; - assert!(model.tracker().unwrap().track(&none, &mut get, &|| false, &mut |_, _| {}).is_err()); - let outside = TrackRequest { first: 100, count: n, prompts: vec![(500, Prompt { points: vec![([1.0, 1.0], true)], bbox: None, mask: None })] }; - assert!(model.tracker().unwrap().track(&outside, &mut get, &|| false, &mut |_, _| {}).is_err()); + assert!(model.tracker().unwrap().track(&req, &mut get, &|| true, &mut |_, _, _| {}).is_err()); + let none = TrackRequest::single(100, n, vec![]); + assert!(model.tracker().unwrap().track(&none, &mut get, &|| false, &mut |_, _, _| {}).is_err()); + let outside = TrackRequest::single(100, n, vec![(500, Prompt { points: vec![([1.0, 1.0], true)], bbox: None, mask: None })]); + assert!(model.tracker().unwrap().track(&outside, &mut get, &|| false, &mut |_, _, _| {}).is_err()); +} + +/// Two objects in one pass: each follows its own disc, every frame is reported once per object, and +/// objects may start on different frames. +#[test] +fn sam2_tracks_two_objects_in_one_pass() { + use effectcraft_segment::{ClipFrame, MaskModel, Prompt, TrackRequest}; + let Some((w, _)) = load() else { return }; + let (fw, fh, n) = (160usize, 120usize, 4usize); + // A red disc moving right and a blue one moving left, never touching. + let a = |t: usize| (36.0 + 6.0 * t as f32, 40.0f32); + let b = |t: usize| (124.0 - 6.0 * t as f32, 84.0f32); + let hit = |c: (f32, f32), x: usize, y: usize| ((x as f32 - c.0).powi(2) + (y as f32 - c.1).powi(2)).sqrt() < 20.0; + let frames: Vec> = (0..n) + .map(|t| { + let mut rgb = vec![[0f32; 3]; fw * fh]; + for y in 0..fh { + for x in 0..fw { + let bg = 0.45 + 0.1 * (x as f32 / 9.0).sin() * (y as f32 / 7.0).cos(); + rgb[y * fw + x] = if hit(a(t), x, y) { + [0.9, 0.15, 0.1] + } else if hit(b(t), x, y) { + [0.1, 0.25, 0.9] + } else { + [bg, bg * 0.9, bg * 0.8] + }; + } + } + std::sync::Arc::new(ClipFrame { rgb, w: fw, h: fh }) + }) + .collect(); + let model = effectcraft_sam2::mask::Sam2Mask::new(std::sync::Arc::new(effectcraft_sam2::session::VideoModel::load(&w).unwrap())); + let click = |c: (f32, f32)| Prompt { points: vec![([c.0, c.1], true)], bbox: None, mask: None }; + // object 0 is prompted on frame 0, object 1 on frame 2 + let req = TrackRequest { first: 10, count: n, objects: vec![vec![(10, click(a(0)))], vec![(12, click(b(2)))]] }; + let mut get = |f: i64| frames.get((f - 10) as usize).cloned(); + let mut seen: std::collections::BTreeMap<(usize, i64), effectcraft_segment::TrackedFrame> = Default::default(); + model.tracker().unwrap().track(&req, &mut get, &|| false, &mut |o, f, t| assert!(seen.insert((o, f), t).is_none(), "object {o} frame {f} twice")).unwrap(); + assert_eq!(seen.len(), 2 * n, "{:?}", seen.keys().collect::>()); + for ((o, f), t) in &seen { + let k = (f - 10) as usize; + let c = if *o == 0 { a(k) } else { b(k) }; + let (mut i, mut u) = (0usize, 0usize); + for y in 0..fh { + for x in 0..fw { + let (p, q) = (t.prob[y * fw + x] > 0.5, hit(c, x, y)); + i += usize::from(p && q); + u += usize::from(p || q); + } + } + let j = i as f64 / u.max(1) as f64; + eprintln!("object {o} frame {f}: IoU {j:.3}, confidence {:.2}", t.confidence); + assert!(j > 0.85, "object {o} frame {f}: IoU {j}"); + } + let too_many = TrackRequest { first: 10, count: n, objects: vec![vec![(10, click(a(0)))]; 40] }; + assert!(model.tracker().unwrap().track(&too_many, &mut get, &|| false, &mut |_, _, _| {}).is_err()); + let empty = TrackRequest { first: 10, count: n, objects: vec![] }; + assert!(model.tracker().unwrap().track(&empty, &mut get, &|| false, &mut |_, _, _| {}).is_err()); } diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index 37b3a4ca0..0a9a0494a 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -59,13 +59,22 @@ pub struct ClipFrame { pub h: usize, } -/// What to track: the clip's frames `first..first + count` (the host's frame numbers) and the -/// prompts of the one object, each on its frame. +/// What to track: the clip's frames `first..first + count` (the host's frame numbers) and, for +/// each object, its prompts on their frames. Objects are tracked in one pass, so the cost of +/// encoding each frame is shared. #[derive(Clone, Debug)] pub struct TrackRequest { pub first: i64, pub count: usize, - pub prompts: Vec<(i64, Prompt)>, + /// One entry per object: its prompts, each on its frame. + pub objects: Vec>, +} + +impl TrackRequest { + /// A request for a single object. + pub fn single(first: i64, count: usize, prompts: Vec<(i64, Prompt)>) -> TrackRequest { + TrackRequest { first, count, objects: vec![prompts] } + } } /// A tracker's answer for one frame. @@ -84,14 +93,15 @@ pub struct TrackedFrame { /// A video tracker with memory (SAM 2): given prompts on some frames, it follows the object /// through the clip in both directions. Optional: a [`MaskModel`] may offer one. pub trait ObjectTracker: Send + Sync { - /// Track the object. `frames(i)` supplies frame `i` (`None`: unavailable); `cancel()` is polled - /// between frames; `out(i, result)` is called as soon as each frame is finished. + /// Track the objects. `frames(i)` supplies frame `i` (`None`: unavailable); `cancel()` is polled + /// between frames; `out(object, i, result)` is called as soon as each object's frame is finished + /// (`object` indexes `req.objects`). fn track( &self, req: &TrackRequest, frames: &mut dyn FnMut(i64) -> Option>, cancel: &dyn Fn() -> bool, - out: &mut dyn FnMut(i64, TrackedFrame), + out: &mut dyn FnMut(usize, i64, TrackedFrame), ) -> Result<()>; } diff --git a/crates/track/src/roto/bench.rs b/crates/track/src/roto/bench.rs index 810bf34e5..2ce0663df 100644 --- a/crates/track/src/roto/bench.rs +++ b/crates/track/src/roto/bench.rs @@ -220,11 +220,11 @@ pub fn run_tracker(sc: &Scenario, tracker: &str, model: &dyn MaskModel) -> Optio let opts = SegOpts::default(); let refs: Vec<&Stroke> = sc.strokes.iter().collect(); let clip: Vec> = sc.frames.iter().map(|i| Arc::new(ClipFrame { rgb: rgb_of(i), w: i.width as usize, h: i.height as usize })).collect(); - let req = TrackRequest { first: 0, count: sc.frames.len(), prompts: vec![(sc.base as i64, tracker_prompt(&refs, 1.0))] }; + let req = TrackRequest::single(0, sc.frames.len(), vec![(sc.base as i64, tracker_prompt(&refs, 1.0))]); let start = Instant::now(); let mut mattes: Vec> = vec![vec![0u8; (W * H) as usize]; sc.frames.len()]; let mut get = |f: i64| usize::try_from(f).ok().and_then(|i| clip.get(i).cloned()); - let r = t.track(&req, &mut get, &|| false, &mut |f, tf| { + let r = t.track(&req, &mut get, &|| false, &mut |_, f, tf| { let (Ok(i), true) = (usize::try_from(f), tf.prob.len() == (W * H) as usize) else { return }; let Some(img) = sc.frames.get(i) else { return }; let strokes: Vec<&Stroke> = if i == sc.base { refs.clone() } else { vec![] }; From 0e0db685bf81e8558e7f1f539e038adc650b9cf3 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:24:08 +0200 Subject: [PATCH 27/77] S2.45: a correction stroke re-tracks only the frames it changes Chain::track tracks the frames of the span without a tracked result, in a window that restarts at the nearest prompted frame; unchanged spans are not tracked again. Test with a recording tracker: first pass 0..20, no second pass, then a correction on frame 8 tracks 8..20 with that prompt only. Co-Authored-By: Claude Sonnet 5.5 --- crates/effects/src/roto.rs | 30 ++++++++++++---- crates/engine/src/tests/roto.rs | 61 +++++++++++++++++++++++++++++++++ docs/roto-brush-sam2.md | 9 +++++ 3 files changed, 94 insertions(+), 6 deletions(-) diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 869e2c132..c31c4836b 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -540,8 +540,9 @@ impl Chain { self.model.as_ref().is_some_and(|m| m.tracker().is_some()) } - /// Follow the object through the whole span with the model's tracker, from the strokes on - /// their frames, and cache every frame's segmentation (under the tracker's keys, with its + /// Follow the object through the span with the model's tracker, from the strokes on their + /// frames (only the frames without a tracked result yet: after a correction stroke, the frames it + /// changes, tracked on from the nearest prompted frame), and cache every frame's segmentation (under the tracker's keys, with its /// confidence). The strokes stay hard constraints and the edges are refined at full /// resolution; frames the tracker can't give are left to [`Chain::ensure`]'s classic engine. /// `progress` gets the number of frames finished. `Ok(None)`: no tracker, nothing done. @@ -553,9 +554,19 @@ impl Chain { ) -> Result, String> { let Some(tracker) = self.model.as_ref().and_then(|m| m.tracker()) else { return Ok(None) }; let (Some(_), true) = (self.data.base, self.data.span[1] >= self.data.span[0]) else { return Err("no base frame".into()) }; + // Frames of the span that already have a tracked result (this session's, or the saved + // tracking) are kept: only the rest is tracked. A correction stroke changes the keys of the + // frames from it on (both directions), so only those come back here. + let (lo_span, hi_span) = (self.data.span[0], self.data.span[1]); + let needed: Vec = (lo_span..=hi_span) + .filter(|f| self.keys.get(f).is_some_and(|k| self.saved_seg(*f, self.opts.scale).is_none() && cached(tracked_key(*k), self.opts.scale).is_none())) + .collect(); + let (Some(&need_lo), Some(&need_hi)) = (needed.first(), needed.last()) else { + return Ok(Some(TrackSummary::default())); + }; // A prompt is a few clicks from the strokes and a box around what the classic engine makes of // them on that frame (clicks alone pick a part of the object, not all of it). - let mut prompts: Vec<(i64, effectcraft_segment::Prompt)> = vec![]; + let mut all_prompts: Vec<(i64, effectcraft_segment::Prompt)> = vec![]; for f in self.data.stroke_frames().into_iter().filter(|f| self.data.in_span(*f)) { let strokes = self.data.strokes_at(f); let p = match frame(f) { @@ -563,13 +574,20 @@ impl Chain { None => rb::tracker_prompt(&strokes, self.opts.scale), }; if p.points.iter().any(|(_, fg)| *fg) { - prompts.push((f, p)); + all_prompts.push((f, p)); } } - if prompts.is_empty() { + if all_prompts.is_empty() { return Err("no foreground strokes to follow".into()); } - let req = TrackRequest::single(self.data.span[0], (self.data.span[1] - self.data.span[0] + 1) as usize, prompts); + // The window to track: from the nearest prompted frame at or before the first frame needed + // (a correction restarts the object there) to the last frame needed, reaching the first + // prompted frame when everything needed lies before it. + let first_prompt = all_prompts.first().map_or(need_lo, |p| p.0); + let lo = all_prompts.iter().map(|p| p.0).filter(|f| *f <= need_lo).max().unwrap_or(need_lo); + let hi = if need_hi < first_prompt { first_prompt } else { need_hi }; + let prompts: Vec<(i64, effectcraft_segment::Prompt)> = all_prompts.into_iter().filter(|p| p.0 >= lo && p.0 <= hi).collect(); + let req = TrackRequest::single(lo, (hi - lo + 1) as usize, prompts); let frame = std::cell::RefCell::new(frame); let mut summary = TrackSummary::default(); tracker.track( diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index 4ba557781..7d18e5c98 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -940,3 +940,64 @@ fn without_a_tracker_the_chain_is_exactly_the_classic_engine() { } forget(&s, clip, uid); } + +/// Records the windows it is asked to track, and follows the disc exactly like [`Exact`]. +#[derive(Default)] +struct Recorder { + requests: std::sync::Mutex)>>, +} + +impl effectcraft_segment::MaskModel for Recorder { + fn info(&self) -> &'static effectcraft_segment::ModelInfo { + &effectcraft_segment::SAM2_TINY + } + fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &effectcraft_segment::Prompt) -> effectcraft_segment::Result> { + Err("no weights".into()) + } + fn tracker(&self) -> Option<&dyn effectcraft_segment::ObjectTracker> { + Some(self) + } +} + +impl effectcraft_segment::ObjectTracker for Recorder { + fn track( + &self, + req: &effectcraft_segment::TrackRequest, + frame: &mut dyn FnMut(i64) -> Option>, + cancel: &dyn Fn() -> bool, + out: &mut dyn FnMut(usize, i64, effectcraft_segment::TrackedFrame), + ) -> effectcraft_segment::Result<()> { + let prompted = req.objects.iter().flatten().map(|(f, _)| *f).collect(); + self.requests.lock().unwrap().push((req.first, req.count, prompted)); + effectcraft_segment::ObjectTracker::track(&Exact, req, frame, cancel, out) + } +} + +/// A correction stroke re-tracks only what it changes: the frames from it on, restarted at the +/// corrected frame; the frames before it keep their tracked results. +#[test] +fn a_correction_re_tracks_only_the_frames_it_changes() { + let (mut s, clip, _) = painted(center, "disk-retrack.mov"); + let rec = std::sync::Arc::new(Recorder::default()); + effectcraft_effects::roto::set_model(Some(rec.clone())); + let run = |s: &mut Session| s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "forward", "wait": true})); + let first = run(&mut s); + assert_eq!(rec.requests.lock().unwrap().as_slice(), [(0, 21, vec![0])], "the first pass tracks the whole span"); + first.unwrap(); + // Nothing changed: nothing to track. + run(&mut s).unwrap(); + assert_eq!(rec.requests.lock().unwrap().len(), 1, "an unchanged span is not tracked again"); + // A correction on frame 8: the frames from 8 on are tracked again, restarting at frame 8. + let c = center(8); + s.execute_checked("roto.stroke", json!({"layer": clip.0, "frame": 8, "points": line([c[0] - 20.0, c[1]], [c[0] + 20.0, c[1]]), "radius": 6})).unwrap(); + let again = run(&mut s); + let reqs = rec.requests.lock().unwrap().clone(); + effectcraft_effects::roto::set_model(None); + again.unwrap(); + assert_eq!(reqs.len(), 2); + assert_eq!(reqs[1], (8, 13, vec![8]), "the window runs from the corrected frame to the end of the span"); + // The frames before the correction are still the first pass's, the rest the second's. + let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 3})).unwrap(); + assert_eq!(st["source"], "tracker", "{st}"); + assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 21, "{st}"); +} diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index d5216e6d4..a4e58a36d 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -29,6 +29,15 @@ strokes and span they came from are unchanged. Changing a stroke or the span mak Choosing another model, or the classic engine, leaves them in place and `roto.status` says so; a classic Propagate replaces them with classic mattes. Object Matte saves nothing. +## Corrections re-track only what they change +A correction stroke changes the keys of the frames from it on (in both directions from the base), so Propagate tracks +only those again: the window starts at the nearest prompted frame at or before the first frame it needs (the corrected +frame, which restarts the object with its own prompt) and runs to the last frame it needs. Frames before the correction +keep their tracked results, and a Propagate with nothing changed does no tracking. The older non-conditioning memory +frames before the corrected frame are not rebuilt, so frames just after it can differ slightly from a from-scratch pass. +Tracked frames of a session are in memory; a project reopened in a new process re-tracks the whole span after a stroke +change (its saved tracking is stale as a whole). + ## How it works `crates/sam2` is a Rust port (on [candle](https://github.com/huggingface/candle), native only) of the Hugging Face `transformers` SAM 2 implementation, Apache-2.0 (see its `NOTICE`). Every stage is checked against PyTorch by From 1ae5bb9839d08a98f32ac8a71d276551f773cc4f Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:29:59 +0200 Subject: [PATCH 28/77] S2.44: confidence on the Roto Brush span bar Frames a tracker followed are coloured by its confidence (green, amber, red; slate where the object was absent) and drawn taller; roto_trust reports them. Not checked visually: no headless way to track without weights in a snapshot. Co-Authored-By: Claude Sonnet 5.5 --- crates/engine/src/analysis/roto.rs | 19 ++++++++++++ crates/engine/src/tests/roto.rs | 8 ++++- crates/ui-egui/src/panels/roto_tool.rs | 43 ++++++++++++++++++++++++-- 3 files changed, 67 insertions(+), 3 deletions(-) diff --git a/crates/engine/src/analysis/roto.rs b/crates/engine/src/analysis/roto.rs index 286b95095..9cc950e07 100644 --- a/crates/engine/src/analysis/roto.rs +++ b/crates/engine/src/analysis/roto.rs @@ -838,6 +838,25 @@ impl Session { changed } + /// How far to trust the frames a tracker followed (SAM 2), by layer frame: its confidence and + /// whether the object was absent. Frames the classic engine made have no entry. + pub fn roto_trust(&self, comp: ItemId, layer: LayerId, effect: Uid) -> BTreeMap { + let mut out = BTreeMap::new(); + let Some(c) = self.project.comp(comp) else { return out }; + let Some(l) = c.layer(layer) else { return out }; + let Some(g) = l.props.find_group(effect) else { return out }; + if is_object_matte(g) { + return out; + } + let ch = fx::Chain::new(¶ms_static(g), layer_size(&self.project, c, l), c.frame_rate.as_f64(), 1.0); + for f in ch.keys.keys() { + if let Some(t) = ch.trust_at(*f) { + out.insert(*f, t); + } + } + out + } + /// The computed frames of an instance (full resolution), by layer frame. pub fn roto_computed(&self, comp: ItemId, layer: LayerId, effect: Uid) -> BTreeMap { let mut out = BTreeMap::new(); diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index 7d18e5c98..66baee148 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -977,7 +977,9 @@ impl effectcraft_segment::ObjectTracker for Recorder { /// corrected frame; the frames before it keep their tracked results. #[test] fn a_correction_re_tracks_only_the_frames_it_changes() { - let (mut s, clip, _) = painted(center, "disk-retrack.mov"); + let (mut s, clip, uid) = painted(center, "disk-retrack.mov"); + let cid = s.active_comp_id().unwrap(); + assert!(s.roto_trust(cid, clip, uid).is_empty(), "no tracker data before tracking"); let rec = std::sync::Arc::new(Recorder::default()); effectcraft_effects::roto::set_model(Some(rec.clone())); let run = |s: &mut Session| s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "forward", "wait": true})); @@ -1000,4 +1002,8 @@ fn a_correction_re_tracks_only_the_frames_it_changes() { let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 3})).unwrap(); assert_eq!(st["source"], "tracker", "{st}"); assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 21, "{st}"); + // The span bar's data: the tracker's confidence for every tracked frame. + let trust = s.roto_trust(cid, clip, uid); + assert_eq!(trust.len(), 21); + assert!(trust.values().all(|t| t.confidence > 0.9 && !t.occluded)); } diff --git a/crates/ui-egui/src/panels/roto_tool.rs b/crates/ui-egui/src/panels/roto_tool.rs index ca3c56ac5..f8cd383c6 100644 --- a/crates/ui-egui/src/panels/roto_tool.rs +++ b/crates/ui-egui/src/panels/roto_tool.rs @@ -243,6 +243,20 @@ fn run_stroke(app: &mut EffectcraftApp, layer: u64, kind: &str, points: &[[f64; } } +/// The span bar's colour for a frame a tracker followed: green when it is sure, amber when it is +/// doubtful, red when it is not, slate when the object was absent. +pub(crate) fn trust_color(confidence: f32, occluded: bool) -> Color32 { + if occluded { + Color32::from_rgb(110, 118, 150) + } else if confidence >= 0.85 { + Color32::from_rgb(70, 200, 90) + } else if confidence >= 0.5 { + Color32::from_rgb(230, 180, 50) + } else { + Color32::from_rgb(220, 70, 60) + } +} + /// The segmentation span bar (above the panel's bottom bar): the layer's frames, the span, the /// base frame (yellow), computed frames (green) and stroke frames. Click to go to a frame; drag /// the span's ends to change it. @@ -264,10 +278,17 @@ pub fn span_bar(app: &mut EffectcraftApp, ui: &mut egui::Ui, p: &egui::Painter, let sp = Rect::from_min_max(pos2(x_of(d.span[0] as f64), track.min.y), pos2(x_of(d.span[1] as f64 + 1.0), track.max.y)); let frozen = fx::is_frozen(¶ms); p.rect_filled(sp, 2.0, if frozen { Color32::from_rgb(60, 110, 190) } else { Color32::from_rgb(90, 90, 96) }); + // Computed frames: green, or (frames a tracker followed) coloured by its confidence, with + // frames where the object was absent in slate. + let trust = app.session.roto_trust(cid, layer.id, uid); for (f, done) in app.session.roto_computed(cid, layer.id, uid) { if done { - let cr = Rect::from_min_max(pos2(x_of(f as f64), track.max.y - 4.0), pos2(x_of(f as f64 + 1.0), track.max.y)); - p.rect_filled(cr, 0.0, Color32::from_rgb(70, 200, 90)); + let (color, h) = match trust.get(&f) { + Some(t) => (trust_color(t.confidence, t.occluded), 6.0), + None => (Color32::from_rgb(70, 200, 90), 4.0), + }; + let cr = Rect::from_min_max(pos2(x_of(f as f64), track.max.y - h), pos2(x_of(f as f64 + 1.0), track.max.y)); + p.rect_filled(cr, 0.0, color); } } for f in d.stroke_frames() { @@ -363,3 +384,21 @@ pub fn bar_buttons(app: &mut EffectcraftApp, ui: &mut egui::Ui, p: &egui::Painte } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[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)); + assert_eq!(g, Color32::from_rgb(70, 200, 90)); + assert_eq!(a, Color32::from_rgb(230, 180, 50)); + assert_eq!(r, Color32::from_rgb(220, 70, 60)); + assert_eq!(s, Color32::from_rgb(110, 118, 150)); + // the boundaries + assert_eq!(trust_color(0.85, false), g); + assert_eq!(trust_color(0.5, false), a); + assert_eq!(trust_color(f32::NAN, false), r); + } +} From 81faf5b4b50dcad69b14ee327334f313856eb27d Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:34:58 +0200 Subject: [PATCH 29/77] S2.66: mask prompts match Sam2VideoModel Antialiased bilinear resize as PyTorch does it, the reference's mask-as-output path (mask is the frame's output, the decoder makes the object pointer), VideoTracker::prompt_mask, and mask-only prompts through the tracker interface. Prompt frame logits exact, tracking forwards and backwards within 0.0019. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/src/lib.rs | 1 + crates/sam2/src/resize.rs | 113 +++++++++++++++++++++++++++++++++ crates/sam2/src/session.rs | 55 ++++++++++++++++ crates/sam2/src/track.rs | 12 ++++ crates/sam2/tests/reference.rs | 88 +++++++++++++++++++++++++ tools/sam2_video_reference.py | 24 +++++++ 6 files changed, 293 insertions(+) create mode 100644 crates/sam2/src/resize.rs diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 6a0eb7fbf..c08ac33b5 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -20,6 +20,7 @@ pub mod model; pub mod neck; pub mod nn; pub mod prompt; +pub mod resize; pub mod session; pub mod track; pub mod weights; diff --git a/crates/sam2/src/resize.rs b/crates/sam2/src/resize.rs new file mode 100644 index 000000000..67e96e9af --- /dev/null +++ b/crates/sam2/src/resize.rs @@ -0,0 +1,113 @@ +// Modified work: the antialiased bilinear resize of PyTorch `F.interpolate(mode="bilinear", +// antialias=True, align_corners=False)`, as used by the Hugging Face `transformers` SAM 2 mask path +// (Apache-2.0 for the model code; the algorithm follows PyTorch's documented behaviour). See NOTICE. +//! Resizing masks to and from the model's resolution. + +use candle_core::{Device, Tensor}; + +use crate::prompt::INPUT_SIZE; +use crate::{Error, Result}; + +/// For each output index: the first input index and the normalised triangle-filter weights. +fn axis_weights(n_in: usize, n_out: usize) -> Vec<(usize, Vec)> { + let scale = n_in as f64 / n_out.max(1) as f64; + let support = if scale >= 1.0 { scale } else { 1.0 }; + let invscale = if scale >= 1.0 { 1.0 / scale } else { 1.0 }; + (0..n_out) + .map(|i| { + let center = scale * (i as f64 + 0.5); + let xmin = ((center - support + 0.5) as i64).max(0) as usize; + let xmax = (((center + support + 0.5) as i64).max(0) as usize).min(n_in); + let mut w: Vec = (xmin..xmax.max(xmin)) + .map(|x| { + let d = ((x as f64 - center + 0.5) * invscale).abs(); + if d < 1.0 { (1.0 - d) as f32 } else { 0.0 } + }) + .collect(); + let total: f32 = w.iter().sum(); + if total > 0.0 { + w.iter_mut().for_each(|v| *v /= total); + } + (xmin, w) + }) + .collect() +} + +/// Antialiased bilinear resize of a row-major `h x w` image to `oh x ow`. +pub fn antialias_bilinear(src: &[f32], (h, w): (usize, usize), (oh, ow): (usize, usize)) -> Result> { + if h == 0 || w == 0 || oh == 0 || ow == 0 || src.len() != h * w { + return Err(Error::Input(format!("cannot resize {} values as {w}x{h} to {ow}x{oh}", src.len()))); + } + let xs = axis_weights(w, ow); + let ys = axis_weights(h, oh); + // horizontal pass: h x ow + let mut mid = vec![0f32; h * ow]; + for y in 0..h { + let row = src.get(y * w..(y + 1) * w).unwrap_or(&[]); + for (x, (x0, wt)) in xs.iter().enumerate() { + let acc: f32 = wt.iter().enumerate().map(|(k, a)| a * row.get(x0 + k).copied().unwrap_or(0.0)).sum(); + if let Some(o) = mid.get_mut(y * ow + x) { + *o = acc; + } + } + } + // vertical pass: oh x ow + let mut out = vec![0f32; oh * ow]; + for (y, (y0, wt)) in ys.iter().enumerate() { + for x in 0..ow { + let acc: f32 = wt.iter().enumerate().map(|(k, a)| a * mid.get((y0 + k) * ow + x).copied().unwrap_or(0.0)).sum(); + if let Some(o) = out.get_mut(y * ow + x) { + *o = acc; + } + } + } + Ok(out) +} + +/// A binary mask of any size as the model's mask prompt: `[1, 1, 1024, 1024]` of 0 and 1 (resized +/// like the reference: antialiased bilinear, then thresholded at one half). +pub fn mask_to_model(mask: &[u8], w: usize, h: usize) -> Result { + let side = INPUT_SIZE; + let f: Vec = mask.iter().map(|v| f32::from(*v != 0)).collect(); + let data = if (w, h) == (side, side) { f } else { antialias_bilinear(&f, (h, w), (side, side))?.into_iter().map(|v| f32::from(v >= 0.5)).collect() }; + Ok(Tensor::from_vec(data, (1, 1, side, side), &Device::Cpu)?) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_constant_stays_constant_and_weights_sum_to_one() { + let out = antialias_bilinear(&[3.0; 64 * 64], (64, 64), (16, 16)).unwrap(); + assert!(out.iter().all(|v| (v - 3.0).abs() < 1e-5)); + let up = antialias_bilinear(&[2.0; 16], (4, 4), (9, 7)).unwrap(); + assert!(up.iter().all(|v| (v - 2.0).abs() < 1e-5)); + } + + #[test] + fn a_half_plane_blurs_over_the_filter_width() { + // 8 x 1: left half 0, right half 1, down to 2 x 1. The triangle filter is four pixels wide on + // each side, so the first output sees the ones at weights 0.375 and 0.125 of a total 3.5. + let src = [0., 0., 0., 0., 1., 1., 1., 1.]; + let out = antialias_bilinear(&src, (1, 8), (1, 2)).unwrap(); + assert!((out[0] - 0.5 / 3.5).abs() < 1e-6 && (out[1] - 3.0 / 3.5).abs() < 1e-6, "{out:?}"); + } + + #[test] + fn hostile_sizes_are_errors() { + assert!(antialias_bilinear(&[1.0; 4], (2, 2), (0, 3)).is_err()); + assert!(antialias_bilinear(&[1.0; 5], (2, 2), (1, 1)).is_err()); + assert!(mask_to_model(&[1; 3], 2, 2).is_err()); + assert!(mask_to_model(&[], 0, 0).is_err()); + } + + #[test] + fn a_small_mask_becomes_a_model_sized_binary_mask() { + let m = mask_to_model(&[0, 1, 1, 0], 2, 2).unwrap(); + assert_eq!(m.dims(), [1, 1, INPUT_SIZE, INPUT_SIZE]); + let v = m.flatten_all().unwrap().to_vec1::().unwrap(); + assert!(v.iter().all(|x| *x == 0.0 || *x == 1.0)); + assert!(v.iter().sum::() > 0.0); + } +} diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 4b20b5c64..c629a7b20 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -19,6 +19,7 @@ use crate::memory::{Memory, MemoryEncoder}; use crate::model::ImageModel; use crate::nn::{apply, linear}; use crate::prompt::{INPUT_SIZE, Prompt}; +use crate::resize::antialias_bilinear; use crate::weights::Weights; use crate::{Error, Result}; @@ -87,6 +88,8 @@ pub struct VideoModel { no_object_pointer: Tensor, pointer_proj: FeedForward, pointer_time_proj: Linear, + /// `mask_downsample`: a 4 x 4 stride-4 convolution taking a mask prompt to the prompt encoder's resolution. + mask_down: (Tensor, Tensor), } impl VideoModel { @@ -99,6 +102,7 @@ impl VideoModel { no_object_pointer: w.get("no_object_pointer", &[1, 256])?, pointer_proj: FeedForward::load(w, "object_pointer_proj", 256, 256, 256, 3, false)?, pointer_time_proj: linear(w, "temporal_positional_encoding_projection_layer", 256, 64)?, + mask_down: (w.get("mask_downsample.weight", &[1, 1, 4, 4])?, w.get("mask_downsample.bias", &[1])?), }) } @@ -331,6 +335,57 @@ impl<'a> VideoTracker<'a> { self.step(object, frame, Some(prompt), false, true) } + /// Prompt an object on a frame with a mask: `[1, 1, 1024, 1024]` of 0 and 1 (see + /// [`crate::resize::mask_to_model`]). The mask becomes the frame's output as it is (no decode of + /// the mask itself); the decoder only makes the object pointer. Like a click prompt, the frame is a + /// conditioning frame unless it has been tracked already. + pub fn prompt_mask(&mut self, frame: usize, object: ObjectId, mask: &Tensor) -> Result { + self.check_frame(frame)?; + if mask.dims() != [1, 1, INPUT_SIZE, INPUT_SIZE] { + return Err(Error::Input(format!("a mask prompt must be [1, 1, {INPUT_SIZE}, {INPUT_SIZE}], got {:?}", mask.dims()))); + } + let values = mask.flatten_all()?.to_vec1::()?; + if values.iter().any(|v| !v.is_finite()) { + return Err(Error::Input("a mask prompt must be finite".into())); + } + let tracked = self.objects.get(&object).is_some_and(|o| o.non_cond.contains_key(&frame)); + self.objects.entry(object).or_default(); + self.step_mask(object, frame, mask, &values, !tracked) + } + + /// The reference's `_use_mask_as_output`: the mask is the output, the decoder makes the pointer. + fn step_mask(&mut self, id: ObjectId, frame: usize, mask: &Tensor, values: &[f32], conditioning: bool) -> Result { + let feats = self.features(frame)?; + let model = self.model.clone(); + let present = values.iter().any(|v| *v > 0.0); + let high = ((mask * 20.0)? - 10.0)?; + let low = antialias_bilinear(&high.flatten_all()?.to_vec1::()?, (INPUT_SIZE, INPUT_SIZE), (4 * GRID, 4 * GRID))?; + // The pointer: the decoder with the mask as its dense prompt and no clicks, on the frame's features + // without the no-memory embedding. + let down = mask.conv2d(&model.mask_down.0, 0, 4, 1, 1)?.broadcast_add(&model.mask_down.1.reshape((1, 1, 1, 1))?)?; + let dense = model.image.prompt.dense_mask(&down)?; + let out = model.image.decode_embedded(&feats.low, &feats.high, &model.image.prompt.sparse_none()?, &dense, true)?; + 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 token = out.sam_tokens.narrow(1, best, 1)?.squeeze(1)?; + let decoder_present = scalar(&out.object_score)? > 0.0; + let pointer = if decoder_present { model.pointer_proj.forward(&token)? } else { model.no_object_pointer.clone() }; + let pointer = if present { pointer } else { model.no_object_pointer.clone() }; + let memory = model.memory_encoder.memory(&feats.low, &high, present)?; + let score = if present { 10.0 } else { -10.0 }; + let result = FrameMask { frame, object: id, logits: low, iou: 1.0, object_score: score, occluded: !present }; + let state = self.objects.entry(id).or_default(); + let entry = FrameOut { memory, pointer, result: conditioning.then(|| result.clone()) }; + if conditioning { + state.non_cond.remove(&frame); + state.cond.insert(frame, entry); + } else { + state.non_cond.insert(frame, entry); + Self::prune(state, frame); + } + Ok(result) + } + /// Track every object from `start` through `max_frames` frames (the rest of the clip when /// `None`) in `direction`, calling `out` for each object on each frame as it is finished. pub fn propagate(&mut self, start: usize, direction: Direction, max_frames: Option, cancel: &Cancel, out: &mut dyn FnMut(FrameMask)) -> Result<()> { diff --git a/crates/sam2/src/track.rs b/crates/sam2/src/track.rs index a46f7f6c6..2b06d08cd 100644 --- a/crates/sam2/src/track.rs +++ b/crates/sam2/src/track.rs @@ -13,6 +13,7 @@ use effectcraft_segment::{ClipFrame, ObjectTracker, TrackRequest, TrackedFrame}; use crate::mask::{Sam2Mask, preprocess}; use crate::prompt::{INPUT_SIZE, Prompt}; +use crate::resize::mask_to_model; use crate::session::{Cancel, Direction, FrameMask, FrameSource, ObjectId, VideoTracker}; use crate::{Error, Result}; @@ -89,6 +90,17 @@ impl Sam2Mask { return Err(Error::Cancelled); } let (w, h) = sizes.borrow().get(&idx).copied().ok_or_else(|| Error::Compute("frame size unknown".into()))?; + // A prompt that is only a mask (probability per pixel of the frame) prompts with the mask. + if let (true, true, Some(m)) = (p.points.is_empty(), p.bbox.is_none(), p.mask.as_ref()) { + if m.len() != w * h { + return Err(Error::Input(format!("a mask prompt of {} values does not fit a {w}x{h} frame", m.len()))); + } + let binary: Vec = m.iter().map(|v| u8::from(*v >= 0.5)).collect(); + tracker.prompt_mask(idx, object_id(oi), &mask_to_model(&binary, w, h)?)?; + let e = first_of.entry(oi).or_insert(idx); + *e = (*e).min(idx); + continue; + } let (sx, sy) = (INPUT_SIZE as f32 / w as f32, INPUT_SIZE as f32 / h as f32); let prompt = Prompt { points: p.points.iter().map(|(c, fg)| ([c[0] * sx, c[1] * sy], *fg)).collect(), diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index da7dd4706..4d4377374 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -456,3 +456,91 @@ fn sam2_tracks_two_objects_in_one_pass() { let empty = TrackRequest { first: 10, count: n, objects: vec![] }; assert!(model.tracker().unwrap().track(&empty, &mut get, &|| false, &mut |_, _, _| {}).is_err()); } + +/// A mask prompt instead of clicks: the mask is the prompt frame's output, and tracking forwards and +/// backwards from it matches `Sam2VideoModel` given the same mask. +#[test] +fn mask_prompted_tracking_matches_sam2_video_model() { + use effectcraft_sam2::session::{Cancel, Direction, ObjectId, VideoModel, VideoTracker}; + let Some((w, r)) = load_video() else { return }; + let model = std::sync::Arc::new(VideoModel::load(&w).unwrap()); + let mut tracker = VideoTracker::new(model, Box::new(Clip(r.get("clip.pixel_values").unwrap().clone()))).unwrap(); + let obj = ObjectId(1); + let first = tracker.prompt_mask(2, obj, r.get("maskprompt.input").unwrap()).unwrap(); + let check = |tag: &str, m: &effectcraft_sam2::session::FrameMask| { + let want = r.get(&format!("maskprompt.{tag}.{}.pred_masks", m.frame)).unwrap().flatten_all().unwrap().to_vec1::().unwrap(); + let want_score = r.get(&format!("maskprompt.{tag}.{}.score", m.frame)).unwrap().flatten_all().unwrap().to_vec1::().unwrap()[0]; + let err = m.logits.iter().zip(&want).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max); + let j = iou(&m.logits, &want); + eprintln!("mask prompt {tag} frame {}: logit error {err:.4}, IoU {j:.4}, score {:.3} (want {want_score:.3})", m.frame, m.object_score); + assert!(j >= VIDEO_MIN_IOU, "{tag} frame {} IoU {j}", m.frame); + assert!(err < VIDEO_LOGIT_TOL, "{tag} frame {} logit error {err}", m.frame); + assert!((m.object_score - want_score).abs() < 0.05, "{tag} frame {} score {} vs {want_score}", m.frame, m.object_score); + }; + check("prompt", &first); + let cancel = Cancel::new(); + let mut fwd = vec![]; + tracker.propagate(2, Direction::Forward, None, &cancel, &mut |m| fwd.push(m)).unwrap(); + fwd.iter().for_each(|m| check("fwd", m)); + let mut rev = vec![]; + tracker.propagate(2, Direction::Backward, None, &cancel, &mut |m| rev.push(m)).unwrap(); + rev.iter().for_each(|m| check("rev", m)); + assert_eq!((fwd.len(), rev.len()), (5, 3)); + // hostile masks are errors + let bad = Tensor::zeros((1, 1, 8, 8), candle_core::DType::F32, &Device::Cpu).unwrap(); + assert!(tracker.prompt_mask(2, obj, &bad).is_err()); + let nan = (Tensor::zeros((1, 1, 1024, 1024), candle_core::DType::F32, &Device::Cpu).unwrap() + f64::NAN).unwrap(); + assert!(tracker.prompt_mask(2, obj, &nan).is_err()); + assert!(tracker.prompt_mask(99, obj, r.get("maskprompt.input").unwrap()).is_err()); +} + +/// The tracker interface takes a mask as a prompt: the disc's true mask on frame 1 tracks the disc. +#[test] +fn a_mask_prompt_through_the_object_tracker_interface() { + use effectcraft_segment::{ClipFrame, MaskModel, Prompt, TrackRequest}; + let Some((w, _)) = load() else { return }; + let (fw, fh, n) = (128usize, 96usize, 4usize); + let centre = |t: usize| (40.0 + 8.0 * t as f32, 48.0f32); + let inside = |t: usize, x: usize, y: usize| ((x as f32 - centre(t).0).powi(2) + (y as f32 - centre(t).1).powi(2)).sqrt() < 26.0; + let frames: Vec> = (0..n) + .map(|t| { + let mut rgb = vec![[0f32; 3]; fw * fh]; + 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(t, x, y) { [0.85, 0.4, 0.15] } else { [bg, bg * 0.8 + 0.1, bg * 0.6 + 0.3] }; + } + } + std::sync::Arc::new(ClipFrame { rgb, w: fw, h: fh }) + }) + .collect(); + let model = effectcraft_sam2::mask::Sam2Mask::new(std::sync::Arc::new(effectcraft_sam2::session::VideoModel::load(&w).unwrap())); + let mask: Vec = (0..fw * fh).map(|i| f32::from(inside(1, i % fw, i / fw))).collect(); + let mut get = |f: i64| frames.get(f as usize).cloned(); + let req = TrackRequest::single(0, n, vec![(1, Prompt { points: vec![], bbox: None, mask: Some(mask) })]); + let mut seen = std::collections::BTreeMap::new(); + model + .tracker() + .unwrap() + .track(&req, &mut get, &|| false, &mut |_, f, t| { + seen.insert(f, t); + }) + .unwrap(); + assert_eq!(seen.len(), n); + for (f, t) in &seen { + let (mut i, mut u) = (0usize, 0usize); + for y in 0..fh { + for x in 0..fw { + let (a, b) = (t.prob[y * fw + x] > 0.5, inside(*f as usize, x, y)); + i += usize::from(a && b); + u += usize::from(a || b); + } + } + let j = i as f64 / u.max(1) as f64; + eprintln!("mask-prompted frame {f}: IoU {j:.3}"); + assert!(j > 0.85 && !t.occluded, "frame {f}: IoU {j}"); + } + // a mask that does not fit the frame is an error + let bad = TrackRequest::single(0, n, vec![(1, Prompt { points: vec![], bbox: None, mask: Some(vec![1.0; 5]) })]); + assert!(model.tracker().unwrap().track(&bad, &mut get, &|| false, &mut |_, _, _| {}).is_err()); +} diff --git a/tools/sam2_video_reference.py b/tools/sam2_video_reference.py index 6e769d955..a05223be9 100644 --- a/tools/sam2_video_reference.py +++ b/tools/sam2_video_reference.py @@ -96,6 +96,30 @@ def main(model_dir, out): for o in model.propagate_in_video_iterator(sess, start_frame_idx=PROMPT_FRAME, reverse=True): record("rev", o) + # The same object prompted with a mask instead of a click (binary, at the model's 1024 x 1024). + msess = Sam2VideoInferenceSession(video=pixels, video_height=SIZE, video_width=SIZE, dtype=torch.float32) + mobj = msess.obj_id_to_idx(1) + yy, xx = np.mgrid[0:1024, 0:1024].astype(np.float32) + cx, cy = (80 + 12 * PROMPT_FRAME) * scale, 128 * scale + disc = torch.from_numpy((np.hypot(xx + 0.5 - cx, yy + 0.5 - cy) < 34 * scale).astype(np.float32))[None, None] + taps["maskprompt.input"] = disc.clone() + msess.add_mask_inputs(mobj, PROMPT_FRAME, disc) + msess.obj_with_new_inputs = [1] + with torch.no_grad(): + o = model(msess, frame_idx=PROMPT_FRAME) + taps[f"maskprompt.prompt.{o.frame_idx}.pred_masks"] = o.pred_masks.detach().float().clone() + taps[f"maskprompt.prompt.{o.frame_idx}.score"] = o.object_score_logits.detach().float().clone() + for o in model.propagate_in_video_iterator(msess, start_frame_idx=PROMPT_FRAME): + taps[f"maskprompt.fwd.{o.frame_idx}.pred_masks"] = o.pred_masks.detach().float().clone() + taps[f"maskprompt.fwd.{o.frame_idx}.score"] = o.object_score_logits.detach().float().clone() + for o in model.propagate_in_video_iterator(msess, start_frame_idx=PROMPT_FRAME, reverse=True): + taps[f"maskprompt.rev.{o.frame_idx}.pred_masks"] = o.pred_masks.detach().float().clone() + taps[f"maskprompt.rev.{o.frame_idx}.score"] = o.object_score_logits.detach().float().clone() + md = msess.output_dict_per_obj[mobj] + for kind in ("cond_frame_outputs", "non_cond_frame_outputs"): + for f, oo in md[kind].items(): + taps[f"maskstate.{f}.object_pointer"] = oo["object_pointer"].detach().float().clone() + d = sess.output_dict_per_obj[obj_idx] for kind in ("cond_frame_outputs", "non_cond_frame_outputs"): for f, o in d[kind].items(): From 378dc97d7fda633f6885eb71f0fe93bdbbcebbc0 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:37:03 +0200 Subject: [PATCH 30/77] S2.32: tracker memory and time budgets; determinism test VideoTracker::set_budget (default 2 GiB of state, optional seconds per propagate call) stops with Error::Budget; state_bytes reports what it holds. Two runs of the same prompts give bit-identical masks. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/src/session.rs | 50 +++++++++++++++++++++++++++++++++- crates/sam2/tests/reference.rs | 48 ++++++++++++++++++++++++++++++++ 2 files changed, 97 insertions(+), 1 deletion(-) diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index c629a7b20..2e45cc407 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -133,11 +133,34 @@ struct FrameFeatures { high: (Tensor, Tensor), } +/// What a tracker may hold and take. A call that would exceed it stops with [`Error::Budget`]. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Budget { + /// Memory banks, results and cached frame features (bytes). + pub max_state_bytes: usize, + /// Wall-clock seconds one `propagate` call may take (`None`: no limit). + pub max_seconds: Option, +} + +impl Default for Budget { + fn default() -> Budget { + Budget { max_state_bytes: 2 << 30, max_seconds: None } + } +} + +/// Bytes of one stored frame: memory features and position encoding `[4096, 64]`, the pointer. +const FRAME_BYTES: usize = 2 * GRID * GRID * 64 * 4 + 256 * 4; +/// Bytes of a conditioning frame's kept result (`256 x 256` logits). +const RESULT_BYTES: usize = 4 * GRID * 4 * GRID * 4; +/// Bytes of one frame's cached neck features (lowest level and position, two skip levels). +const FEATURE_BYTES: usize = (2 * 256 * GRID * GRID + 32 * 16 * GRID * GRID + 64 * 4 * GRID * GRID) * 4; + pub struct VideoTracker<'a> { model: Arc, frames: Box, objects: BTreeMap, cache: VecDeque<(usize, Arc)>, + budget: Budget, } fn sine_embedding_1d(positions: &[f32], dim: usize) -> Result { @@ -161,13 +184,32 @@ impl<'a> VideoTracker<'a> { if n == 0 || n > MAX_FRAMES { return Err(Error::Input(format!("a clip needs between 1 and {MAX_FRAMES} frames, got {n}"))); } - Ok(VideoTracker { model, frames, objects: BTreeMap::new(), cache: VecDeque::new() }) + Ok(VideoTracker { model, frames, objects: BTreeMap::new(), cache: VecDeque::new(), budget: Budget::default() }) } pub fn num_frames(&self) -> usize { self.frames.num_frames() } + /// Replace the limits (see [`Budget`]). + pub fn set_budget(&mut self, budget: Budget) { + self.budget = budget; + } + + /// Bytes held now: stored frames, kept results and cached features. + pub fn state_bytes(&self) -> usize { + let frames: usize = self.objects.values().map(|o| (o.cond.len() + o.non_cond.len()) * FRAME_BYTES + o.cond.len() * RESULT_BYTES).sum(); + frames + self.cache.len() * FEATURE_BYTES + } + + fn check_budget(&self) -> Result<()> { + let used = self.state_bytes(); + if used > self.budget.max_state_bytes { + return Err(Error::Budget(format!("the tracker holds {} MB, over its {} MB limit", used >> 20, self.budget.max_state_bytes >> 20))); + } + Ok(()) + } + /// Forget every object and cached frame (the clip stays). pub fn reset(&mut self) { self.objects.clear(); @@ -298,6 +340,7 @@ impl<'a> VideoTracker<'a> { state.non_cond.insert(frame, entry); Self::prune(state, frame); } + self.check_budget()?; Ok(result) } @@ -383,6 +426,7 @@ impl<'a> VideoTracker<'a> { state.non_cond.insert(frame, entry); Self::prune(state, frame); } + self.check_budget()?; Ok(result) } @@ -410,6 +454,7 @@ impl<'a> VideoTracker<'a> { if ids.is_empty() { return Err(Error::Input("no object has been prompted".into())); } + let started = std::time::Instant::now(); for k in 0..=span { let frame = if reverse { start.checked_sub(k) } else { Some(start + k).filter(|f| *f < n) }; let Some(frame) = frame else { break }; @@ -417,6 +462,9 @@ impl<'a> VideoTracker<'a> { if cancel.is_cancelled() { return Err(Error::Cancelled); } + if let Some(limit) = self.budget.max_seconds.filter(|l| started.elapsed().as_secs_f64() > *l) { + return Err(Error::Budget(format!("tracking took more than {limit} s"))); + } let cached = self.objects.get(id).is_some_and(|o| o.cond.contains_key(&frame)); let stored = self.objects.get(id).and_then(|o| o.cond.get(&frame)).and_then(|o| o.result.clone()); let mask = match stored { diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index 4d4377374..230de7389 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -544,3 +544,51 @@ fn a_mask_prompt_through_the_object_tracker_interface() { let bad = TrackRequest::single(0, n, vec![(1, Prompt { points: vec![], bbox: None, mask: Some(vec![1.0; 5]) })]); assert!(model.tracker().unwrap().track(&bad, &mut get, &|| false, &mut |_, _, _| {}).is_err()); } + +/// The same prompts on the same frames give the same masks, bit for bit, run after run (CPU). +#[test] +fn tracking_is_deterministic() { + use effectcraft_sam2::prompt::Prompt; + use effectcraft_sam2::session::{Cancel, Direction, ObjectId, VideoModel, VideoTracker}; + let Some((w, r)) = load_video() else { return }; + let model = std::sync::Arc::new(VideoModel::load(&w).unwrap()); + let clip = r.get("clip.pixel_values").unwrap().narrow(0, 1, 3).unwrap(); + let run = || { + let mut t = VideoTracker::new(model.clone(), Box::new(Clip(clip.clone()))).unwrap(); + t.prompt(0, ObjectId(1), &Prompt { points: vec![([416.0, 512.0], true)], bbox: None }).unwrap(); + let mut out = vec![]; + t.propagate(0, Direction::Forward, None, &Cancel::new(), &mut |m| out.push(m)).unwrap(); + out + }; + let (a, b) = (run(), run()); + assert_eq!(a.len(), 3); + for (x, y) in a.iter().zip(&b) { + assert_eq!(x.logits, y.logits, "frame {} differs between runs", x.frame); + assert_eq!((x.iou.to_bits(), x.object_score.to_bits()), (y.iou.to_bits(), y.object_score.to_bits())); + } +} + +/// Hard limits: a tracker over its memory budget or past its time budget stops with Error::Budget. +#[test] +fn the_tracker_stops_at_its_budget() { + use effectcraft_sam2::Error; + use effectcraft_sam2::prompt::Prompt; + use effectcraft_sam2::session::{Budget, Cancel, Direction, ObjectId, VideoModel, VideoTracker}; + let Some((w, r)) = load_video() else { return }; + let model = std::sync::Arc::new(VideoModel::load(&w).unwrap()); + let clip = r.get("clip.pixel_values").unwrap().narrow(0, 1, 3).unwrap(); + let click = Prompt { points: vec![([416.0, 512.0], true)], bbox: None }; + // memory: one frame's worth is more than a 1 MB limit + let mut t = VideoTracker::new(model.clone(), Box::new(Clip(clip.clone()))).unwrap(); + t.set_budget(Budget { max_state_bytes: 1 << 20, max_seconds: None }); + assert!(matches!(t.prompt(0, ObjectId(1), &click), Err(Error::Budget(_)))); + // default limit: fine, and the bytes held are reported + let mut t = VideoTracker::new(model.clone(), Box::new(Clip(clip.clone()))).unwrap(); + t.prompt(0, ObjectId(1), &click).unwrap(); + let held = t.state_bytes(); + assert!(held > 2 << 20 && held < 200 << 20, "{held}"); + // time: a limit of zero stops the first propagation + t.set_budget(Budget { max_state_bytes: 2 << 30, max_seconds: Some(0.0) }); + let r = t.propagate(0, Direction::Forward, None, &Cancel::new(), &mut |_| {}); + assert!(matches!(r, Err(Error::Budget(_))), "{r:?}"); +} From 6ecd6793a41f92e27c2e9e551946313d8b707708 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:41:01 +0200 Subject: [PATCH 31/77] S2.32: unload an idle Roto Brush model after ten minutes MaskModel::idle_secs (SAM 2 tracks its use); the engine frees a model idle for ten minutes when no job runs, keeps the choice, and reloads it before a stroke or a Propagate needs it. Docs: memory and limits. Co-Authored-By: Claude Sonnet 5.5 --- crates/engine/src/analysis/models.rs | 32 +++++++++++++++++++++++- crates/engine/src/cmd/roto_cmds.rs | 2 ++ crates/engine/src/tests/models.rs | 37 ++++++++++++++++++++++++++++ crates/sam2/src/mask.rs | 13 +++++++++- crates/sam2/src/track.rs | 5 +++- crates/segment/src/lib.rs | 5 ++++ docs/roto-brush-sam2.md | 7 ++++++ 7 files changed, 98 insertions(+), 3 deletions(-) diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index bad9fd5ff..a584aebbb 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -63,7 +63,7 @@ impl Models { self.face_model.lock().ok().and_then(|m| m.clone()) } - fn set_active(&self, task: Task, model: Option) { + pub(crate) fn set_active(&self, task: Task, model: Option) { activate(self.lane(task), &self.face_model, task, model); } } @@ -128,6 +128,9 @@ fn kind(task: Task) -> &'static Kind { } } +/// How long a Roto Brush model may sit unused before its memory is freed. +const ROTO_MODEL_IDLE: Duration = Duration::from_secs(600); + /// How long a job that needs the Roto Brush model waits for it to load (or download) before it /// gives up. A worker that never reports back must not hang the command. const ROTO_MODEL_WAIT: Duration = Duration::from_secs(600); @@ -290,8 +293,35 @@ impl Session { Err(format!("{} is still loading: try again when it is ready, or pass wait: true", info.name)) } + /// Free a Roto Brush model that has been idle for [`ROTO_MODEL_IDLE`] (and is not under a job). The + /// choice stays: it loads again when a job or a stroke needs it. + fn unload_idle_roto_model(&mut self) { + if self.is_roto_running() || self.models.status(Task::Mask).busy.is_some() { + return; + } + let Some(m) = crate::effects::roto::installed_model() else { return }; + if m.idle_secs().is_some_and(|s| s >= ROTO_MODEL_IDLE.as_secs_f64()) { + let name = m.info().name; + drop(m); + self.models.set_active(Task::Mask, None); + self.events.push(crate::Event::Toast { + message: format!("{name} unloaded after {} minutes idle; it loads again when it is used", ROTO_MODEL_IDLE.as_secs() / 60), + error: false, + }); + } + } + + /// Load the chosen Roto Brush model again if it was unloaded for being idle (before a stroke uses it). + pub fn reload_roto_model(&mut self) { + let chosen = self.chosen_model(Task::Mask).to_string(); + if chosen != seg::CLASSICAL && self.models.status(Task::Mask).active.is_none() && self.models.status(Task::Mask).busy.is_none() { + self.apply_model(Task::Mask, true); + } + } + /// Finish what the model workers did: a toast, and loading a model that just arrived. pub fn poll_models(&mut self) { + self.unload_idle_roto_model(); for task in [Task::Mask, Task::Face] { let done = self.models.lane(task).lock().ok().and_then(|mut s| s.done.take()); if let Some((message, error)) = done { diff --git a/crates/engine/src/cmd/roto_cmds.rs b/crates/engine/src/cmd/roto_cmds.rs index 0931236a2..481a23b5d 100644 --- a/crates/engine/src/cmd/roto_cmds.rs +++ b/crates/engine/src/cmd/roto_cmds.rs @@ -196,6 +196,8 @@ fn points_p(p: &Value, cmd: &str) -> Result> { fn stroke(s: &mut Session, p: &Value) -> Result { let cmd = "roto.stroke"; + // A model unloaded for being idle comes back before the stroke is segmented with it. + s.reload_roto_model(); let (cid, lid) = match roto_p(s, p, cmd) { Ok((c, l, _)) => (c, l), Err(_) => layer_p(s, p, cmd)?, diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index 0b714d717..fab025a3a 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -105,3 +105,40 @@ fn models_are_listed_chosen_and_verified() { } let _ = std::fs::remove_dir_all(d); } + +/// A Roto Brush model that has been idle for ten minutes is unloaded (the choice stays), one in use is not, +/// and a model that does not track use is never unloaded. +#[test] +fn an_idle_roto_model_is_unloaded() { + use effectcraft_segment::{Loaded, MaskModel, ModelInfo, Prompt}; + struct Fake(Option); + impl MaskModel for Fake { + fn info(&self) -> &'static ModelInfo { + &effectcraft_segment::SAM2_TINY + } + fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &Prompt) -> effectcraft_segment::Result> { + Err("no weights".into()) + } + fn idle_secs(&self) -> Option { + self.0 + } + } + crate::tests_roto::hold_roto_cache_test_lock(); + let (mut s, d) = session("idle"); + for (idle, unloaded) in [(Some(5.0), false), (None, false), (Some(601.0), true)] { + let m: Arc = Arc::new(Fake(idle)); + s.models.set_active(Task::Mask, Some(Loaded::Mask(m))); + assert_eq!(s.models.status(Task::Mask).active.as_deref(), Some("sam2.1-hiera-tiny")); + s.poll_models(); + assert_eq!(s.models.status(Task::Mask).active.is_none(), unloaded, "idle {idle:?}"); + assert_eq!(crate::effects::roto::installed_model().is_none(), unloaded, "the effect's model follows"); + if unloaded { + assert!(s.events.iter().any(|e| matches!(e, crate::Event::Toast { message, .. } if message.contains("unloaded after 10 minutes")))); + } + } + // the choice survives, and with the classic engine chosen reloading does nothing + s.reload_roto_model(); + assert!(s.models.status(Task::Mask).active.is_none()); + crate::effects::roto::set_model(None); + let _ = std::fs::remove_dir_all(d); +} diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index a0eac9cd6..f917e29ef 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -54,11 +54,17 @@ fn fingerprint(rgb: &[[f32; 3]], w: usize, h: usize) -> u64 { pub struct Sam2Mask { model: Arc, last: Mutex)>>, + used: Mutex, } impl Sam2Mask { pub fn new(model: Arc) -> Sam2Mask { - Sam2Mask { model, last: Mutex::new(None) } + Sam2Mask { model, last: Mutex::new(None), used: Mutex::new(std::time::Instant::now()) } + } + + /// Note that the model is being used (for [`MaskModel::idle_secs`]). + pub(crate) fn touch(&self) { + *self.used.lock().unwrap_or_else(PoisonError::into_inner) = std::time::Instant::now(); } /// The video model, for trackers built on the same weights. @@ -103,12 +109,17 @@ impl MaskModel for Sam2Mask { } fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &effectcraft_segment::Prompt) -> effectcraft_segment::Result> { + self.touch(); self.probability(rgb, w, h, prompt).map_err(|e| e.to_string()) } fn tracker(&self) -> Option<&dyn effectcraft_segment::ObjectTracker> { Some(self) } + + fn idle_secs(&self) -> Option { + Some(self.used.lock().unwrap_or_else(PoisonError::into_inner).elapsed().as_secs_f64()) + } } /// Verify the installed file against the registry and build the model. diff --git a/crates/sam2/src/track.rs b/crates/sam2/src/track.rs index 2b06d08cd..2476ddb74 100644 --- a/crates/sam2/src/track.rs +++ b/crates/sam2/src/track.rs @@ -151,6 +151,9 @@ impl ObjectTracker for Sam2Mask { cancel: &dyn Fn() -> bool, out: &mut dyn FnMut(usize, i64, TrackedFrame), ) -> effectcraft_segment::Result<()> { - self.run_track(req, frames, cancel, out).map_err(|e| e.to_string()) + self.touch(); + let r = self.run_track(req, frames, cancel, out).map_err(|e| e.to_string()); + self.touch(); + r } } diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index 0a9a0494a..bd0b1e672 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -49,6 +49,11 @@ pub trait MaskModel: Send + Sync { fn tracker(&self) -> Option<&dyn ObjectTracker> { None } + /// Seconds since the model was last used (`None`: it does not track use, so it is never unloaded + /// for being idle). Large models free their memory after a quiet spell and load again on use. + fn idle_secs(&self) -> Option { + None + } } /// One decoded frame handed to a tracker (straight RGB 0–1, row-major `w×h`). diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index a4e58a36d..e61f49be7 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -38,6 +38,13 @@ frames before the corrected frame are not rebuilt, so frames just after it can d Tracked frames of a session are in memory; a project reopened in a new process re-tracks the whole span after a stroke change (its saved tracking is stale as a whole). +## Memory and limits +The weights take about 156 MB. A model idle for 10 minutes is unloaded (a toast says so; the choice stays) and loads +again when a Propagate or a stroke needs it. The tracker holds the memory banks of the objects it follows (about 2 MB per +stored frame, pruned to the frames later ones can still look at) and refuses to grow past 2 GiB, or past a time limit when +one is set, with a clear error instead of running out of memory. Two runs of the same prompts give bit-identical masks on +the CPU. + ## How it works `crates/sam2` is a Rust port (on [candle](https://github.com/huggingface/candle), native only) of the Hugging Face `transformers` SAM 2 implementation, Apache-2.0 (see its `NOTICE`). Every stage is checked against PyTorch by From aab00e43fecb5be09b68cb8c7cc007067a614d62 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:44:55 +0200 Subject: [PATCH 32/77] S2.60: parallel exact GELU: encode 3.4 s -> 2.6 s per frame Candle's gelu_erf runs on one thread and was a third of the Hiera block time. A rayon CustomOp1 calls the same erf per element (bit-identical, tested). All reference tests unchanged. op_bench example times the block's operations. Co-Authored-By: Claude Sonnet 5.5 --- Cargo.lock | 1 + crates/sam2/Cargo.toml | 1 + crates/sam2/examples/op_bench.rs | 52 +++++++++++++++++++++++++++++ crates/sam2/src/decoder.rs | 8 +++-- crates/sam2/src/memory.rs | 6 ++-- crates/sam2/src/nn.rs | 56 ++++++++++++++++++++++++++++++-- crates/sam2/src/prompt.rs | 6 ++-- 7 files changed, 119 insertions(+), 11 deletions(-) create mode 100644 crates/sam2/examples/op_bench.rs diff --git a/Cargo.lock b/Cargo.lock index 423abb85a..a9b24748d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1835,6 +1835,7 @@ dependencies = [ "candle-core", "candle-nn", "effectcraft-segment", + "rayon", ] [[package]] diff --git a/crates/sam2/Cargo.toml b/crates/sam2/Cargo.toml index 4b6dd7e9f..61afb4117 100644 --- a/crates/sam2/Cargo.toml +++ b/crates/sam2/Cargo.toml @@ -11,6 +11,7 @@ workspace = true [dependencies] effectcraft-segment = { workspace = true } +rayon = { 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/examples/op_bench.rs b/crates/sam2/examples/op_bench.rs new file mode 100644 index 000000000..754248835 --- /dev/null +++ b/crates/sam2/examples/op_bench.rs @@ -0,0 +1,52 @@ +//! Times the tensor operations a Hiera block is made of, at the sizes of its first (biggest) block. +//! `cargo run --release -p effectcraft-sam2 --example op_bench` + +use std::time::Instant; + +use candle_core::{DType, Device, Tensor}; +use candle_nn::{Linear, Module}; + +fn time(name: &str, mut f: impl FnMut() -> T) { + let _ = f(); + let n = 5; + let t = Instant::now(); + for _ in 0..n { + let _ = f(); + } + println!("{name:<44} {:>8.1} ms", t.elapsed().as_secs_f64() * 1e3 / n as f64); +} + +fn main() -> Result<(), Box> { + let d = Device::Cpu; + // block 0: 1 x 256 x 256 x 96 activations, windows of 8 x 8 + let x = Tensor::randn(0f32, 1.0, (1, 256, 256, 96), &d)?; + let w = Tensor::randn(0f32, 0.1, (288, 96), &d)?; + let b = Tensor::zeros(288, DType::F32, &d)?; + let lin = Linear::new(w, Some(b)); + let flat = x.reshape((65536, 96))?; + time("linear 65536x96 -> 288 (qkv)", || lin.forward(&flat).unwrap()); + let big = Tensor::randn(0f32, 1.0, (65536, 384), &d)?; + time("gelu_erf on 65536 x 384", || big.gelu_erf().unwrap()); + time("relu on 65536 x 384", || big.relu().unwrap()); + time("add 65536 x 384", || (&big + &big).unwrap()); + let xs = x.clone(); + time("permute+contiguous (window partition)", || x.reshape((1, 32, 8, 32, 8, 96)).unwrap().permute((0, 1, 3, 2, 4, 5)).unwrap().contiguous().unwrap()); + let lnw = Tensor::ones(96, DType::F32, &d)?; + let lnb = Tensor::zeros(96, DType::F32, &d)?; + time("fused layer_norm 256x256x96", || candle_nn::ops::layer_norm(&xs, &lnw, &lnb, 1e-6).unwrap()); + // windowed attention: 1024 windows x 64 tokens x 96 dims, one head + let q = Tensor::randn(0f32, 1.0, (1024, 1, 64, 96), &d)?; + let kt = Tensor::randn(0f32, 1.0, (1024, 1, 96, 64), &d)?; + time("batched matmul 1024 x (64x96)(96x64)", || q.matmul(&kt).unwrap()); + let sc = Tensor::randn(0f32, 1.0, (1024, 1, 64, 64), &d)?; + time("softmax_last_dim 1024x64x64", || candle_nn::ops::softmax_last_dim(&sc).unwrap()); + let v = Tensor::randn(0f32, 1.0, (1024, 1, 64, 96), &d)?; + time("batched matmul 1024 x (64x64)(64x96)", || sc.matmul(&v).unwrap()); + // the MLP matmuls + let w1 = Linear::new(Tensor::randn(0f32, 0.1, (384, 96), &d)?, Some(Tensor::zeros(384, DType::F32, &d)?)); + time("linear 65536x96 -> 384 (mlp in)", || w1.forward(&flat).unwrap()); + let w2 = Linear::new(Tensor::randn(0f32, 0.1, (96, 384), &d)?, Some(Tensor::zeros(96, DType::F32, &d)?)); + time("linear 65536x384 -> 96 (mlp out)", || w2.forward(&big).unwrap()); + println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get())); + Ok(()) +} diff --git a/crates/sam2/src/decoder.rs b/crates/sam2/src/decoder.rs index d7a7a6dcf..d6b25a7e4 100644 --- a/crates/sam2/src/decoder.rs +++ b/crates/sam2/src/decoder.rs @@ -8,7 +8,7 @@ use candle_core::{D, Tensor}; use candle_nn::{Linear, ops::sigmoid, ops::softmax}; -use crate::nn::{LayerNorm, apply, layer_norm, linear}; +use crate::nn::{LayerNorm, apply, gelu, layer_norm, linear}; use crate::weights::Weights; use crate::{Error, Result}; @@ -218,8 +218,10 @@ impl MaskDecoder { let (s0, s1) = high; let up = add_bias(emb.conv_transpose2d(&self.up1.0, 0, 0, 2, 1)?, &self.up1.1)?; let up = (up + s1)?; - let up = self.up_ln.forward(&up.permute((0, 2, 3, 1))?)?.permute((0, 3, 1, 2))?.contiguous()?.gelu_erf()?; - let up = (add_bias(up.conv_transpose2d(&self.up2.0, 0, 0, 2, 1)?, &self.up2.1)? + s0)?.gelu_erf()?; + let up = self.up_ln.forward(&up.permute((0, 2, 3, 1))?)?.permute((0, 3, 1, 2))?.contiguous()?; + let up = gelu(&up)?; + let up = (add_bias(up.conv_transpose2d(&self.up2.0, 0, 0, 2, 1)?, &self.up2.1)? + s0)?; + let up = gelu(&up)?; let hyper: Vec = self.hyper.iter().enumerate().map(|(i, m)| m.forward(&mask_tokens_out.narrow(1, i, 1)?.squeeze(1)?)).collect::>()?; let hyper = Tensor::stack(&hyper, 1)?; diff --git a/crates/sam2/src/memory.rs b/crates/sam2/src/memory.rs index 51e3d3c92..735c4a563 100644 --- a/crates/sam2/src/memory.rs +++ b/crates/sam2/src/memory.rs @@ -8,7 +8,7 @@ use candle_core::Tensor; use crate::Result; use crate::neck::sine_position; -use crate::nn::{LayerNorm, apply, layer_norm, linear}; +use crate::nn::{LayerNorm, apply, gelu, layer_norm, linear}; use crate::weights::Weights; use candle_nn::Linear; @@ -49,7 +49,7 @@ impl FuserBlock { fn forward(&self, x: &Tensor) -> Result { let h = norm_channels_first(&self.norm, &self.depthwise.forward(x)?)?; let h = h.permute((0, 2, 3, 1))?.contiguous()?; - let h = apply(&self.pw2, &apply(&self.pw1, &h)?.gelu_erf()?)?; + let h = apply(&self.pw2, &gelu(&apply(&self.pw1, &h)?)?)?; let h = h.broadcast_mul(&self.scale)?.permute((0, 3, 1, 2))?.contiguous()?; Ok((x + h)?) } @@ -108,7 +108,7 @@ impl MemoryEncoder { pub fn encode(&self, features: &Tensor, mask: &Tensor) -> Result<(Tensor, Tensor)> { let mut m = mask.clone(); for (conv, ln) in &self.down { - m = norm_channels_first(ln, &conv.forward(&m)?)?.gelu_erf()?; + m = gelu(&norm_channels_first(ln, &conv.forward(&m)?)?)?; } let m = self.final_conv.forward(&m)?; let mut x = (self.feature_projection.forward(features)? + m)?; diff --git a/crates/sam2/src/nn.rs b/crates/sam2/src/nn.rs index f90958eb1..76bafabac 100644 --- a/crates/sam2/src/nn.rs +++ b/crates/sam2/src/nn.rs @@ -1,7 +1,7 @@ // Modified work: small layers used by the ported SAM 2 modules (Apache-2.0, see NOTICE). //! Linear, layer norm and a few tensor helpers on candle. -use candle_core::Tensor; +use candle_core::{CpuStorage, CustomOp1, Layout, Shape, Tensor}; use candle_nn::{Linear, Module}; use crate::weights::Weights; @@ -41,7 +41,7 @@ impl Mlp { } pub fn forward(&self, x: &Tensor) -> Result { - apply(&self.proj_out, &apply(&self.proj_in, x)?.gelu_erf()?) + apply(&self.proj_out, &gelu(&apply(&self.proj_in, x)?)?) } } @@ -57,3 +57,55 @@ pub fn apply(l: &Linear, x: &Tensor) -> Result { out.push(y.dim(1)?); Ok(y.reshape(out)?) } + +/// Exact (erf) GELU on all cores. Candle's own `gelu_erf` runs on one thread, which made it a third of +/// the time of the Hiera blocks; this gives the same values (it calls the same `erf` per element). +struct GeluErf; + +impl CustomOp1 for GeluErf { + fn name(&self) -> &'static str { + "gelu_erf_parallel" + } + + fn cpu_fwd(&self, storage: &CpuStorage, layout: &Layout) -> candle_core::Result<(CpuStorage, Shape)> { + use rayon::prelude::*; + let CpuStorage::F32(data) = storage else { candle_core::bail!("gelu_erf_parallel needs f32") }; + let Some((start, end)) = layout.contiguous_offsets() else { candle_core::bail!("gelu_erf_parallel needs a contiguous input") }; + let src = data.get(start..end).ok_or_else(|| candle_core::Error::Msg("gelu_erf_parallel: layout outside the storage".into()))?; + let mut out = vec![0f32; src.len()]; + const CHUNK: usize = 16 * 1024; + out.par_chunks_mut(CHUNK).zip(src.par_chunks(CHUNK)).for_each(|(o, i)| { + for (a, v) in o.iter_mut().zip(i) { + *a = (candle_core::cpu::erf::erf_f32(v * std::f32::consts::FRAC_1_SQRT_2) + 1.0) * 0.5 * v; + } + }); + Ok((CpuStorage::F32(out), layout.shape().clone())) + } +} + +/// Exact GELU of `x`, in parallel for f32 tensors. +pub fn gelu(x: &Tensor) -> Result { + if x.dtype() != candle_core::DType::F32 || !x.device().is_cpu() { + return Ok(x.gelu_erf()?); + } + Ok(x.contiguous()?.apply_op1_no_bwd(&GeluErf)?) +} + +#[cfg(test)] +mod tests { + use super::*; + use candle_core::Device; + + #[test] + fn the_parallel_gelu_equals_candles_bit_for_bit() { + let v: Vec = (0..101_000).map(|i| ((i as f32) * 0.00037).sin() * 6.0 - 0.5).collect(); + let x = Tensor::from_vec(v, (101, 1000), &Device::Cpu).unwrap(); + let (a, b) = (gelu(&x).unwrap(), x.gelu_erf().unwrap()); + assert_eq!(a.flatten_all().unwrap().to_vec1::().unwrap(), b.flatten_all().unwrap().to_vec1::().unwrap()); + // a transposed (non-contiguous) input works, and an empty one does too + let t = x.t().unwrap(); + assert_eq!(gelu(&t).unwrap().dims(), t.dims()); + let e = Tensor::zeros((0, 4), candle_core::DType::F32, &Device::Cpu).unwrap(); + assert_eq!(gelu(&e).unwrap().elem_count(), 0); + } +} diff --git a/crates/sam2/src/prompt.rs b/crates/sam2/src/prompt.rs index 51aade1c7..91496bba3 100644 --- a/crates/sam2/src/prompt.rs +++ b/crates/sam2/src/prompt.rs @@ -5,7 +5,7 @@ use candle_core::{Device, Tensor}; -use crate::nn::{LayerNorm, layer_norm}; +use crate::nn::{LayerNorm, gelu, layer_norm}; use crate::weights::Weights; use crate::{Error, Result}; @@ -166,8 +166,8 @@ impl PromptEncoder { Ok(x.conv2d(&c.0, 0, c.3, 1, 1)?.broadcast_add(&c.1.reshape((1, c.1.dim(0)?, 1, 1))?)?) }; let cl = |x: &Tensor, ln: &LayerNorm| -> Result { Ok(ln.forward(&x.permute((0, 2, 3, 1))?)?.permute((0, 3, 1, 2))?.contiguous()?) }; - let x = cl(&step(mask, c1)?, &self.mask.ln1)?.gelu_erf()?; - let x = cl(&step(&x, c2)?, &self.mask.ln2)?.gelu_erf()?; + let x = gelu(&cl(&step(mask, c1)?, &self.mask.ln1)?)?; + let x = gelu(&cl(&step(&x, c2)?, &self.mask.ln2)?)?; step(&x, c3) } } From b7c78065dba4e0644281688a3163bb5ef8592f76 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:51:17 +0200 Subject: [PATCH 33/77] S2.61: cache the sine positions, pool without permutes Sine position encodings depend only on the sizes: built once (150 ms of a frame). do_pool takes the maximum over grouped cells instead of permute + max_pool2d + permute (tested equal, odd sizes included). Encode 2.6 s -> 2.4 s per frame in release on 16 threads (3.4 s at the start of the speed work); reference tests unchanged. Chunked global attention was tried and did not help. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/examples/op_bench.rs | 16 ++++++++++++++++ crates/sam2/examples/sam2_bench.rs | 10 ++++++++++ crates/sam2/src/hiera.rs | 23 ++++++++++++++++++++--- crates/sam2/src/neck.rs | 19 +++++++++++++++++++ 4 files changed, 65 insertions(+), 3 deletions(-) diff --git a/crates/sam2/examples/op_bench.rs b/crates/sam2/examples/op_bench.rs index 754248835..4e2b5b4a3 100644 --- a/crates/sam2/examples/op_bench.rs +++ b/crates/sam2/examples/op_bench.rs @@ -47,6 +47,22 @@ fn main() -> Result<(), Box> { time("linear 65536x96 -> 384 (mlp in)", || w1.forward(&flat).unwrap()); let w2 = Linear::new(Tensor::randn(0f32, 0.1, (96, 384), &d)?, Some(Tensor::zeros(96, DType::F32, &d)?)); time("linear 65536x384 -> 96 (mlp out)", || w2.forward(&big).unwrap()); + // block 1 (the first stage change): pooling and the wide projections + let p1 = Tensor::randn(0f32, 1.0, (1, 192, 256, 256), &d)?; + time("max_pool2d 2x2 on 1 x 192 x 256 x 256", || p1.max_pool2d_with_stride(2, 2).unwrap()); + let p2 = Tensor::randn(0f32, 1.0, (1024, 192, 8, 8), &d)?; + time("max_pool2d 2x2 on 1024 x 192 x 8 x 8", || p2.max_pool2d_with_stride(2, 2).unwrap()); + let nhwc = Tensor::randn(0f32, 1.0, (1, 256, 256, 192), &d)?; + time("permute NHWC -> NCHW + contiguous (192ch)", || nhwc.permute((0, 3, 1, 2)).unwrap().contiguous().unwrap()); + let w3 = Linear::new(Tensor::randn(0f32, 0.1, (192, 96), &d)?, Some(Tensor::zeros(192, DType::F32, &d)?)); + time("linear 65536x96 -> 192 (block 1 proj)", || w3.forward(&flat).unwrap()); + time("NHWC 2x2 max via reshape + max(2), max(3)", || nhwc.reshape((1, 128, 2, 128, 2, 192)).unwrap().max(2).unwrap().max(3).unwrap()); + time("NHWC 2x2 max via two strided maxes of narrowed", || { + let a = nhwc.reshape((1, 128, 2, 256, 192)).unwrap(); + let m = a.narrow(2, 0, 1).unwrap().maximum(&a.narrow(2, 1, 1).unwrap()).unwrap().squeeze(2).unwrap(); + let b = m.reshape((1, 128, 128, 2, 192)).unwrap(); + b.narrow(3, 0, 1).unwrap().maximum(&b.narrow(3, 1, 1).unwrap()).unwrap().squeeze(3).unwrap() + }); println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get())); Ok(()) } diff --git a/crates/sam2/examples/sam2_bench.rs b/crates/sam2/examples/sam2_bench.rs index 7bc678b7c..ea871a4f0 100644 --- a/crates/sam2/examples/sam2_bench.rs +++ b/crates/sam2/examples/sam2_bench.rs @@ -46,6 +46,16 @@ fn run() -> Result<(), Box> { x = b.forward(&x)?; println!("block {i}: {:.0} ms, out {:?}", t.elapsed().as_secs_f64() * 1e3, x.dims()); } + // The rest of the encode: the neck (1x1 convolutions, top-down sums) and its sine position encodings. + let stages = hiera.forward(&pixels)?; + let t = Instant::now(); + let _ = model.vision.neck.forward(&stages)?; + println!("neck (convolutions + positions): {:.0} ms", t.elapsed().as_secs_f64() * 1e3); + let t = Instant::now(); + for (h, w) in [(32, 32), (64, 64), (128, 128), (256, 256)] { + let _ = effectcraft_sam2::neck::sine_position(128, h, w)?; + } + println!("sine positions for the four levels: {:.0} ms", t.elapsed().as_secs_f64() * 1e3); println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get())); Ok(()) } diff --git a/crates/sam2/src/hiera.rs b/crates/sam2/src/hiera.rs index 90833f7da..61becd418 100644 --- a/crates/sam2/src/hiera.rs +++ b/crates/sam2/src/hiera.rs @@ -48,10 +48,16 @@ fn at(v: &[T], i: usize, what: &str) -> Result { v.get(i).copied().ok_or_else(|| Error::Weights(format!("Hiera config has no {what} for stage {i}"))) } -/// Max pool a `[B, H, W, C]` tensor by `stride` (kernel = stride, no padding, floor). +/// Max pool a `[B, H, W, C]` tensor by `stride` (kernel = stride, no padding, floor), without the +/// permutes a pooling layer would need: it groups the cells and takes the maximum over each group. fn do_pool(x: &Tensor, stride: usize) -> Result { - let x = x.permute((0, 3, 1, 2))?.contiguous()?.max_pool2d_with_stride(stride, stride)?; - Ok(x.permute((0, 2, 3, 1))?.contiguous()?) + let (b, h, w, c) = x.dims4()?; + if stride == 0 { + return Err(Error::Input("pool stride 0".into())); + } + let (oh, ow) = (h / stride, w / stride); + let x = if oh * stride != h || ow * stride != w { x.narrow(1, 0, oh * stride)?.narrow(2, 0, ow * stride)? } else { x.clone() }; + Ok(x.contiguous()?.reshape((b, oh, stride, ow, stride, c))?.max(2)?.max(3)?.contiguous()?) } /// Split `[B, H, W, C]` into `[B * nH * nW, ws, ws, C]` windows, zero padding the bottom/right. @@ -353,6 +359,17 @@ mod tests { assert_eq!(d, 0.0); } + #[test] + fn pooling_equals_a_max_pool_layer_including_odd_sizes() { + let x = Tensor::randn(0f32, 1.0, (2, 7, 9, 5), &Device::Cpu).unwrap(); + let want = x.permute((0, 3, 1, 2)).unwrap().contiguous().unwrap().max_pool2d_with_stride(2, 2).unwrap().permute((0, 2, 3, 1)).unwrap(); + let got = do_pool(&x, 2).unwrap(); + assert_eq!(got.dims(), want.dims()); + let d = (got - want).unwrap().abs().unwrap().sum_all().unwrap().to_scalar::().unwrap(); + assert_eq!(d, 0.0); + assert!(do_pool(&x, 0).is_err()); + } + #[test] fn zero_window_is_an_error_not_a_panic() { let x = Tensor::zeros((1, 4, 4, 2), candle_core::DType::F32, &Device::Cpu).unwrap(); diff --git a/crates/sam2/src/neck.rs b/crates/sam2/src/neck.rs index 0a82ec7a9..02a253db3 100644 --- a/crates/sam2/src/neck.rs +++ b/crates/sam2/src/neck.rs @@ -3,6 +3,9 @@ // HuggingFace Inc. team) to Rust on candle. See NOTICE. //! The feature pyramid neck on top of Hiera, and the vision encoder that joins them. +use std::collections::HashMap; +use std::sync::{Mutex, OnceLock, PoisonError}; + use candle_core::{Device, Tensor}; use crate::hiera::{Hiera, HieraConfig}; @@ -10,7 +13,23 @@ use crate::weights::Weights; use crate::{Error, Result}; /// Sine position encoding `[1, 2 * num_pos_features, h, w]`, normalised (`Sam2SinePositionEmbedding`). +/// It depends only on the sizes, so each is built once (it was 150 ms of a 2.7 s frame). pub fn sine_position(num_pos_features: usize, h: usize, w: usize) -> Result { + static CACHE: OnceLock>> = OnceLock::new(); + let cache = CACHE.get_or_init(Default::default); + if let Some(t) = cache.lock().unwrap_or_else(PoisonError::into_inner).get(&(num_pos_features, h, w)) { + return Ok(t.clone()); + } + let t = build_sine_position(num_pos_features, h, w)?; + let mut map = cache.lock().unwrap_or_else(PoisonError::into_inner); + if map.len() >= 16 { + map.clear(); + } + map.insert((num_pos_features, h, w), t.clone()); + Ok(t) +} + +fn build_sine_position(num_pos_features: usize, h: usize, w: usize) -> Result { let eps = 1e-6f32; let scale = 2.0 * std::f32::consts::PI; let dim_t: Vec = (0..num_pos_features).map(|i| 10000f32.powf(2.0 * (i / 2) as f32 / num_pos_features as f32)).collect(); From d6fcd4a62b37d089dd9514c10143b5457ac1716c Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:53:31 +0200 Subject: [PATCH 34/77] S2.60: measured tracked-frame cost; bench times memory attention A tracked frame is the encode (2.4 s) plus memory attention (2.4 s with a full bank), about 5 s on 16 threads, not one encode as the docs said. Docs corrected; chunked attention tried and reverted. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/examples/sam2_bench.rs | 14 ++++++++++++++ docs/roto-brush-sam2.md | 17 +++++++++++++---- 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/crates/sam2/examples/sam2_bench.rs b/crates/sam2/examples/sam2_bench.rs index ea871a4f0..a4e78feb0 100644 --- a/crates/sam2/examples/sam2_bench.rs +++ b/crates/sam2/examples/sam2_bench.rs @@ -56,6 +56,20 @@ fn run() -> Result<(), Box> { let _ = effectcraft_sam2::neck::sine_position(128, h, w)?; } println!("sine positions for the four levels: {:.0} ms", t.elapsed().as_secs_f64() * 1e3); + // What a tracked frame adds to the encode: memory attention over the memory bank (7 memory frames and + // 16 object pointers of 4 tokens), then the memory encoder. + let weights = effectcraft_sam2::weights::Weights::open(&file)?; + let attn = effectcraft_sam2::memattn::MemoryAttention::load(&weights)?; + let cur = Tensor::randn(0f32, 1.0, (4096, 256), &Device::Cpu)?; + let cur_pos = Tensor::randn(0f32, 1.0, (4096, 256), &Device::Cpu)?; + let m = 7 * 4096 + 64; + let mem = Tensor::randn(0f32, 1.0, (m, 64), &Device::Cpu)?; + let mem_pos = Tensor::randn(0f32, 1.0, (m, 64), &Device::Cpu)?; + for run in 0..2 { + let t = Instant::now(); + let _ = attn.forward(&cur, &cur_pos, &mem, &mem_pos, 64)?; + println!("memory attention (run {run}, {m} memory tokens): {:.0} ms", t.elapsed().as_secs_f64() * 1e3); + } println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get())); Ok(()) } diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index e61f49be7..580553655 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -64,10 +64,19 @@ the CPU. SAM 2 is prompted with at most three foreground and three background points per frame: it was trained on a handful of clicks, and measured confidence collapses with many (0.99 with 1–7 points, 0.25 with 16, 0.00 with 32). -## Speed (this PC: 16 threads, CPU only) -A 1024 × 1024 frame takes about **3.4 s** to encode (release build); a click on an encoded frame about 120 ms. Tracking a -clip costs one encode per frame, so a 300-frame clip takes on the order of 17 minutes on a CPU. Candle has no AMD/Vulkan -backend; GPU work is Phase 6 of the plan (the early Hiera blocks are memory-bound and have headroom on the CPU too). +## Speed (this PC: 16 threads, CPU only, release build) +| Step | Time | +|---|---| +| Encode one 1024 × 1024 frame | about **2.4 s** (3.4 s before the speed work: a parallel GELU, cached position encodings, pooling without permutes) | +| A click on an encoded frame | about 105 ms | +| Memory attention on a tracked frame (full bank of 7 memory frames) | about **2.4 s** | +| A tracked frame in all | about **5 s** (the first six have fewer memories) | + +A 300-frame clip therefore takes on the order of 25 minutes on this CPU. (An earlier version of this page said one encode +per frame; the memory attention is as large as the encode.) What was tried and did not help: chunking the queries of the +global and memory attention to keep scores in cache. Candle has no AMD/Vulkan backend, so a GPU path is future work (a Burn +or hand-written wgpu route); `cargo run --release -p effectcraft-sam2 --example sam2_bench` and `--example op_bench` time the +pieces. ## Benchmark `effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt]` runs synthetic clips From 921bb00ba6906b8209f043503bf9c52b1b8057f9 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 21:01:17 +0200 Subject: [PATCH 35/77] S2.61: parallel bias add fused with GELU; GPU matmul spike Candle adds a bias on one thread; the Linear layers now add it (and apply GELU for the MLP in) in one parallel pass. Encode 2.4 s -> 2.1 s, memory attention 2.4 -> 2.3 s; reference tests unchanged. gpu_matmul example: a tiled WGSL matmul on the RX 7900 GRE reaches 1.9-3.6 TFLOPS against 0.16-0.47 for the CPU. Co-Authored-By: Claude Sonnet 5.5 --- Cargo.lock | 3 + crates/sam2/Cargo.toml | 6 ++ crates/sam2/examples/gpu_matmul.rs | 164 +++++++++++++++++++++++++++++ crates/sam2/src/memory.rs | 4 +- crates/sam2/src/nn.rs | 76 +++++++++++-- docs/roto-brush-sam2.md | 15 +-- 6 files changed, 251 insertions(+), 17 deletions(-) create mode 100644 crates/sam2/examples/gpu_matmul.rs diff --git a/Cargo.lock b/Cargo.lock index a9b24748d..9b174e6cf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1832,10 +1832,13 @@ dependencies = [ name = "effectcraft-sam2" version = "0.18.0" dependencies = [ + "bytemuck", "candle-core", "candle-nn", "effectcraft-segment", + "pollster 0.4.0", "rayon", + "wgpu", ] [[package]] diff --git a/crates/sam2/Cargo.toml b/crates/sam2/Cargo.toml index 61afb4117..d1d74ebc6 100644 --- a/crates/sam2/Cargo.toml +++ b/crates/sam2/Cargo.toml @@ -17,3 +17,9 @@ rayon = { workspace = true } [target.'cfg(not(target_arch = "wasm32"))'.dependencies] candle-core = "=0.9.2" candle-nn = "=0.9.2" + +# The GPU matmul spike (examples/gpu_matmul.rs) only. +[dev-dependencies] +wgpu = { workspace = true } +pollster = { workspace = true } +bytemuck = { workspace = true } diff --git a/crates/sam2/examples/gpu_matmul.rs b/crates/sam2/examples/gpu_matmul.rs new file mode 100644 index 000000000..1429ebf02 --- /dev/null +++ b/crates/sam2/examples/gpu_matmul.rs @@ -0,0 +1,164 @@ +//! A spike for the GPU path (plan phase 6): how fast is a plain tiled WGSL matrix multiply on this +//! machine's GPU, against the candle CPU matmul the encoder uses today? Candle has no AMD/Vulkan +//! backend; this measures what hand-written wgpu kernels could reach on the shapes SAM 2 needs. +//! +//! `cargo run --release -p effectcraft-sam2 --example gpu_matmul` + +use std::time::Instant; + +use candle_core::{Device, Tensor}; +use wgpu::util::DeviceExt; + +const SHADER: &str = r#" +struct P { m: u32, n: u32, k: u32, pad: u32 } +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b: array; +@group(0) @binding(2) var c: array; +@group(0) @binding(3) var p: P; +var sa: array; +var sb: array; +@compute @workgroup_size(16, 16) +fn main(@builtin(workgroup_id) wg: vec3, @builtin(local_invocation_id) li: vec3) { + let row0 = wg.y * 64u; + let col0 = wg.x * 64u; + let tid = li.y * 16u + li.x; + var acc: array; + for (var i = 0u; i < 16u; i++) { acc[i] = 0.0; } + for (var kt = 0u; kt < p.k; kt += 16u) { + for (var i = 0u; i < 4u; i++) { + let idx = tid + i * 256u; + let r = idx / 16u; + let cc = idx % 16u; + let gr = row0 + r; + let gc = kt + cc; + var va = 0.0; + if (gr < p.m && gc < p.k) { va = a[gr * p.k + gc]; } + sa[idx] = va; + let r2 = idx / 64u; + let c2 = idx % 64u; + let gr2 = kt + r2; + let gc2 = col0 + c2; + var vb = 0.0; + if (gr2 < p.k && gc2 < p.n) { vb = b[gr2 * p.n + gc2]; } + sb[idx] = vb; + } + workgroupBarrier(); + for (var kk = 0u; kk < 16u; kk++) { + var ar: array; + var br: array; + for (var i = 0u; i < 4u; i++) { + ar[i] = sa[(li.y * 4u + i) * 16u + kk]; + br[i] = sb[kk * 64u + li.x * 4u + i]; + } + for (var i = 0u; i < 4u; i++) { + for (var j = 0u; j < 4u; j++) { acc[i * 4u + j] += ar[i] * br[j]; } + } + } + workgroupBarrier(); + } + for (var i = 0u; i < 4u; i++) { + for (var j = 0u; j < 4u; j++) { + let gr = row0 + li.y * 4u + i; + let gc = col0 + li.x * 4u + j; + if (gr < p.m && gc < p.n) { c[gr * p.n + gc] = acc[i * 4u + j]; } + } + } +} +"#; + +fn main() -> Result<(), Box> { + let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle()); + let adapter = pollster::block_on( + instance.request_adapter(&wgpu::RequestAdapterOptions { power_preference: wgpu::PowerPreference::HighPerformance, ..Default::default() }), + )?; + println!("adapter: {:?}", adapter.get_info().name); + let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor { required_limits: adapter.limits(), ..Default::default() }))?; + let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: None, source: wgpu::ShaderSource::Wgsl(SHADER.into()) }); + let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: None, + layout: None, + module: &module, + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + // (name, m, k, n): the matmul shapes of SAM 2's encoder and memory attention + let shapes = [ + ("hiera block 0 qkv", 65536usize, 96usize, 288usize), + ("hiera block 0 mlp in", 65536, 96, 384), + ("hiera global attention q.kT (one head)", 4096, 96, 4096), + ("memory attention q.kT (full bank)", 4096, 256, 28736), + ]; + for (name, m, k, n) in shapes { + let a = Tensor::randn(0f32, 1.0, (m, k), &Device::Cpu)?; + let b = Tensor::randn(0f32, 1.0, (k, n), &Device::Cpu)?; + let av = a.flatten_all()?.to_vec1::()?; + let bv = b.flatten_all()?.to_vec1::()?; + let ba = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: bytemuck::cast_slice(&av), + usage: wgpu::BufferUsages::STORAGE, + }); + let bb = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: bytemuck::cast_slice(&bv), + usage: wgpu::BufferUsages::STORAGE, + }); + let bc = device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: (m * n * 4) as u64, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + let params = [m as u32, n as u32, k as u32, 0]; + let bp = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: bytemuck::cast_slice(¶ms), + usage: wgpu::BufferUsages::UNIFORM, + }); + let bg = device.create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &pipeline.get_bind_group_layout(0), + entries: &[ + wgpu::BindGroupEntry { binding: 0, resource: ba.as_entire_binding() }, + wgpu::BindGroupEntry { binding: 1, resource: bb.as_entire_binding() }, + wgpu::BindGroupEntry { binding: 2, resource: bc.as_entire_binding() }, + wgpu::BindGroupEntry { binding: 3, resource: bp.as_entire_binding() }, + ], + }); + let run = |reps: usize| { + let mut enc = device.create_command_encoder(&Default::default()); + for _ in 0..reps { + let mut pass = enc.begin_compute_pass(&Default::default()); + pass.set_pipeline(&pipeline); + pass.set_bind_group(0, &bg, &[]); + pass.dispatch_workgroups(n.div_ceil(64) as u32, m.div_ceil(64) as u32, 1); + } + queue.submit([enc.finish()]); + let _ = device.poll(wgpu::PollType::wait_indefinitely()); + }; + run(1); + let reps = 5; + let t = Instant::now(); + run(reps); + let gpu = t.elapsed().as_secs_f64() / reps as f64; + let flops = 2.0 * m as f64 * k as f64 * n as f64; + // the CPU: candle's matmul, as the encoder uses it + let _ = a.matmul(&b)?; + let t = Instant::now(); + let cpu_reps = 3; + for _ in 0..cpu_reps { + let _ = a.matmul(&b)?; + } + let cpu = t.elapsed().as_secs_f64() / cpu_reps as f64; + println!( + "{name:<42} gpu {:>7.1} ms ({:>6.0} GFLOPS) cpu {:>7.1} ms ({:>5.0} GFLOPS) x{:.1}", + gpu * 1e3, + flops / gpu / 1e9, + cpu * 1e3, + flops / cpu / 1e9, + cpu / gpu + ); + } + Ok(()) +} diff --git a/crates/sam2/src/memory.rs b/crates/sam2/src/memory.rs index 735c4a563..45b470a3f 100644 --- a/crates/sam2/src/memory.rs +++ b/crates/sam2/src/memory.rs @@ -8,7 +8,7 @@ use candle_core::Tensor; use crate::Result; use crate::neck::sine_position; -use crate::nn::{LayerNorm, apply, gelu, layer_norm, linear}; +use crate::nn::{LayerNorm, apply, apply_gelu, gelu, layer_norm, linear}; use crate::weights::Weights; use candle_nn::Linear; @@ -49,7 +49,7 @@ impl FuserBlock { fn forward(&self, x: &Tensor) -> Result { let h = norm_channels_first(&self.norm, &self.depthwise.forward(x)?)?; let h = h.permute((0, 2, 3, 1))?.contiguous()?; - let h = apply(&self.pw2, &gelu(&apply(&self.pw1, &h)?)?)?; + let h = apply(&self.pw2, &apply_gelu(&self.pw1, &h)?)?; let h = h.broadcast_mul(&self.scale)?.permute((0, 3, 1, 2))?.contiguous()?; Ok((x + h)?) } diff --git a/crates/sam2/src/nn.rs b/crates/sam2/src/nn.rs index 76bafabac..256a60da9 100644 --- a/crates/sam2/src/nn.rs +++ b/crates/sam2/src/nn.rs @@ -1,8 +1,8 @@ // Modified work: small layers used by the ported SAM 2 modules (Apache-2.0, see NOTICE). //! Linear, layer norm and a few tensor helpers on candle. -use candle_core::{CpuStorage, CustomOp1, Layout, Shape, Tensor}; -use candle_nn::{Linear, Module}; +use candle_core::{CpuStorage, CustomOp1, CustomOp2, Layout, Shape, Tensor}; +use candle_nn::Linear; use crate::weights::Weights; use crate::{Error, Result}; @@ -41,23 +41,83 @@ impl Mlp { } pub fn forward(&self, x: &Tensor) -> Result { - apply(&self.proj_out, &gelu(&apply(&self.proj_in, x)?)?) + apply(&self.proj_out, &apply_gelu(&self.proj_in, x)?) } } -/// Apply a linear layer to `[.., in]` through a flat 2-D matmul. Candle's `Linear::forward` -/// broadcasts the weight over every leading dimension (copying it per window), which is slow -/// on the 4-D activations of the backbone. -pub fn apply(l: &Linear, x: &Tensor) -> Result { +/// `x + bias` (and, for [`bias_gelu`], then GELU) over rows of `x`, on all cores. Candle adds a bias +/// on one thread, which cost more than the matmul it followed on the wide activations of the early +/// Hiera blocks. +struct BiasAdd { + gelu: bool, +} + +impl CustomOp2 for BiasAdd { + fn name(&self) -> &'static str { + "bias_add_parallel" + } + + fn cpu_fwd(&self, xs: &CpuStorage, xl: &Layout, bs: &CpuStorage, bl: &Layout) -> candle_core::Result<(CpuStorage, Shape)> { + use rayon::prelude::*; + let (CpuStorage::F32(x), CpuStorage::F32(b)) = (xs, bs) else { candle_core::bail!("bias_add_parallel needs f32") }; + let (Some((x0, x1)), Some((b0, b1))) = (xl.contiguous_offsets(), bl.contiguous_offsets()) else { + candle_core::bail!("bias_add_parallel needs contiguous inputs") + }; + let (x, b) = match (x.get(x0..x1), b.get(b0..b1)) { + (Some(x), Some(b)) if !b.is_empty() && x.len() % b.len() == 0 => (x, b), + _ => candle_core::bail!("bias_add_parallel: sizes do not fit"), + }; + let cols = b.len(); + let mut out = vec![0f32; x.len()]; + let rows_per = (64 * 1024 / cols).max(1); + let gelu = self.gelu; + out.par_chunks_mut(rows_per * cols).zip(x.par_chunks(rows_per * cols)).for_each(|(o, i)| { + for (orow, irow) in o.chunks_mut(cols).zip(i.chunks(cols)) { + for ((a, v), bias) in orow.iter_mut().zip(irow).zip(b) { + let y = v + bias; + *a = if gelu { (candle_core::cpu::erf::erf_f32(y * std::f32::consts::FRAC_1_SQRT_2) + 1.0) * 0.5 * y } else { y }; + } + } + }); + Ok((CpuStorage::F32(out), xl.shape().clone())) + } +} + +fn bias_op(x: &Tensor, bias: &Tensor, gelu: bool) -> Result { + if x.dtype() != candle_core::DType::F32 || !x.device().is_cpu() { + let y = x.broadcast_add(bias)?; + return Ok(if gelu { y.gelu_erf()? } else { y }); + } + Ok(x.contiguous()?.apply_op2_no_bwd(&bias.contiguous()?, &BiasAdd { gelu })?) +} + +fn matmul_bias(l: &Linear, x: &Tensor, gelu: bool) -> Result { let dims = x.dims(); let (&last, lead) = dims.split_last().ok_or_else(|| Error::Compute("linear on a scalar".into()))?; let rows: usize = lead.iter().product(); - let y = l.forward(&x.reshape((rows, last))?)?; + let y = x.reshape((rows, last))?.matmul(&l.weight().t()?)?; + let y = match l.bias() { + Some(b) => bias_op(&y, b, gelu)?, + None if gelu => y.gelu_erf()?, + None => y, + }; let mut out: Vec = lead.to_vec(); out.push(y.dim(1)?); Ok(y.reshape(out)?) } +/// Apply a linear layer to `[.., in]` through a flat 2-D matmul and a parallel bias add. Candle's +/// `Linear::forward` broadcasts the weight over every leading dimension (copying it per window) and +/// adds the bias on one thread: both slow on the 4-D activations of the backbone. +pub fn apply(l: &Linear, x: &Tensor) -> Result { + matmul_bias(l, x, false) +} + +/// [`apply`] followed by the exact GELU, in the same pass over the output. +pub fn apply_gelu(l: &Linear, x: &Tensor) -> Result { + matmul_bias(l, x, true) +} + /// Exact (erf) GELU on all cores. Candle's own `gelu_erf` runs on one thread, which made it a third of /// the time of the Hiera blocks; this gives the same values (it calls the same `erf` per element). struct GeluErf; diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index 580553655..29a9e83a7 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -67,15 +67,16 @@ of clicks, and measured confidence collapses with many (0.99 with 1–7 points, ## Speed (this PC: 16 threads, CPU only, release build) | Step | Time | |---|---| -| Encode one 1024 × 1024 frame | about **2.4 s** (3.4 s before the speed work: a parallel GELU, cached position encodings, pooling without permutes) | -| A click on an encoded frame | about 105 ms | -| Memory attention on a tracked frame (full bank of 7 memory frames) | about **2.4 s** | -| A tracked frame in all | about **5 s** (the first six have fewer memories) | +| Encode one 1024 × 1024 frame | about **2.1 s** (3.4 s before the speed work: a parallel GELU, cached position encodings, pooling without permutes, a parallel bias add) | +| A click on an encoded frame | about 100 ms | +| Memory attention on a tracked frame (full bank of 7 memory frames) | about **2.3 s** | +| A tracked frame in all | about **4.6 s** (the first six have fewer memories) | -A 300-frame clip therefore takes on the order of 25 minutes on this CPU. (An earlier version of this page said one encode +A 300-frame clip therefore takes on the order of 23 minutes on this CPU. (An earlier version of this page said one encode per frame; the memory attention is as large as the encode.) What was tried and did not help: chunking the queries of the -global and memory attention to keep scores in cache. Candle has no AMD/Vulkan backend, so a GPU path is future work (a Burn -or hand-written wgpu route); `cargo run --release -p effectcraft-sam2 --example sam2_bench` and `--example op_bench` time the +global and memory attention to keep scores in cache. Candle has no AMD/Vulkan backend, so a GPU path is future work. A spike (`--example gpu_matmul`) measured a plain tiled +WGSL matmul on the RX 7900 GRE at 1,900–3,600 GFLOPS against 160–470 for candle's CPU matmul on the same shapes (6–13×), +so hand-written wgpu kernels could bring a tracked frame to a second or less; `cargo run --release -p effectcraft-sam2 --example sam2_bench` and `--example op_bench` time the pieces. ## Benchmark From 5c7eb3c2ae0b6516307b0babd59aa89f7b2591c6 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 21:10:14 +0200 Subject: [PATCH 36/77] S2.62: GPU encoder (wgpu): Hiera + neck on any GPU, matches PyTorch to 7e-6 Strided batched matmul with bias, GELU/ReLU and residual in the epilogue, layer norm, softmax, window partition, pooling, patch embedding, RoPE and neck kernels (WGSL, all tested against CPU definitions). GpuEncoder keeps activations on the GPU and returns the 64x64 level and the decoder's skip features. On the RX 7900 GRE one frame encodes in 0.29 s (debug build) against 2.1 s on the CPU in release; errors 6.8e-6 vs PyTorch, 1.2e-5 vs the CPU path. Co-Authored-By: Claude Sonnet 5.5 --- Cargo.lock | 1 + crates/sam2/Cargo.toml | 5 +- crates/sam2/src/gpu/encoder.rs | 286 ++++++++++++++++++++ crates/sam2/src/gpu/kernels.wgsl | 98 +++++++ crates/sam2/src/gpu/kernels_ops.wgsl | 173 ++++++++++++ crates/sam2/src/gpu/mod.rs | 376 +++++++++++++++++++++++++++ crates/sam2/src/gpu/ops.rs | 167 ++++++++++++ crates/sam2/src/gpu/tests.rs | 301 +++++++++++++++++++++ crates/sam2/src/hiera.rs | 54 +++- crates/sam2/src/lib.rs | 1 + crates/sam2/src/neck.rs | 4 +- crates/sam2/tests/reference.rs | 40 +++ 12 files changed, 1490 insertions(+), 16 deletions(-) create mode 100644 crates/sam2/src/gpu/encoder.rs create mode 100644 crates/sam2/src/gpu/kernels.wgsl create mode 100644 crates/sam2/src/gpu/kernels_ops.wgsl create mode 100644 crates/sam2/src/gpu/mod.rs create mode 100644 crates/sam2/src/gpu/ops.rs create mode 100644 crates/sam2/src/gpu/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 9b174e6cf..ce98aeadc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1836,6 +1836,7 @@ dependencies = [ "candle-core", "candle-nn", "effectcraft-segment", + "log", "pollster 0.4.0", "rayon", "wgpu", diff --git a/crates/sam2/Cargo.toml b/crates/sam2/Cargo.toml index d1d74ebc6..ff8e64c04 100644 --- a/crates/sam2/Cargo.toml +++ b/crates/sam2/Cargo.toml @@ -12,14 +12,13 @@ workspace = true [dependencies] effectcraft-segment = { workspace = true } rayon = { workspace = true } +log = { workspace = true } # Native only (D1): candle is not built for wasm32, where this crate is empty. [target.'cfg(not(target_arch = "wasm32"))'.dependencies] candle-core = "=0.9.2" candle-nn = "=0.9.2" - -# The GPU matmul spike (examples/gpu_matmul.rs) only. -[dev-dependencies] +# The optional GPU path (wgpu compute, any vendor); the CPU path stays the reference. wgpu = { workspace = true } pollster = { workspace = true } bytemuck = { workspace = true } diff --git a/crates/sam2/src/gpu/encoder.rs b/crates/sam2/src/gpu/encoder.rs new file mode 100644 index 000000000..9a6caeeb4 --- /dev/null +++ b/crates/sam2/src/gpu/encoder.rs @@ -0,0 +1,286 @@ +// Original work for EffectCraft; the maths follows the Hugging Face `transformers` SAM 2 +// implementation (Apache-2.0, see NOTICE). +//! The vision encoder on the GPU: patch embedding, the Hiera blocks and the neck, with the decoder's two +//! skip convolutions. Activations stay on the GPU between the kernels; only the three tensors the rest of +//! the model needs come back. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex, PoisonError}; + +use candle_core::{Device, Tensor}; + +use super::{Act, Buf, GLinear, GNorm, Gpu, Matmul, Strides}; +use crate::hiera::{BlockPlan, Hiera, HieraConfig}; +use crate::weights::Weights; +use crate::{Error, Result}; + +struct GBlock { + plan: BlockPlan, + norm1: GNorm, + qkv: GLinear, + proj: GLinear, + norm2: GNorm, + mlp_in: GLinear, + mlp_out: GLinear, + block_proj: Option, +} + +/// What the rest of the model takes from the encoder (CPU tensors, NCHW). +pub struct GpuFeatures { + /// The lowest-resolution neck level `[1, 256, 64, 64]`. + pub low: Tensor, + /// The decoder's skip features: `conv_s0` of the highest level `[1, 32, 256, 256]` and `conv_s1` of the + /// next `[1, 64, 128, 128]`. + pub high: (Tensor, Tensor), +} + +pub struct GpuEncoder { + gpu: Arc, + cfg: HieraConfig, + patch_w: Buf, + patch_b: Buf, + blocks: Vec, + stage_ends: Vec, + neck: Vec, + top_down: Vec, + skip0: GLinear, + skip1: GLinear, + /// The position embedding of each grid size, uploaded on first use. + pos: Mutex>, + /// The CPU backbone, which owns the position embedding tables. + cpu: Arc, + eps: f32, +} + +impl GpuEncoder { + /// Upload the backbone, the neck and the skip convolutions. `cpu` supplies the position embedding. + pub fn load(gpu: Arc, w: &Weights, cpu: Arc) -> Result { + let cfg = cpu.config().clone(); + let flat = |t: Tensor| -> Result> { Ok(t.flatten_all()?.to_vec1::()?) }; + let hidden = cfg.embed_dims.first().copied().unwrap_or(96); + let mut blocks = vec![]; + for (total, plan) in BlockPlan::all(&cfg)?.into_iter().enumerate() { + let p = format!("vision_encoder.backbone.blocks.{total}"); + blocks.push(GBlock { + plan, + norm1: GNorm::load(&gpu, w, &format!("{p}.layer_norm1"), plan.dim)?, + qkv: GLinear::load(&gpu, w, &format!("{p}.attn.qkv"), plan.dim, plan.dim_out * 3)?, + proj: GLinear::load(&gpu, w, &format!("{p}.attn.proj"), plan.dim_out, plan.dim_out)?, + norm2: GNorm::load(&gpu, w, &format!("{p}.layer_norm2"), plan.dim_out)?, + mlp_in: GLinear::load(&gpu, w, &format!("{p}.mlp.proj_in"), plan.dim_out, plan.dim_out * cfg.mlp_ratio)?, + mlp_out: GLinear::load(&gpu, w, &format!("{p}.mlp.proj_out"), plan.dim_out * cfg.mlp_ratio, plan.dim_out)?, + block_proj: (plan.dim != plan.dim_out).then(|| GLinear::load(&gpu, w, &format!("{p}.proj"), plan.dim, plan.dim_out)).transpose()?, + }); + } + let mut stage_ends = vec![]; + let mut n = 0; + for b in &cfg.blocks { + n += b; + stage_ends.push(n.saturating_sub(1)); + } + let channels: Vec = cfg.embed_dims.iter().rev().copied().collect(); + let neck = channels + .iter() + .enumerate() + .map(|(j, c)| GLinear::load_conv1x1(&gpu, w, &format!("vision_encoder.neck.convs.{j}"), *c, 256)) + .collect::>>()?; + Ok(GpuEncoder { + patch_w: gpu.upload(&flat(w.get("vision_encoder.backbone.patch_embed.projection.weight", &[hidden, 3, 7, 7])?)?), + patch_b: gpu.upload(&flat(w.get("vision_encoder.backbone.patch_embed.projection.bias", &[hidden])?)?), + skip0: GLinear::load_conv1x1(&gpu, w, "mask_decoder.conv_s0", 256, 32)?, + skip1: GLinear::load_conv1x1(&gpu, w, "mask_decoder.conv_s1", 256, 64)?, + gpu, + eps: cfg.eps as f32, + cfg, + blocks, + stage_ends, + neck, + top_down: vec![2, 3], + pos: Mutex::default(), + cpu, + }) + } + + fn linear(&self, x: &Buf, rows: usize, l: &GLinear, out: &Buf, act: Act, residual: Option<&Buf>) -> Result<()> { + self.gpu.linear(x, rows, l, out, act, residual) + } + + fn norm(&self, x: &Buf, rows: usize, dim: usize, n: &GNorm) -> Result { + let out = self.gpu.zeros(rows * dim)?; + self.gpu.layernorm(x, &n.w, &n.b, &out, rows, dim, self.eps)?; + Ok(out) + } + + fn pos_buffer(&self, h: usize, w: usize) -> Result { + let mut map = self.pos.lock().unwrap_or_else(PoisonError::into_inner); + if let Some(b) = map.get(&(h, w)) { + return Ok(b.clone()); + } + let table = self.cpu.pos(h, w)?.flatten_all()?.to_vec1::()?; + let buf = self.gpu.upload(&table); + map.insert((h, w), buf.clone()); + Ok(buf) + } + + /// One Hiera block on `[h * w, dim]` tokens; returns the new tokens and grid size. + fn block(&self, b: &GBlock, x: &Buf, h: usize, w: usize) -> Result<(Buf, usize, usize)> { + let g = &self.gpu; + let BlockPlan { dim, dim_out, heads, window, query_stride } = b.plan; + let n = h * w; + let normed = self.norm(x, n, dim, &b.norm1)?; + // the skip connection: from the block's input, or (stage change) its projection, pooled with the queries + let (oh, ow) = query_stride.map_or((h, w), |s| (h / s, w / s)); + if query_stride.is_some_and(|s| s != 2) { + return Err(Error::Input("only a query stride of 2 is supported on the GPU".into())); + } + let residual = match (&b.block_proj, query_stride) { + (Some(p), Some(_)) => { + let proj = g.zeros(n * dim_out)?; + self.linear(&normed, n, p, &proj, Act::None, None)?; + let pooled = g.zeros(oh * ow * dim_out)?; + g.maxpool2(&proj, &pooled, 1, h, w, dim_out, (n * dim_out, w * dim_out, dim_out))?; + pooled + } + (Some(p), None) => { + let proj = g.zeros(n * dim_out)?; + self.linear(&normed, n, p, &proj, Act::None, None)?; + proj + } + (None, _) => x.clone(), + }; + // windows (or the whole grid for global attention) + let (hw, nwin, t, win_h, win_w) = if window > 0 { + let (nwx, nwy) = (w.div_ceil(window), h.div_ceil(window)); + let buf = g.zeros(nwx * nwy * window * window * dim)?; + g.window_partition(&normed, &buf, h, w, dim, window)?; + (buf, nwx * nwy, window * window, window, window) + } else { + (normed, 1, n, h, w) + }; + let tok = nwin * t; + let qkv = g.zeros(tok * 3 * dim_out)?; + self.linear(&hw, tok, &b.qkv, &qkv, Act::None, None)?; + // queries, pooled when the stage changes + let (q_buf, q_rs, tq) = match query_stride { + Some(s) => { + let (qh, qw) = (win_h / s, win_w / s); + let pooled = g.zeros(nwin * qh * qw * dim_out)?; + g.maxpool2(&qkv, &pooled, nwin, win_h, win_w, dim_out, (t * 3 * dim_out, win_w * 3 * dim_out, 3 * dim_out))?; + (pooled, dim_out, qh * qw) + } + None => (qkv.clone(), 3 * dim_out, t), + }; + let hd = dim_out / heads.max(1); + let scale = (hd as f32).powf(-0.5); + let scores = g.zeros(nwin * tq * t)?; + let attn = g.zeros(nwin * tq * dim_out)?; + for hh in 0..heads { + g.matmul(&Matmul { + m: tq, + n: t, + k: hd, + batch: nwin, + a: (&q_buf, Strides { offset: hh * hd, batch: tq * q_rs, row: q_rs, col: 1 }), + b: (&qkv, Strides { offset: dim_out + hh * hd, batch: t * 3 * dim_out, row: 1, col: 3 * dim_out }), + bias: None, + c: (&scores, Strides { offset: 0, batch: tq * t, row: t, col: 1 }), + residual: None, + alpha: scale, + act: Act::None, + })?; + g.softmax(&scores, nwin * tq, t)?; + g.matmul(&Matmul { + m: tq, + n: hd, + k: t, + batch: nwin, + a: (&scores, Strides { offset: 0, batch: tq * t, row: t, col: 1 }), + b: (&qkv, Strides { offset: 2 * dim_out + hh * hd, batch: t * 3 * dim_out, row: 3 * dim_out, col: 1 }), + bias: None, + c: (&attn, Strides { offset: hh * hd, batch: tq * dim_out, row: dim_out, col: 1 }), + residual: None, + alpha: 1.0, + act: Act::None, + })?; + } + let projected = g.zeros(nwin * tq * dim_out)?; + self.linear(&attn, nwin * tq, &b.proj, &projected, Act::None, None)?; + let attn_img = if window > 0 { + let ws = window / query_stride.unwrap_or(1); + let img = g.zeros(oh * ow * dim_out)?; + g.window_unpartition(&projected, &img, oh, ow, dim_out, ws)?; + img + } else { + projected + }; + let n2 = oh * ow; + let x2 = g.zeros(n2 * dim_out)?; + g.add(&residual, &attn_img, &x2, n2 * dim_out)?; + let normed2 = self.norm(&x2, n2, dim_out, &b.norm2)?; + let hidden = g.zeros(n2 * b.mlp_in.out)?; + self.linear(&normed2, n2, &b.mlp_in, &hidden, Act::Gelu, None)?; + let out = g.zeros(n2 * dim_out)?; + self.linear(&hidden, n2, &b.mlp_out, &out, Act::None, Some(&x2))?; + Ok((out, oh, ow)) + } + + /// Encode normalised `[1, 3, H, W]` pixels (`H` and `W` multiples of 32). + pub fn encode(&self, pixels: &Tensor) -> Result { + let g = &self.gpu; + let (_, c, ph, pw) = pixels.dims4()?; + let window0 = self.cfg.windows.first().copied().unwrap_or(8); + if c != 3 || ph == 0 || pw == 0 || ph % 4 != 0 || pw % 4 != 0 || (ph / 4) % window0 != 0 || (pw / 4) % window0 != 0 { + return Err(Error::Input(format!("image size {pw}x{ph} does not fit the encoder"))); + } + let hidden = self.cfg.embed_dims.first().copied().unwrap_or(96); + let px = g.upload(&pixels.contiguous()?.flatten_all()?.to_vec1::()?); + let (mut h, mut w) = (ph / 4, pw / 4); + let tokens = g.zeros(h * w * hidden)?; + g.patch_embed(&px, &self.patch_w, &self.patch_b, &tokens, ph, pw, hidden)?; + let mut x = g.zeros(h * w * hidden)?; + g.add(&tokens, &self.pos_buffer(h, w)?, &x, h * w * hidden)?; + let mut stages: Vec<(Buf, usize, usize, usize)> = vec![]; + for (i, b) in self.blocks.iter().enumerate() { + let (nx, nh, nw) = self.block(b, &x, h, w)?; + (x, h, w) = (nx, nh, nw); + if self.stage_ends.contains(&i) { + stages.push((x.clone(), h, w, b.plan.dim_out)); + } + } + // the neck, from the lowest resolution up: lateral 1x1 convolutions, and the top-down sum on the levels that have it + let n = self.neck.len(); + if stages.len() != n { + return Err(Error::Compute("the backbone returned the wrong number of stages".into())); + } + let mut levels: Vec<(Buf, usize, usize)> = vec![]; + let mut prev: Option<(Buf, usize, usize)> = None; + for i in (0..n).rev() { + let (buf, sh, sw, _) = &stages[i]; + let conv = self.neck.get(n - 1 - i).ok_or_else(|| Error::Compute("missing neck convolution".into()))?; + let lateral = g.zeros(sh * sw * 256)?; + self.linear(buf, sh * sw, conv, &lateral, Act::None, None)?; + let cur = match (&prev, self.top_down.contains(&i) && i != n - 1) { + (Some((p, _, _)), true) => { + let sum = g.zeros(sh * sw * 256)?; + g.upsample_add(&lateral, p, &sum, *sh, *sw, 256)?; + sum + } + _ => lateral, + }; + prev = Some((cur.clone(), *sh, *sw)); + levels.push((cur, *sh, *sw)); + } + // levels: lowest resolution first. The decoder needs the third (64 x 64) and the skip convolutions of the last two. + let [_, (l2, h2, w2), (l1, h1, w1), (l0, h0, w0)] = <[_; 4]>::try_from(levels).map_err(|_| Error::Compute("the neck must have four levels".into()))?; + let s0 = g.zeros(h0 * w0 * 32)?; + self.linear(&l0, h0 * w0, &self.skip0, &s0, Act::None, None)?; + let s1 = g.zeros(h1 * w1 * 64)?; + self.linear(&l1, h1 * w1, &self.skip1, &s1, Act::None, None)?; + g.flush()?; + let nchw = |buf: &Buf, hh: usize, ww: usize, ch: usize| -> Result { + let v = g.download(buf)?; + Ok(Tensor::from_vec(v, (1, hh, ww, ch), &Device::Cpu)?.permute((0, 3, 1, 2))?.contiguous()?) + }; + Ok(GpuFeatures { low: nchw(&l2, h2, w2, 256)?, high: (nchw(&s0, h0, w0, 32)?, nchw(&s1, h1, w1, 64)?) }) + } +} diff --git a/crates/sam2/src/gpu/kernels.wgsl b/crates/sam2/src/gpu/kernels.wgsl new file mode 100644 index 000000000..cce4c2a05 --- /dev/null +++ b/crates/sam2/src/gpu/kernels.wgsl @@ -0,0 +1,98 @@ +// SAM 2 compute kernels (f32). Original work for EffectCraft; the maths follows the Hugging Face +// `transformers` SAM 2 implementation (Apache-2.0, see NOTICE). +// +// Every kernel reads its parameters from one uniform: 24 u32 words (floats are bit-cast). + +struct P { v: array, 6> } +@group(0) @binding(0) var in0: array; +@group(0) @binding(1) var in1: array; +@group(0) @binding(2) var in2: array; +@group(0) @binding(3) var out0: array; +@group(0) @binding(4) var in3: array; +@group(0) @binding(5) var p: P; + +fn u(i: u32) -> u32 { return p.v[i / 4u][i % 4u]; } +fn f(i: u32) -> f32 { return bitcast(u(i)); } + +// erf, Abramowitz and Stegun 7.1.26 (abs error 1.5e-7). +fn erf_(x: f32) -> f32 { + let s = select(1.0, -1.0, x < 0.0); + let a = abs(x); + let t = 1.0 / (1.0 + 0.3275911 * a); + let y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-a * a); + return s * y; +} +fn gelu_(x: f32) -> f32 { return (erf_(x * 0.70710678) + 1.0) * 0.5 * x; } + +// ---------------------------------------------------------------- strided batched matmul +// C[b] = act(alpha * A[b] (m x k) . B[b] (k x n) + bias) (+ residual), with free strides for A and B. +// words: 0 m, 1 n, 2 k, 3 batch, 4 a_off, 5 a_bs, 6 a_rs, 7 a_cs, 8 b_off, 9 b_bs, 10 b_rs, 11 b_cs, +// 12 c_off, 13 c_bs, 14 c_rs, 15 flags (bit0 bias, bits 1-2 act: 1 gelu 2 relu, bit3 residual), +// 16 alpha (f32) +// in0 = A, in1 = B, in2 = bias (n), out0 = C, in3 = residual (same layout as C) +var sa: array; +var sb: array; + +@compute @workgroup_size(16, 16) +fn matmul(@builtin(workgroup_id) wg: vec3, @builtin(local_invocation_id) li: vec3) { + let m = u(0u); let n = u(1u); let k = u(2u); + let a_base = u(4u) + wg.z * u(5u); + let b_base = u(8u) + wg.z * u(9u); + let a_rs = u(6u); let a_cs = u(7u); let b_rs = u(10u); let b_cs = u(11u); + let row0 = wg.y * 64u; + let col0 = wg.x * 64u; + let tid = li.y * 16u + li.x; + let a_k_contig = a_cs == 1u; + let b_n_contig = b_cs == 1u; + var acc: array; + for (var i = 0u; i < 16u; i++) { acc[i] = 0.0; } + for (var kt = 0u; kt < k; kt += 16u) { + for (var i = 0u; i < 4u; i++) { + let idx = tid + i * 256u; + var r: u32; var c: u32; + if (a_k_contig) { r = idx / 16u; c = idx % 16u; } else { r = idx % 64u; c = idx / 64u; } + let gr = row0 + r; let gc = kt + c; + var va = 0.0; + if (gr < m && gc < k) { va = in0[a_base + gr * a_rs + gc * a_cs]; } + sa[r * 16u + c] = va; + var r2: u32; var c2: u32; + if (b_n_contig) { r2 = idx / 64u; c2 = idx % 64u; } else { r2 = idx % 16u; c2 = idx / 16u; } + let gr2 = kt + r2; let gc2 = col0 + c2; + var vb = 0.0; + if (gr2 < k && gc2 < n) { vb = in1[b_base + gr2 * b_rs + gc2 * b_cs]; } + sb[r2 * 64u + c2] = vb; + } + workgroupBarrier(); + for (var kk = 0u; kk < 16u; kk++) { + var ar: array; + var br: array; + for (var i = 0u; i < 4u; i++) { + ar[i] = sa[(li.y * 4u + i) * 16u + kk]; + br[i] = sb[kk * 64u + li.x * 4u + i]; + } + for (var i = 0u; i < 4u; i++) { + for (var j = 0u; j < 4u; j++) { acc[i * 4u + j] += ar[i] * br[j]; } + } + } + workgroupBarrier(); + } + let flags = u(15u); + let alpha = f(16u); + let c_base = u(12u) + wg.z * u(13u); + let c_rs = u(14u); + let act = (flags >> 1u) & 3u; + for (var i = 0u; i < 4u; i++) { + for (var j = 0u; j < 4u; j++) { + let gr = row0 + li.y * 4u + i; + let gc = col0 + li.x * 4u + j; + if (gr < m && gc < n) { + var v = acc[i * 4u + j] * alpha; + if ((flags & 1u) != 0u) { v += in2[gc]; } + if (act == 1u) { v = gelu_(v); } else if (act == 2u) { v = max(v, 0.0); } + let idx = c_base + gr * c_rs + gc; + if ((flags & 8u) != 0u) { v += in3[idx]; } + out0[idx] = v; + } + } + } +} diff --git a/crates/sam2/src/gpu/kernels_ops.wgsl b/crates/sam2/src/gpu/kernels_ops.wgsl new file mode 100644 index 000000000..aae0c8323 --- /dev/null +++ b/crates/sam2/src/gpu/kernels_ops.wgsl @@ -0,0 +1,173 @@ + +// ---------------------------------------------------------------- helpers +// A 1-D index over n items spread on a 2-D grid of 256-thread groups (a dimension holds 65535 groups). +fn gid1(wg: vec3, nw: vec3, li: vec3) -> u32 { return (wg.y * nw.x + wg.x) * 256u + li.x; } + +var red: array; + +// Sum of `red` over the workgroup (64 threads), for a caller that has filled it; all threads get the result. +fn reduce_sum(li: u32) -> f32 { + workgroupBarrier(); + for (var st = 32u; st > 0u; st = st / 2u) { + if (li < st) { red[li] += red[li + st]; } + workgroupBarrier(); + } + let r = red[0]; + workgroupBarrier(); + return r; +} + +fn reduce_max(li: u32) -> f32 { + workgroupBarrier(); + for (var st = 32u; st > 0u; st = st / 2u) { + if (li < st) { red[li] = max(red[li], red[li + st]); } + workgroupBarrier(); + } + let r = red[0]; + workgroupBarrier(); + return r; +} + +// ---------------------------------------------------------------- layer norm over the last dimension +// words: 0 rows, 1 cols, 2 eps (f32). in0 x, in1 weight, in2 bias, out0 y. One workgroup per row. +@compute @workgroup_size(64) +fn layernorm(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: vec3, @builtin(local_invocation_id) li: vec3) { + let row = wg.y * nw.x + wg.x; + let rows = u(0u); let cols = u(1u); + if (row >= rows) { return; } + let base = row * cols; + var s = 0.0; + for (var c = li.x; c < cols; c += 64u) { s += in0[base + c]; } + red[li.x] = s; + let mean = reduce_sum(li.x) / f32(cols); + var q = 0.0; + for (var c = li.x; c < cols; c += 64u) { let d = in0[base + c] - mean; q += d * d; } + red[li.x] = q; + let inv = inverseSqrt(reduce_sum(li.x) / f32(cols) + f(2u)); + for (var c = li.x; c < cols; c += 64u) { out0[base + c] = (in0[base + c] - mean) * inv * in1[c] + in2[c]; } +} + +// ---------------------------------------------------------------- softmax over rows, in place on out0 +// words: 0 rows, 1 cols. One workgroup per row. +@compute @workgroup_size(64) +fn softmax(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: vec3, @builtin(local_invocation_id) li: vec3) { + let row = wg.y * nw.x + wg.x; + let rows = u(0u); let cols = u(1u); + if (row >= rows) { return; } + let base = row * cols; + var mx = -3.0e38; + for (var c = li.x; c < cols; c += 64u) { mx = max(mx, out0[base + c]); } + red[li.x] = mx; + let m = reduce_max(li.x); + var s = 0.0; + for (var c = li.x; c < cols; c += 64u) { let e = exp(out0[base + c] - m); out0[base + c] = e; s += e; } + red[li.x] = s; + let inv = 1.0 / reduce_sum(li.x); + for (var c = li.x; c < cols; c += 64u) { out0[base + c] = out0[base + c] * inv; } +} + +// ---------------------------------------------------------------- elementwise add: out0 = in0 + in1 (n items) +@compute @workgroup_size(256) +fn add(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: vec3, @builtin(local_invocation_id) li: vec3) { + let i = gid1(wg, nw, li); + if (i < u(0u)) { out0[i] = in0[i] + in1[i]; } +} + +// ---------------------------------------------------------------- windows +// words: 0 n (output elements), 1 H, 2 W, 3 C, 4 ws, 5 windows across (nwx) +// window_partition: [H, W, C] -> [nwin, ws, ws, C] (zero padded); window_unpartition: the inverse. +@compute @workgroup_size(256) +fn window_partition(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: vec3, @builtin(local_invocation_id) li: vec3) { + let i = gid1(wg, nw, li); + if (i >= u(0u)) { return; } + let H = u(1u); let W = u(2u); let C = u(3u); let ws = u(4u); let nwx = u(5u); + let c = i % C; var t = i / C; + let x = t % ws; t = t / ws; + let y = t % ws; let win = t / ws; + let gy = (win / nwx) * ws + y; + let gx = (win % nwx) * ws + x; + var v = 0.0; + if (gy < H && gx < W) { v = in0[(gy * W + gx) * C + c]; } + out0[i] = v; +} + +@compute @workgroup_size(256) +fn window_unpartition(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: vec3, @builtin(local_invocation_id) li: vec3) { + let i = gid1(wg, nw, li); + if (i >= u(0u)) { return; } + let W = u(2u); let C = u(3u); let ws = u(4u); let nwx = u(5u); + let c = i % C; let t = i / C; + let gx = t % W; let gy = t / W; + let win = (gy / ws) * nwx + gx / ws; + out0[i] = in0[((win * ws + gy % ws) * ws + gx % ws) * C + c]; +} + +// ---------------------------------------------------------------- 2 x 2 max pool over [nb, h, w, c] tokens +// words: 0 n (output elements), 1 nb, 2 h, 3 w, 4 c, 5 src image stride, 6 src row stride, 7 src token stride +// (so it can read the q third of a qkv buffer in place). Output [nb, h/2, w/2, c] contiguous. +@compute @workgroup_size(256) +fn maxpool2(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: vec3, @builtin(local_invocation_id) li: vec3) { + let i = gid1(wg, nw, li); + if (i >= u(0u)) { return; } + let h = u(2u); let w = u(3u); let c = u(4u); + let oh = h / 2u; let ow = w / 2u; + let ch = i % c; var t = i / c; + let ox = t % ow; t = t / ow; + let oy = t % oh; let img = t / oh; + let base = img * u(5u) + (oy * 2u) * u(6u) + (ox * 2u) * u(7u) + ch; + out0[i] = max(max(in0[base], in0[base + u(7u)]), max(in0[base + u(6u)], in0[base + u(6u) + u(7u)])); +} + +// ---------------------------------------------------------------- patch embedding: 7 x 7 stride 4 pad 3 convolution +// words: 0 n (output elements), 1 H, 2 W, 3 Co, 4 Ho, 5 Wo. in0 pixels [3, H, W], in1 weight [Co, 3, 7, 7], +// in2 bias [Co]. Output [Ho, Wo, Co]. +@compute @workgroup_size(256) +fn patch_embed(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: vec3, @builtin(local_invocation_id) li: vec3) { + let i = gid1(wg, nw, li); + if (i >= u(0u)) { return; } + let H = u(1u); let W = u(2u); let Co = u(3u); let Wo = u(5u); + let co = i % Co; let t = i / Co; + let ox = t % Wo; let oy = t / Wo; + var acc = in2[co]; + for (var ci = 0u; ci < 3u; ci++) { + for (var ky = 0u; ky < 7u; ky++) { + let iy = i32(oy * 4u + ky) - 3; + if (iy < 0 || iy >= i32(H)) { continue; } + for (var kx = 0u; kx < 7u; kx++) { + let ix = i32(ox * 4u + kx) - 3; + if (ix < 0 || ix >= i32(W)) { continue; } + acc += in0[ci * H * W + u32(iy) * W + u32(ix)] * in1[((co * 3u + ci) * 7u + ky) * 7u + kx]; + } + } + } + out0[i] = acc; +} + +// ---------------------------------------------------------------- rotary position encoding on [rows, cols] +// words: 0 n (= rows * cols), 1 cols, 2 table rows, 3 rows rotated (the rest is copied). in0 x, in1 cos, in2 sin +// (tables [table rows, cols]); pairs (x0, x1) -> (x0 c - x1 s, x1 c + x0 s). +@compute @workgroup_size(256) +fn rope(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: vec3, @builtin(local_invocation_id) li: vec3) { + let i = gid1(wg, nw, li); + if (i >= u(0u)) { return; } + let cols = u(1u); + let row = i / cols; let col = i % cols; + if (row >= u(3u)) { out0[i] = in0[i]; return; } + let t = (row % u(2u)) * cols + col; + let x = in0[i]; + var pair = 0.0; + if ((col & 1u) == 1u) { pair = in0[i - 1u]; } else { pair = in0[i + 1u]; } + if ((col & 1u) == 0u) { out0[i] = x * in1[t] - pair * in2[t]; } else { out0[i] = x * in1[t] + pair * in2[t]; } +} + +// ---------------------------------------------------------------- neck: lateral + nearest 2x upsample of prev +// words: 0 n (output elements), 1 h, 2 w, 3 c. in0 lateral [h, w, c], in1 prev [h/2, w/2, c]. +@compute @workgroup_size(256) +fn upsample_add(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: vec3, @builtin(local_invocation_id) li: vec3) { + let i = gid1(wg, nw, li); + if (i >= u(0u)) { return; } + let w = u(2u); let c = u(3u); + let ch = i % c; let t = i / c; + let x = t % w; let y = t / w; + out0[i] = in0[i] + in1[(((y / 2u) * (w / 2u)) + x / 2u) * c + ch]; +} diff --git a/crates/sam2/src/gpu/mod.rs b/crates/sam2/src/gpu/mod.rs new file mode 100644 index 000000000..a14df4b16 --- /dev/null +++ b/crates/sam2/src/gpu/mod.rs @@ -0,0 +1,376 @@ +// Original work for EffectCraft (wgpu compute kernels); the maths follows the Hugging Face +// `transformers` SAM 2 implementation (Apache-2.0, see NOTICE). +//! An optional GPU path for the heavy parts of SAM 2 (the Hiera backbone, the neck and the memory +//! attention), as wgpu compute kernels: it runs on AMD, NVIDIA and Intel alike, where candle has +//! no backend for the first. The CPU path stays the reference: every kernel is tested against it, +//! and any GPU failure (no adapter, an out-of-memory error, a lost device) is an `Err` the caller +//! answers by running the CPU path. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex, PoisonError}; + +use wgpu::util::DeviceExt; + +use crate::{Error, Result}; + +const SHADER: &str = concat!(include_str!("kernels.wgsl"), include_str!("kernels_ops.wgsl")); +/// Words in a kernel's parameter block. +const PARAM_WORDS: usize = 24; + +/// A buffer of f32 on the GPU (cheap to clone: it is a handle). +#[derive(Clone)] +pub struct Buf { + pub(crate) buffer: wgpu::Buffer, + /// Elements (f32). + pub len: usize, +} + +/// One strided matrix operand: `(offset, batch stride, row stride, column stride)` in elements. +#[derive(Clone, Copy, Debug)] +pub struct Strides { + pub offset: usize, + pub batch: usize, + pub row: usize, + pub col: usize, +} + +impl Strides { + /// A row-major `[rows, cols]` matrix (one batch). + pub fn row_major(cols: usize) -> Strides { + Strides { offset: 0, batch: 0, row: cols, col: 1 } + } +} + +/// What to do to a matrix product after the multiplication. +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub enum Act { + #[default] + None, + Gelu, + Relu, +} + +/// A batched matrix product `C[b] = act(alpha * A[b] . B[b] + bias) (+ residual)`. +pub struct Matmul<'a> { + pub m: usize, + pub n: usize, + pub k: usize, + pub batch: usize, + pub a: (&'a Buf, Strides), + pub b: (&'a Buf, Strides), + pub bias: Option<&'a Buf>, + pub c: (&'a Buf, Strides), + pub residual: Option<&'a Buf>, + pub alpha: f32, + pub act: Act, +} + +/// The GPU device and its compiled kernels. +pub struct Gpu { + device: wgpu::Device, + queue: wgpu::Queue, + module: wgpu::ShaderModule, + pipelines: Mutex>>, + encoder: Mutex>, + error: Arc>>, + adapter: String, +} + +fn words(v: &[u32]) -> [u32; PARAM_WORDS] { + let mut out = [0u32; PARAM_WORDS]; + for (o, x) in out.iter_mut().zip(v) { + *o = *x; + } + out +} + +fn u32_of(v: usize, what: &str) -> Result { + u32::try_from(v).map_err(|_| Error::Compute(format!("{what} {v} does not fit a GPU index"))) +} + +impl Gpu { + /// A device on the best real GPU (a software adapter is refused: the CPU path is faster). + pub fn new() -> Result { + let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle()); + let adapter = pollster::block_on( + instance.request_adapter(&wgpu::RequestAdapterOptions { power_preference: wgpu::PowerPreference::HighPerformance, ..Default::default() }), + ) + .map_err(|e| Error::Compute(format!("no GPU adapter: {e}")))?; + let info = adapter.get_info(); + if info.device_type == wgpu::DeviceType::Cpu { + return Err(Error::Compute(format!("only a software adapter ({}) is available", info.name))); + } + let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor { + label: Some("effectcraft sam2"), + required_limits: adapter.limits(), + ..Default::default() + })) + .map_err(|e| Error::Compute(format!("no GPU device: {e}")))?; + let error: Arc>> = Arc::default(); + let sink = error.clone(); + device.on_uncaptured_error(Arc::new(move |e| { + log::error!("sam2 gpu: {e}"); + sink.lock().unwrap_or_else(PoisonError::into_inner).get_or_insert_with(|| e.to_string()); + })); + let sink = error.clone(); + device.set_device_lost_callback(move |reason, message| { + sink.lock().unwrap_or_else(PoisonError::into_inner).get_or_insert_with(|| format!("GPU device lost ({reason:?}): {message}")); + }); + let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("sam2 kernels"), source: wgpu::ShaderSource::Wgsl(SHADER.into()) }); + Ok(Gpu { device, queue, module, pipelines: Mutex::default(), encoder: Mutex::new(None), error, adapter: info.name }) + } + + pub fn adapter_name(&self) -> &str { + &self.adapter + } + + fn take_error(&self) -> Result<()> { + match self.error.lock().unwrap_or_else(PoisonError::into_inner).take() { + Some(e) => Err(Error::Compute(format!("GPU: {e}"))), + None => Ok(()), + } + } + + pub fn upload(&self, data: &[f32]) -> Buf { + let buffer = self.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: bytemuck::cast_slice(data), + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + }); + Buf { buffer, len: data.len() } + } + + /// A zeroed buffer of `len` elements. + pub fn zeros(&self, len: usize) -> Result { + let bytes = (len.max(1) as u64).checked_mul(4).ok_or_else(|| Error::Budget("GPU buffer size overflows".into()))?; + if bytes > self.device.limits().max_buffer_size || bytes > self.device.limits().max_storage_buffer_binding_size { + return Err(Error::Budget(format!("a {} MB GPU buffer is over this GPU's limit", bytes >> 20))); + } + let buffer = self.device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: bytes, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + Ok(Buf { buffer, len }) + } + + fn pipeline(&self, entry: &'static str) -> Arc { + let mut map = self.pipelines.lock().unwrap_or_else(PoisonError::into_inner); + map.entry(entry) + .or_insert_with(|| { + Arc::new(self.device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some(entry), + layout: None, + module: &self.module, + entry_point: Some(entry), + compilation_options: Default::default(), + cache: None, + })) + }) + .clone() + } + + /// Queue one dispatch of `entry`. `bindings` are the kernel's buffers by binding index. + pub(crate) fn run(&self, entry: &'static str, bindings: &[(u32, &Buf)], params: &[u32], groups: (u32, u32, u32)) { + let pipeline = self.pipeline(entry); + let uniform = self.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: bytemuck::cast_slice(&words(params)), + usage: wgpu::BufferUsages::UNIFORM, + }); + let mut entries: Vec> = + bindings.iter().map(|(i, b)| wgpu::BindGroupEntry { binding: *i, resource: b.buffer.as_entire_binding() }).collect(); + entries.push(wgpu::BindGroupEntry { binding: 5, resource: uniform.as_entire_binding() }); + let group = self.device.create_bind_group(&wgpu::BindGroupDescriptor { label: None, layout: &pipeline.get_bind_group_layout(0), entries: &entries }); + let mut slot = self.encoder.lock().unwrap_or_else(PoisonError::into_inner); + let enc = slot.get_or_insert_with(|| self.device.create_command_encoder(&Default::default())); + let mut pass = enc.begin_compute_pass(&Default::default()); + pass.set_pipeline(&pipeline); + pass.set_bind_group(0, &group, &[]); + pass.dispatch_workgroups(groups.0, groups.1, groups.2); + } + + /// Submit what is queued and wait for it to finish. + pub fn flush(&self) -> Result<()> { + let enc = self.encoder.lock().unwrap_or_else(PoisonError::into_inner).take(); + if let Some(enc) = enc { + self.queue.submit([enc.finish()]); + } + let _ = self.device.poll(wgpu::PollType::wait_indefinitely()); + self.take_error() + } + + /// Copy a buffer back to the CPU (after submitting everything queued). + pub fn download(&self, buf: &Buf) -> Result> { + let bytes = buf.len as u64 * 4; + if bytes == 0 { + return Ok(vec![]); + } + let staging = self.device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: bytes, + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let mut enc = + self.encoder.lock().unwrap_or_else(PoisonError::into_inner).take().unwrap_or_else(|| self.device.create_command_encoder(&Default::default())); + enc.copy_buffer_to_buffer(&buf.buffer, 0, &staging, 0, bytes); + self.queue.submit([enc.finish()]); + let (tx, rx) = std::sync::mpsc::channel(); + staging.map_async(wgpu::MapMode::Read, .., move |r| { + let _ = tx.send(r); + }); + let _ = self.device.poll(wgpu::PollType::wait_indefinitely()); + self.take_error()?; + rx.recv().map_err(|_| Error::Compute("GPU readback never completed".into()))?.map_err(|e| Error::Compute(format!("GPU readback: {e}")))?; + let out = { + let view = staging.get_mapped_range(..).map_err(|e| Error::Compute(format!("GPU readback: {e}")))?; + bytemuck::cast_slice::(&view).to_vec() + }; + staging.unmap(); + Ok(out) + } + + /// Queue `C[b] = act(alpha * A[b] . B[b] + bias) (+ residual)`. + pub fn matmul(&self, op: &Matmul<'_>) -> Result<()> { + let (a, sa) = op.a; + let (b, sb) = op.b; + let (c, sc) = op.c; + if op.m == 0 || op.n == 0 || op.batch == 0 { + return Ok(()); + } + // 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; + 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 op.bias.is_some_and(|bias| bias.len < op.n) || op.residual.is_some_and(|r| r.len < c.len) { + return Err(Error::Input("matmul bias or residual is too small".into())); + } + let flags = u32::from(op.bias.is_some()) + | (match op.act { + Act::None => 0, + Act::Gelu => 1, + Act::Relu => 2, + } << 1) + | (u32::from(op.residual.is_some()) << 3); + let params = [ + u32_of(op.m, "m")?, + u32_of(op.n, "n")?, + u32_of(op.k, "k")?, + u32_of(op.batch, "batch")?, + u32_of(sa.offset, "a offset")?, + u32_of(sa.batch, "a batch")?, + u32_of(sa.row, "a row")?, + u32_of(sa.col, "a col")?, + u32_of(sb.offset, "b offset")?, + u32_of(sb.batch, "b batch")?, + u32_of(sb.row, "b row")?, + u32_of(sb.col, "b col")?, + u32_of(sc.offset, "c offset")?, + u32_of(sc.batch, "c batch")?, + u32_of(sc.row, "c row")?, + flags, + op.alpha.to_bits(), + ]; + let dummy; + let bias = match op.bias { + Some(b) => b, + None => { + dummy = self.upload(&[0.0]); + &dummy + } + }; + // a buffer is never bound twice in one dispatch (read and read-write), so no residual means a dummy + let dummy_res; + let res = match op.residual { + Some(r) => r, + None => { + dummy_res = self.upload(&[0.0]); + &dummy_res + } + }; + self.run( + "matmul", + &[(0, a), (1, b), (2, bias), (3, c), (4, res)], + ¶ms, + (u32_of(op.n.div_ceil(64), "groups")?, u32_of(op.m.div_ceil(64), "groups")?, u32_of(op.batch, "batch")?), + ); + Ok(()) + } +} + +pub mod encoder; +mod ops; + +/// A linear layer's weight `[out, in]` and bias `[out]` on the GPU. +pub struct GLinear { + pub(crate) w: Buf, + pub(crate) b: Buf, + pub(crate) inp: usize, + pub(crate) out: usize, +} + +/// A layer norm's scale and shift on the GPU. +pub struct GNorm { + pub(crate) w: Buf, + pub(crate) b: Buf, +} + +fn flat(t: candle_core::Tensor) -> Result> { + Ok(t.flatten_all()?.to_vec1::()?) +} + +impl GLinear { + /// From `{name}.weight` (`[out, in]`) and `{name}.bias`. + pub fn load(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize) -> Result { + Ok(GLinear { + w: g.upload(&flat(w.get(&format!("{name}.weight"), &[out, inp])?)?), + b: g.upload(&flat(w.get(&format!("{name}.bias"), &[out])?)?), + inp, + out, + }) + } + + /// A 1 x 1 convolution (`[out, in, 1, 1]`) as a linear layer over channels-last tokens. + pub fn load_conv1x1(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize) -> Result { + Ok(GLinear { + w: g.upload(&flat(w.get(&format!("{name}.weight"), &[out, inp, 1, 1])?)?), + b: g.upload(&flat(w.get(&format!("{name}.bias"), &[out])?)?), + inp, + out, + }) + } +} + +impl GNorm { + pub fn load(g: &Gpu, w: &crate::weights::Weights, name: &str, dim: usize) -> Result { + Ok(GNorm { w: g.upload(&flat(w.get(&format!("{name}.weight"), &[dim])?)?), b: g.upload(&flat(w.get(&format!("{name}.bias"), &[dim])?)?) }) + } +} + +impl Gpu { + /// `out = act(x . W^T + bias) (+ residual)` for `rows` rows of `l.inp` values. + pub fn linear(&self, x: &Buf, rows: usize, l: &GLinear, out: &Buf, act: Act, residual: Option<&Buf>) -> Result<()> { + // the weight is stored [out, in]: B[k][j] = w[j * in + k] + self.matmul(&Matmul { + m: rows, + n: l.out, + k: l.inp, + batch: 1, + a: (x, Strides::row_major(l.inp)), + b: (&l.w, Strides { offset: 0, batch: 0, row: 1, col: l.inp }), + bias: Some(&l.b), + c: (out, Strides::row_major(l.out)), + residual, + alpha: 1.0, + act, + }) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/sam2/src/gpu/ops.rs b/crates/sam2/src/gpu/ops.rs new file mode 100644 index 000000000..7bb99c495 --- /dev/null +++ b/crates/sam2/src/gpu/ops.rs @@ -0,0 +1,167 @@ +//! The non-matmul kernels, as methods of [`Gpu`] (each queues one dispatch). + +use super::{Buf, Gpu}; +use crate::{Error, Result}; + +fn n32(v: usize, what: &str) -> Result { + u32::try_from(v).map_err(|_| Error::Compute(format!("{what} {v} does not fit a GPU index"))) +} + +/// A 2-D grid of groups covering `groups` (a dimension holds 65535). +fn grid(groups: usize) -> Result<(u32, u32, u32)> { + let gx = groups.clamp(1, 65535); + Ok((n32(gx, "groups")?, n32(groups.div_ceil(gx).max(1), "groups")?, 1)) +} + +/// One 256-thread group per 256 items. +fn items(n: usize) -> Result<(u32, u32, u32)> { + grid(n.div_ceil(256)) +} + +fn need(buf: &Buf, len: usize, what: &str) -> Result<()> { + if buf.len < len { + return Err(Error::Input(format!("{what}: a buffer of {} elements is smaller than the {len} needed", buf.len))); + } + Ok(()) +} + +impl Gpu { + /// `y = layer_norm(x) * weight + bias` over the last dimension of a `[rows, cols]` buffer. + pub fn layernorm(&self, x: &Buf, weight: &Buf, bias: &Buf, y: &Buf, rows: usize, cols: usize, eps: f32) -> Result<()> { + need(x, rows * cols, "layernorm x")?; + need(y, rows * cols, "layernorm y")?; + need(weight, cols, "layernorm weight")?; + need(bias, cols, "layernorm bias")?; + self.run("layernorm", &[(0, x), (1, weight), (2, bias), (3, y)], &[n32(rows, "rows")?, n32(cols, "cols")?, eps.to_bits()], grid(rows)?); + Ok(()) + } + + /// Softmax over each of `rows` rows of `cols`, in place. + pub fn softmax(&self, x: &Buf, rows: usize, cols: usize) -> Result<()> { + need(x, rows * cols, "softmax")?; + self.run("softmax", &[(3, x)], &[n32(rows, "rows")?, n32(cols, "cols")?], grid(rows)?); + Ok(()) + } + + /// `out = a + b` over `n` elements. + pub fn add(&self, a: &Buf, b: &Buf, out: &Buf, n: usize) -> Result<()> { + need(a, n, "add a")?; + need(b, n, "add b")?; + need(out, n, "add out")?; + self.run("add", &[(0, a), (1, b), (3, out)], &[n32(n, "n")?], items(n)?); + Ok(()) + } + + /// `[h, w, c]` to `[windows, ws, ws, c]`, zero padded at the bottom and right. + pub fn window_partition(&self, x: &Buf, out: &Buf, h: usize, w: usize, c: usize, ws: usize) -> Result<()> { + if ws == 0 { + return Err(Error::Input("window size 0".into())); + } + let (nwx, nwy) = (w.div_ceil(ws), h.div_ceil(ws)); + let n = nwx * nwy * ws * ws * c; + need(x, h * w * c, "window_partition input")?; + need(out, n, "window_partition output")?; + self.run( + "window_partition", + &[(0, x), (3, out)], + &[n32(n, "n")?, n32(h, "h")?, n32(w, "w")?, n32(c, "c")?, n32(ws, "ws")?, n32(nwx, "windows")?], + items(n)?, + ); + Ok(()) + } + + /// The inverse of [`Gpu::window_partition`] (the padding is dropped). + pub fn window_unpartition(&self, x: &Buf, out: &Buf, h: usize, w: usize, c: usize, ws: usize) -> Result<()> { + if ws == 0 { + return Err(Error::Input("window size 0".into())); + } + let (nwx, nwy) = (w.div_ceil(ws), h.div_ceil(ws)); + need(x, nwx * nwy * ws * ws * c, "window_unpartition input")?; + need(out, h * w * c, "window_unpartition output")?; + self.run( + "window_unpartition", + &[(0, x), (3, out)], + &[n32(h * w * c, "n")?, n32(h, "h")?, n32(w, "w")?, n32(c, "c")?, n32(ws, "ws")?, n32(nwx, "windows")?], + items(h * w * c)?, + ); + Ok(()) + } + + /// 2 x 2 max pool of `nb` images of `[h, w, c]` tokens (`h` and `w` even) read with the given strides + /// (image, row, token) so it can pool one third of a qkv buffer in place; output contiguous. + #[allow(clippy::too_many_arguments)] + pub fn maxpool2(&self, x: &Buf, out: &Buf, nb: usize, h: usize, w: usize, c: usize, strides: (usize, usize, usize)) -> Result<()> { + if !h.is_multiple_of(2) || !w.is_multiple_of(2) { + return Err(Error::Input(format!("maxpool2 needs even sizes, got {w}x{h}"))); + } + let n = nb * (h / 2) * (w / 2) * c; + let last = (nb.saturating_sub(1)) * strides.0 + (h - 1) * strides.1 + (w - 1) * strides.2 + c; + need(x, last, "maxpool2 input")?; + need(out, n, "maxpool2 output")?; + self.run( + "maxpool2", + &[(0, x), (3, out)], + &[ + n32(n, "n")?, + n32(nb, "nb")?, + n32(h, "h")?, + n32(w, "w")?, + n32(c, "c")?, + n32(strides.0, "stride")?, + n32(strides.1, "stride")?, + n32(strides.2, "stride")?, + ], + items(n)?, + ); + Ok(()) + } + + /// The patch embedding: `[3, h, w]` pixels, weight `[co, 3, 7, 7]`, bias `[co]`; output `[h/4, w/4, co]`. + pub fn patch_embed(&self, pixels: &Buf, weight: &Buf, bias: &Buf, out: &Buf, h: usize, w: usize, co: usize) -> Result<()> { + let (ho, wo) = ((h + 6 - 7) / 4 + 1, (w + 6 - 7) / 4 + 1); + need(pixels, 3 * h * w, "patch_embed pixels")?; + need(weight, co * 3 * 49, "patch_embed weight")?; + need(bias, co, "patch_embed bias")?; + need(out, ho * wo * co, "patch_embed output")?; + let n = ho * wo * co; + self.run( + "patch_embed", + &[(0, pixels), (1, weight), (2, bias), (3, out)], + &[n32(n, "n")?, n32(h, "h")?, n32(w, "w")?, n32(co, "co")?, n32(ho, "ho")?, n32(wo, "wo")?], + items(n)?, + ); + Ok(()) + } + + /// Rotary encoding of `[rows, cols]`: rows below `rotated` are rotated with the `[table_rows, cols]` + /// tables (row `r` uses table row `r % table_rows`), the rest copied. + #[allow(clippy::too_many_arguments)] + pub fn rope(&self, x: &Buf, cos: &Buf, sin: &Buf, out: &Buf, rows: usize, cols: usize, table_rows: usize, rotated: usize) -> Result<()> { + if !cols.is_multiple_of(2) || table_rows == 0 { + return Err(Error::Input("rope needs an even width and a table".into())); + } + need(x, rows * cols, "rope input")?; + need(out, rows * cols, "rope output")?; + need(cos, table_rows * cols, "rope cos")?; + need(sin, table_rows * cols, "rope sin")?; + self.run( + "rope", + &[(0, x), (1, cos), (2, sin), (3, out)], + &[n32(rows * cols, "n")?, n32(cols, "cols")?, n32(table_rows, "table")?, n32(rotated.min(rows), "rotated")?], + items(rows * cols)?, + ); + Ok(()) + } + + /// `out = lateral + nearest_upsample_2x(prev)` for `[h, w, c]` (`h` and `w` even). + pub fn upsample_add(&self, lateral: &Buf, prev: &Buf, out: &Buf, h: usize, w: usize, c: usize) -> Result<()> { + if !h.is_multiple_of(2) || !w.is_multiple_of(2) { + return Err(Error::Input(format!("upsample_add needs even sizes, got {w}x{h}"))); + } + need(lateral, h * w * c, "upsample_add lateral")?; + need(prev, h / 2 * (w / 2) * c, "upsample_add prev")?; + need(out, h * w * c, "upsample_add out")?; + self.run("upsample_add", &[(0, lateral), (1, prev), (3, out)], &[n32(h * w * c, "n")?, n32(h, "h")?, n32(w, "w")?, n32(c, "c")?], items(h * w * c)?); + Ok(()) + } +} diff --git a/crates/sam2/src/gpu/tests.rs b/crates/sam2/src/gpu/tests.rs new file mode 100644 index 000000000..92ad20f41 --- /dev/null +++ b/crates/sam2/src/gpu/tests.rs @@ -0,0 +1,301 @@ +//! Each kernel against a plain CPU computation. They skip when there is no GPU. + +use std::sync::{Mutex, PoisonError}; + +use super::*; + +/// The one device the tests share (several devices at once upset the Vulkan loader), held for the whole +/// test so the tests take turns. +pub(crate) struct Shared(#[allow(dead_code)] std::sync::MutexGuard<'static, ()>, &'static Gpu); + +impl std::ops::Deref for Shared { + type Target = Gpu; + fn deref(&self) -> &Gpu { + self.1 + } +} + +pub(crate) fn gpu() -> Option { + static TURN: Mutex<()> = Mutex::new(()); + static DEVICE: std::sync::OnceLock> = std::sync::OnceLock::new(); + let guard = TURN.lock().unwrap_or_else(PoisonError::into_inner); + let device = DEVICE.get_or_init(|| match Gpu::new() { + Ok(g) => Some(g), + Err(e) => { + eprintln!("skipped: {e}"); + None + } + }); + device.as_ref().map(|g| Shared(guard, g)) +} + +/// Deterministic pseudo-random values in -1..1. +pub(crate) fn noise(n: usize, seed: u32) -> Vec { + let mut s = seed.wrapping_mul(2654435761).wrapping_add(1); + (0..n) + .map(|_| { + s ^= s << 13; + s ^= s >> 17; + s ^= s << 5; + (s % 20001) as f32 / 10000.0 - 1.0 + }) + .collect() +} + +pub(crate) fn assert_close(got: &[f32], want: &[f32], tol: f32, what: &str) { + assert_eq!(got.len(), want.len(), "{what}: length"); + let worst = got.iter().zip(want).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max); + assert!(worst <= tol, "{what}: max abs error {worst} (tolerance {tol})"); +} + +fn gelu(x: f32) -> f32 { + use candle_core::{Device, Tensor}; + let t = Tensor::new(&[x], &Device::Cpu).and_then(|t| t.gelu_erf()).and_then(|t| t.to_vec1::()); + t.map(|v| v[0]).unwrap_or(f32::NAN) +} + +/// `C[b][i][j] = sum_k A[b][i][k] * B[b][k][j]` by the definition, with the strides given. +#[allow(clippy::too_many_arguments)] +fn naive(a: &[f32], sa: Strides, b: &[f32], sb: Strides, m: usize, n: usize, k: usize, batch: usize) -> Vec { + let mut c = vec![0f32; batch * m * n]; + for z in 0..batch { + for i in 0..m { + for j in 0..n { + let mut acc = 0f32; + for l in 0..k { + acc += a[sa.offset + z * sa.batch + i * sa.row + l * sa.col] * b[sb.offset + z * sb.batch + l * sb.row + j * sb.col]; + } + c[z * m * n + i * n + j] = acc; + } + } + } + c +} + +#[test] +fn matmul_matches_the_definition_for_awkward_sizes_and_strides() { + let Some(g) = gpu() else { return }; + // (m, n, k, batch): not multiples of the 64 x 64 x 16 tile, and degenerate ones + for (m, n, k, batch) in [(100, 50, 37, 1), (1, 1, 1, 1), (65, 129, 17, 2), (64, 64, 16, 1), (3, 200, 96, 3), (130, 7, 300, 1)] { + let (av, bv) = (noise(batch * m * k, 1), noise(batch * k * n, 2)); + let (a, b) = (g.upload(&av), g.upload(&bv)); + let c = g.zeros(batch * m * n).unwrap(); + let (sa, sb) = (Strides { offset: 0, batch: m * k, row: k, col: 1 }, Strides { offset: 0, batch: k * n, row: n, col: 1 }); + let sc = Strides { offset: 0, batch: m * n, row: n, col: 1 }; + g.matmul(&Matmul { m, n, k, batch, a: (&a, sa), b: (&b, sb), bias: None, c: (&c, sc), residual: None, alpha: 1.0, act: Act::None }).unwrap(); + g.flush().unwrap(); + assert_close(&g.download(&c).unwrap(), &naive(&av, sa, &bv, sb, m, n, k, batch), 2e-4 * (k as f32).sqrt(), &format!("{m}x{n}x{k} x{batch}")); + } +} + +#[test] +fn matmul_with_transposed_operands_bias_activation_and_residual() { + let Some(g) = gpu() else { return }; + let (m, n, k) = (70, 90, 41); + let (av, wv, bias, res) = (noise(m * k, 3), noise(n * k, 4), noise(n, 5), noise(m * n, 6)); + let (a, w, bi, r) = (g.upload(&av), g.upload(&wv), g.upload(&bias), g.upload(&res)); + let c = g.zeros(m * n).unwrap(); + // B is a weight stored [n, k] (a linear layer): B[k][j] = w[j * k + l] + let sb = Strides { offset: 0, batch: 0, row: 1, col: k }; + let sa = Strides::row_major(k); + g.matmul(&Matmul { + m, + n, + k, + batch: 1, + a: (&a, sa), + b: (&w, sb), + bias: Some(&bi), + c: (&c, Strides::row_major(n)), + residual: Some(&r), + alpha: 0.5, + act: Act::Gelu, + }) + .unwrap(); + g.flush().unwrap(); + let base = naive(&av, sa, &wv, sb, m, n, k, 1); + let want: Vec = base.iter().enumerate().map(|(i, v)| gelu(v * 0.5 + bias[i % n]) + res[i]).collect(); + assert_close(&g.download(&c).unwrap(), &want, 1e-3, "linear + gelu + residual"); + // relu, and a transposed A (m contiguous) + let at: Vec = (0..k * m).map(|i| av[(i % m) * k + i / m]).collect(); + let (a2, c2) = (g.upload(&at), g.zeros(m * n).unwrap()); + g.matmul(&Matmul { + m, + n, + k, + batch: 1, + a: (&a2, Strides { offset: 0, batch: 0, row: 1, col: m }), + b: (&w, sb), + bias: None, + c: (&c2, Strides::row_major(n)), + residual: None, + alpha: 1.0, + act: Act::Relu, + }) + .unwrap(); + g.flush().unwrap(); + let want: Vec = base.iter().map(|v| v.max(0.0)).collect(); + assert_close(&g.download(&c2).unwrap(), &want, 1e-3, "transposed A, relu"); +} + +#[test] +fn matmul_refuses_operands_that_overrun_their_buffers() { + let Some(g) = gpu() else { return }; + let (a, b, c) = (g.upload(&[1.0; 6]), g.upload(&[1.0; 6]), g.zeros(4).unwrap()); + let ok = Matmul { + m: 2, + n: 2, + k: 3, + batch: 1, + a: (&a, Strides::row_major(3)), + b: (&b, Strides::row_major(2)), + bias: None, + c: (&c, Strides::row_major(2)), + residual: None, + alpha: 1.0, + act: Act::None, + }; + assert!(g.matmul(&ok).is_ok()); + let too_big = Matmul { m: 2, n: 2, k: 4, ..ok }; + assert!(matches!(g.matmul(&too_big), Err(Error::Input(_)))); + assert!(g.zeros(usize::MAX / 2).is_err()); + g.flush().unwrap(); +} + +// ---------------------------------------------------------------- the other kernels + +fn run_and_read(g: &Gpu, out: &Buf) -> Vec { + g.flush().unwrap(); + g.download(out).unwrap() +} + +#[test] +fn layernorm_and_softmax_match_the_definition() { + let Some(g) = gpu() else { return }; + for (rows, cols) in [(7, 96), (1, 1), (300, 768), (70_000, 8)] { + let x = noise(rows * cols, 7); + let (w, b) = (noise(cols, 8), noise(cols, 9)); + let (gx, gw, gb, gy) = (g.upload(&x), g.upload(&w), g.upload(&b), g.zeros(rows * cols).unwrap()); + g.layernorm(&gx, &gw, &gb, &gy, rows, cols, 1e-6).unwrap(); + let want: Vec = x + .chunks(cols) + .flat_map(|r| { + let mean = r.iter().sum::() / cols as f32; + let var = r.iter().map(|v| (v - mean).powi(2)).sum::() / cols as f32; + r.iter().enumerate().map(|(i, v)| (v - mean) / (var + 1e-6).sqrt() * w[i] + b[i]).collect::>() + }) + .collect(); + assert_close(&run_and_read(&g, &gy), &want, 2e-4, &format!("layernorm {rows}x{cols}")); + } + for (rows, cols) in [(5, 64), (3, 1), (100, 4096), (2, 28736), (70_000, 64)] { + let x = noise(rows * cols, 10).iter().map(|v| v * 6.0).collect::>(); + let gx = g.upload(&x); + g.softmax(&gx, rows, cols).unwrap(); + let want: Vec = x + .chunks(cols) + .flat_map(|r| { + let m = r.iter().cloned().fold(f32::MIN, f32::max); + let e: Vec = r.iter().map(|v| (v - m).exp()).collect(); + let s: f32 = e.iter().sum(); + e.into_iter().map(|v| v / s).collect::>() + }) + .collect(); + assert_close(&run_and_read(&g, &gx), &want, 1e-5, &format!("softmax {rows}x{cols}")); + } +} + +#[test] +fn windows_pooling_and_add_match_naive_loops() { + let Some(g) = gpu() else { return }; + let (h, w, c, ws) = (13, 22, 5, 8); + let x = noise(h * w * c, 11); + let (nwx, nwy) = (w.div_ceil(ws), h.div_ceil(ws)); + let (gx, parts, back) = (g.upload(&x), g.zeros(nwx * nwy * ws * ws * c).unwrap(), g.zeros(h * w * c).unwrap()); + g.window_partition(&gx, &parts, h, w, c, ws).unwrap(); + g.window_unpartition(&parts, &back, h, w, c, ws).unwrap(); + let mut want = vec![0f32; nwx * nwy * ws * ws * c]; + for win in 0..nwx * nwy { + for y in 0..ws { + for xx in 0..ws { + let (gy, gxx) = ((win / nwx) * ws + y, (win % nwx) * ws + xx); + for ch in 0..c { + if gy < h && gxx < w { + want[((win * ws + y) * ws + xx) * c + ch] = x[(gy * w + gxx) * c + ch]; + } + } + } + } + } + g.flush().unwrap(); + assert_close(&g.download(&parts).unwrap(), &want, 0.0, "window_partition"); + assert_close(&g.download(&back).unwrap(), &x, 0.0, "window round trip"); + assert!(g.window_partition(&gx, &parts, h, w, c, 0).is_err()); + + // pooling a q third of a [nb, 8, 6, 3 * c] buffer in place + let (nb, ph, pw, pc) = (3, 8, 6, 4); + let qkv = noise(nb * ph * pw * 3 * pc, 12); + let (gq, out) = (g.upload(&qkv), g.zeros(nb * (ph / 2) * (pw / 2) * pc).unwrap()); + g.maxpool2(&gq, &out, nb, ph, pw, pc, (ph * pw * 3 * pc, pw * 3 * pc, 3 * pc)).unwrap(); + let mut want = vec![]; + for img in 0..nb { + for oy in 0..ph / 2 { + for ox in 0..pw / 2 { + for ch in 0..pc { + let at = |y: usize, xx: usize| qkv[((img * ph + y) * pw + xx) * 3 * pc + ch]; + want.push(at(2 * oy, 2 * ox).max(at(2 * oy, 2 * ox + 1)).max(at(2 * oy + 1, 2 * ox)).max(at(2 * oy + 1, 2 * ox + 1))); + } + } + } + } + assert_close(&run_and_read(&g, &out), &want, 0.0, "maxpool2"); + assert!(g.maxpool2(&gq, &out, nb, 7, pw, pc, (1, 1, 1)).is_err()); + + let (a, b) = (noise(100_001, 13), noise(100_001, 14)); + let (ga, gb, go) = (g.upload(&a), g.upload(&b), g.zeros(a.len()).unwrap()); + g.add(&ga, &gb, &go, a.len()).unwrap(); + let want: Vec = a.iter().zip(&b).map(|(x, y)| x + y).collect(); + assert_close(&run_and_read(&g, &go), &want, 0.0, "add"); +} + +#[test] +fn patch_embed_rope_and_upsample_match_the_cpu_definitions() { + use candle_core::{Device, Tensor}; + let Some(g) = gpu() else { return }; + // the 7 x 7 stride 4 convolution against candle's conv2d + let (h, w, co) = (64usize, 96usize, 8usize); + let (px, wt, bias) = (noise(3 * h * w, 15), noise(co * 147, 16), noise(co, 17)); + let (gp, gw, gb) = (g.upload(&px), g.upload(&wt), g.upload(&bias)); + let (ho, wo) = (h / 4, w / 4); + let out = g.zeros(ho * wo * co).unwrap(); + g.patch_embed(&gp, &gw, &gb, &out, h, w, co).unwrap(); + let t = Tensor::from_vec(px, (1, 3, h, w), &Device::Cpu).unwrap(); + let k = Tensor::from_vec(wt, (co, 3, 7, 7), &Device::Cpu).unwrap(); + let want = + t.conv2d(&k, 3, 4, 1, 1).unwrap().broadcast_add(&Tensor::from_vec(bias, (1, co, 1, 1), &Device::Cpu).unwrap()).unwrap().permute((0, 2, 3, 1)).unwrap(); + assert_close(&run_and_read(&g, &out), &want.flatten_all().unwrap().to_vec1::().unwrap(), 2e-4, "patch_embed"); + + // rotary encoding: pairs rotated with row % table rows, later rows copied + let (rows, cols, table, rotated) = (37usize, 16usize, 10usize, 30usize); + let (x, cs, sn) = (noise(rows * cols, 18), noise(table * cols, 19), noise(table * cols, 20)); + let (gx, gc, gs, go) = (g.upload(&x), g.upload(&cs), g.upload(&sn), g.zeros(rows * cols).unwrap()); + g.rope(&gx, &gc, &gs, &go, rows, cols, table, rotated).unwrap(); + let mut want = x.clone(); + for r in 0..rotated { + for j in 0..cols / 2 { + let (x0, x1) = (x[r * cols + 2 * j], x[r * cols + 2 * j + 1]); + let (t0, t1) = ((r % table) * cols + 2 * j, (r % table) * cols + 2 * j + 1); + want[r * cols + 2 * j] = x0 * cs[t0] - x1 * sn[t0]; + want[r * cols + 2 * j + 1] = x1 * cs[t1] + x0 * sn[t1]; + } + } + assert_close(&run_and_read(&g, &go), &want, 1e-6, "rope"); + assert!(g.rope(&gx, &gc, &gs, &go, rows, 15, table, rotated).is_err()); + + // neck: lateral + nearest upsample + let (h, w, c) = (6usize, 10usize, 7usize); + let (lat, prev) = (noise(h * w * c, 21), noise(h / 2 * (w / 2) * c, 22)); + let (gl, gp2, go) = (g.upload(&lat), g.upload(&prev), g.zeros(h * w * c).unwrap()); + g.upsample_add(&gl, &gp2, &go, h, w, c).unwrap(); + let want: Vec = (0..h * w * c).map(|i| lat[i] + prev[((i / c / w / 2) * (w / 2) + (i / c % w) / 2) * c + i % c]).collect(); + assert_close(&run_and_read(&g, &go), &want, 0.0, "upsample_add"); +} diff --git a/crates/sam2/src/hiera.rs b/crates/sam2/src/hiera.rs index 61becd418..246448e86 100644 --- a/crates/sam2/src/hiera.rs +++ b/crates/sam2/src/hiera.rs @@ -116,6 +116,39 @@ impl Attention { } } +/// What a block is made of: its widths, heads, window (0: global attention) and query pooling. +#[derive(Clone, Copy, Debug, PartialEq)] +pub(crate) struct BlockPlan { + pub dim: usize, + pub dim_out: usize, + pub heads: usize, + pub window: usize, + pub query_stride: Option, +} + +impl BlockPlan { + pub(crate) fn new(cfg: &HieraConfig, stage: usize, block_in_stage: usize, total: usize) -> Result { + let first = stage > 0 && block_in_stage == 0; + let dim = at(&cfg.embed_dims, if first { stage - 1 } else { stage }, "embed dim")?; + let dim_out = at(&cfg.embed_dims, stage, "embed dim")?; + let window = if first { at(&cfg.windows, stage - 1, "window")? } else { at(&cfg.windows, stage, "window")? }; + let window = if cfg.global_blocks.contains(&total) { 0 } else { window }; + let query_stride = (stage > 0 && stage <= cfg.query_pool_stages && block_in_stage == 0).then_some(cfg.query_stride); + Ok(BlockPlan { dim, dim_out, heads: at(&cfg.heads, stage, "heads")?, window, query_stride }) + } + + /// Every block of the backbone, in order. + pub(crate) fn all(cfg: &HieraConfig) -> Result> { + let mut out = vec![]; + for (stage, n) in cfg.blocks.iter().enumerate() { + for b in 0..*n { + out.push(BlockPlan::new(cfg, stage, b, out.len())?); + } + } + Ok(out) + } +} + /// One Hiera block: pre-norm attention (window or global, optionally pooling the queries) and MLP. pub struct Block { norm1: LayerNorm, @@ -130,13 +163,8 @@ pub struct Block { impl Block { fn load(w: &Weights, cfg: &HieraConfig, stage: usize, block_in_stage: usize, total: usize) -> Result { let prefix = format!("vision_encoder.backbone.blocks.{total}"); - let first = stage > 0 && block_in_stage == 0; - let dim = at(&cfg.embed_dims, if first { stage - 1 } else { stage }, "embed dim")?; - let dim_out = at(&cfg.embed_dims, stage, "embed dim")?; - let window = if first { at(&cfg.windows, stage - 1, "window")? } else { at(&cfg.windows, stage, "window")? }; - let window = if cfg.global_blocks.contains(&total) { 0 } else { window }; - let query_stride = (stage > 0 && stage <= cfg.query_pool_stages && block_in_stage == 0).then_some(cfg.query_stride); - let heads = at(&cfg.heads, stage, "heads")?; + let plan = BlockPlan::new(cfg, stage, block_in_stage, total)?; + let (dim, dim_out, heads) = (plan.dim, plan.dim_out, plan.heads); Ok(Block { norm1: layer_norm(w, &format!("{prefix}.layer_norm1"), dim, cfg.eps)?, attn: Attention { @@ -144,13 +172,13 @@ impl Block { proj: linear(w, &format!("{prefix}.attn.proj"), dim_out, dim_out)?, heads, scale: ((dim_out / heads.max(1)) as f64).powf(-0.5), - query_stride, + query_stride: plan.query_stride, }, norm2: layer_norm(w, &format!("{prefix}.layer_norm2"), dim_out, cfg.eps)?, mlp: Mlp::load(w, &format!("{prefix}.mlp"), dim_out, dim_out * cfg.mlp_ratio)?, proj: (dim != dim_out).then(|| linear(w, &format!("{prefix}.proj"), dim, dim_out)).transpose()?, - window, - query_stride, + window: plan.window, + query_stride: plan.query_stride, }) } @@ -269,6 +297,10 @@ impl Hiera { } /// Patch embedding of `[B, 3, H, W]` pixels: `[B, H/4, W/4, C]` (before the positional embedding). + pub(crate) fn config(&self) -> &HieraConfig { + &self.cfg + } + pub fn patch_embed(&self, pixels: &Tensor) -> Result { let x = pixels.conv2d(&self.patch_w, 3, 4, 1, 1)?; let x = x.broadcast_add(&self.patch_b.reshape((1, self.hidden, 1, 1))?)?; @@ -276,7 +308,7 @@ impl Hiera { } /// Background embedding resized to the grid plus the tiled window embedding, `[1, H, W, C]`. - fn pos(&self, h: usize, w: usize) -> Result { + pub(crate) fn pos(&self, h: usize, w: usize) -> Result { let bg = self.cfg.pos_background; let big = bicubic_resize(&self.pos_embed, self.hidden, (bg, bg), (h, w)); let ws = self.window0.max(1); diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index c08ac33b5..59eceddb3 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -12,6 +12,7 @@ use std::fmt; pub mod decoder; +pub mod gpu; pub mod hiera; pub mod mask; pub mod memattn; diff --git a/crates/sam2/src/neck.rs b/crates/sam2/src/neck.rs index 02a253db3..01ab5523b 100644 --- a/crates/sam2/src/neck.rs +++ b/crates/sam2/src/neck.rs @@ -129,7 +129,7 @@ impl Neck { /// Hiera plus neck: `Sam2VisionModel`. pub struct VisionEncoder { - pub hiera: Hiera, + pub hiera: std::sync::Arc, pub neck: Neck, levels: usize, } @@ -144,7 +144,7 @@ pub struct ImageFeatures { impl VisionEncoder { pub fn load_tiny(w: &Weights) -> Result { let cfg = HieraConfig::tiny(); - Ok(VisionEncoder { hiera: Hiera::load(w, &cfg)?, neck: Neck::load(w, &[768, 384, 192, 96], 256, &[2, 3])?, levels: 3 }) + Ok(VisionEncoder { hiera: std::sync::Arc::new(Hiera::load(w, &cfg)?), neck: Neck::load(w, &[768, 384, 192, 96], 256, &[2, 3])?, levels: 3 }) } pub fn forward(&self, pixels: &Tensor) -> Result { diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index 230de7389..59b09b7f1 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -592,3 +592,43 @@ fn the_tracker_stops_at_its_budget() { let r = t.propagate(0, Direction::Forward, None, &Cancel::new(), &mut |_| {}); assert!(matches!(r, Err(Error::Budget(_))), "{r:?}"); } + +/// The GPU encoder (wgpu kernels, any vendor) against the CPU path and PyTorch on the same image. The two +/// float implementations differ by accumulation order and the GPU's erf (1.5e-7), compounded through 12 +/// blocks. Measured (RX 7900 GRE, 2026-10-15): 6.8e-6 against PyTorch, 1.2e-5 against the CPU path. +const GPU_TOL: f32 = 1e-4; + +#[test] +fn gpu_encoder_matches_the_cpu_path_and_pytorch() { + use effectcraft_sam2::gpu::{Gpu, encoder::GpuEncoder}; + use effectcraft_sam2::model::ImageModel; + let Some((w, r)) = load() else { return }; + let gpu = match Gpu::new() { + Ok(g) => std::sync::Arc::new(g), + Err(e) => { + eprintln!("skipped: {e}"); + return; + } + }; + eprintln!("GPU: {}", gpu.adapter_name()); + let model = ImageModel::load(&w).unwrap(); + let pixels = r.get("input.pixel_values").unwrap(); + let cpu = model.encode(pixels).unwrap(); + let enc = GpuEncoder::load(gpu, &w, model.vision.hiera.clone()).unwrap(); + let t = std::time::Instant::now(); + let got = enc.encode(pixels).unwrap(); + eprintln!("GPU encode (first run, includes shader compile): {:.0} ms", t.elapsed().as_secs_f64() * 1e3); + let t = std::time::Instant::now(); + let got = (enc.encode(pixels).unwrap(), got).0; + eprintln!("GPU encode: {:.0} ms", t.elapsed().as_secs_f64() * 1e3); + let low_cpu = cpu.features.features.get(2).unwrap(); + let e = max_abs(&got.low, low_cpu); + eprintln!("low level vs CPU: max abs error {e:e}"); + assert!(e < GPU_TOL, "low level error {e}"); + let e = max_abs(&got.low, r.get("vision_encoder.neck.1").unwrap()); + eprintln!("low level vs PyTorch: max abs error {e:e}"); + assert!(e < GPU_TOL, "low level vs PyTorch error {e}"); + let (e0, e1) = (max_abs(&got.high.0, &cpu.high.0), max_abs(&got.high.1, &cpu.high.1)); + eprintln!("skip features vs CPU: {e0:e} / {e1:e}"); + assert!(e0 < GPU_TOL && e1 < GPU_TOL, "skip feature errors {e0} / {e1}"); +} From ee3498b24df12cffa38c774c317b5c516a97bd9d Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 21:43:18 +0200 Subject: [PATCH 37/77] S2.63: GPU memory attention and memory encoder; one shared device; CPU fallback Memory attention (1.8e-5 vs PyTorch) and the memory encoder (2.6e-6) as wgpu kernels, convolution and layer-norm+GELU kernels, VideoModel::load_auto picks the GPU when there is one (EFFECTCRAFT_SAM2_DEVICE overrides) and falls back to the CPU on any GPU error; disable_gpu switches it off mid-clip. One process-wide device (several at once crashed the Vulkan loader). A tracked frame: 4.3 s CPU, 0.7 s on the RX 7900 GRE. Clips tracked on the GPU match Sam2VideoModel (worst logit error 0.0032). Co-Authored-By: Claude Sonnet 5.5 --- crates/engine/src/cmd/roto_cmds.rs | 2 +- crates/sam2/examples/sam2_bench.rs | 51 ++++++++ crates/sam2/src/gpu/encoder.rs | 3 +- crates/sam2/src/gpu/kernels_ops.wgsl | 37 +++++- crates/sam2/src/gpu/memattn.rs | 184 +++++++++++++++++++++++++++ crates/sam2/src/gpu/memenc.rs | 127 ++++++++++++++++++ crates/sam2/src/gpu/mod.rs | 39 +++++- crates/sam2/src/gpu/ops.rs | 56 +++++++- crates/sam2/src/gpu/tests.rs | 61 +++++++-- crates/sam2/src/mask.rs | 8 +- crates/sam2/src/memattn.rs | 2 +- crates/sam2/src/memory.rs | 16 ++- crates/sam2/src/session.rs | 155 +++++++++++++++++++++- crates/sam2/tests/reference.rs | 130 ++++++++++++++++++- crates/segment/src/lib.rs | 4 + docs/roto-brush-sam2.md | 34 +++-- 16 files changed, 861 insertions(+), 48 deletions(-) create mode 100644 crates/sam2/src/gpu/memattn.rs create mode 100644 crates/sam2/src/gpu/memenc.rs diff --git a/crates/engine/src/cmd/roto_cmds.rs b/crates/engine/src/cmd/roto_cmds.rs index 481a23b5d..5849fa2c8 100644 --- a/crates/engine/src/cmd/roto_cmds.rs +++ b/crates/engine/src/cmd/roto_cmds.rs @@ -478,7 +478,7 @@ fn status(s: &mut Session, p: &Value) -> Result { .keys() .filter_map(|f| chain.trust_at(*f).map(|t| (f.to_string(), json!({"confidence": t.confidence, "occluded": t.occluded})))) .collect(); - out["tracker"] = json!({"model": s.models_roto_active(), "saved": chain.saved_model(), "frames": tracked}); + out["tracker"] = json!({"model": s.models_roto_active(), "device": fx::installed_model().and_then(|m| m.device_name()), "saved": chain.saved_model(), "frames": tracked}); let notice = s.roto_model_notice(s.uses_roto_model(cid, lid, uid)).or_else(|| s.saved_tracking_notice(chain.saved_model())); out["notice"] = json!(notice); out["source"] = json!(if tracked.contains_key(&frame.to_string()) { diff --git a/crates/sam2/examples/sam2_bench.rs b/crates/sam2/examples/sam2_bench.rs index a4e78feb0..0b51dd668 100644 --- a/crates/sam2/examples/sam2_bench.rs +++ b/crates/sam2/examples/sam2_bench.rs @@ -70,6 +70,57 @@ fn run() -> Result<(), Box> { let _ = attn.forward(&cur, &cur_pos, &mem, &mem_pos, 64)?; println!("memory attention (run {run}, {m} memory tokens): {:.0} ms", t.elapsed().as_secs_f64() * 1e3); } + // The same two costs on the GPU (wgpu kernels), when there is one. + match effectcraft_sam2::gpu::Gpu::shared() { + Ok(gpu) => { + println!("GPU: {}", gpu.adapter_name()); + let enc = effectcraft_sam2::gpu::encoder::GpuEncoder::load(gpu.clone(), &weights, model.vision.hiera.clone())?; + let gattn = effectcraft_sam2::gpu::memattn::GpuMemoryAttention::load(gpu, &weights)?; + for run in 0..3 { + let t = Instant::now(); + let _ = enc.encode(&pixels)?; + let encode = t.elapsed().as_secs_f64() * 1e3; + let t = Instant::now(); + let _ = gattn.forward(&cur, &cur_pos, &mem, &mem_pos, 64)?; + println!("GPU run {run}: encode {encode:.0} ms, memory attention (full bank) {:.0} ms", t.elapsed().as_secs_f64() * 1e3); + } + } + Err(e) => println!("no GPU: {e}"), + } + // A whole tracked frame (encode, memory attention, decode, memory encoding) through the video session, on each backend. + struct Same(Tensor); + impl effectcraft_sam2::session::FrameSource for Same { + fn num_frames(&self) -> usize { + 12 + } + fn frame(&self, _i: usize) -> effectcraft_sam2::Result { + Ok(self.0.clone()) + } + } + for backend in [effectcraft_sam2::session::Backend::Cpu, effectcraft_sam2::session::Backend::Gpu] { + let Ok(vm) = effectcraft_sam2::session::VideoModel::load_with(&weights, backend) else { continue }; + let vm = std::sync::Arc::new(vm); + let mut tracker = effectcraft_sam2::session::VideoTracker::new(vm.clone(), Box::new(Same(pixels.clone())))?; + let obj = effectcraft_sam2::session::ObjectId(1); + let t = Instant::now(); + tracker.prompt(0, obj, &Prompt { points: vec![([512.0, 512.0], true)], bbox: None })?; + println!("{:?} ({}): prompt frame {:.0} ms", backend, vm.device_name(), t.elapsed().as_secs_f64() * 1e3); + let mut last = Instant::now(); + let mut times = vec![]; + tracker.propagate(0, effectcraft_sam2::session::Direction::Forward, None, &effectcraft_sam2::session::Cancel::new(), &mut |m| { + if m.frame > 0 { + times.push(last.elapsed().as_secs_f64() * 1e3); + } + last = Instant::now(); + })?; + println!( + "{:?}: tracked frames {:.0} ms each (first {:.0}, last {:.0})", + backend, + times.iter().sum::() / times.len().max(1) as f64, + times.first().copied().unwrap_or(0.0), + times.last().copied().unwrap_or(0.0) + ); + } println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get())); Ok(()) } diff --git a/crates/sam2/src/gpu/encoder.rs b/crates/sam2/src/gpu/encoder.rs index 9a6caeeb4..630f4f09b 100644 --- a/crates/sam2/src/gpu/encoder.rs +++ b/crates/sam2/src/gpu/encoder.rs @@ -107,7 +107,7 @@ impl GpuEncoder { fn norm(&self, x: &Buf, rows: usize, dim: usize, n: &GNorm) -> Result { let out = self.gpu.zeros(rows * dim)?; - self.gpu.layernorm(x, &n.w, &n.b, &out, rows, dim, self.eps)?; + self.gpu.layernorm(x, &n.w, &n.b, &out, rows, dim, self.eps, false)?; Ok(out) } @@ -227,6 +227,7 @@ impl GpuEncoder { /// Encode normalised `[1, 3, H, W]` pixels (`H` and `W` multiples of 32). pub fn encode(&self, pixels: &Tensor) -> Result { let g = &self.gpu; + let _turn = g.exclusive(); let (_, c, ph, pw) = pixels.dims4()?; let window0 = self.cfg.windows.first().copied().unwrap_or(8); if c != 3 || ph == 0 || pw == 0 || ph % 4 != 0 || pw % 4 != 0 || (ph / 4) % window0 != 0 || (pw / 4) % window0 != 0 { diff --git a/crates/sam2/src/gpu/kernels_ops.wgsl b/crates/sam2/src/gpu/kernels_ops.wgsl index aae0c8323..419476600 100644 --- a/crates/sam2/src/gpu/kernels_ops.wgsl +++ b/crates/sam2/src/gpu/kernels_ops.wgsl @@ -29,7 +29,7 @@ fn reduce_max(li: u32) -> f32 { } // ---------------------------------------------------------------- layer norm over the last dimension -// words: 0 rows, 1 cols, 2 eps (f32). in0 x, in1 weight, in2 bias, out0 y. One workgroup per row. +// words: 0 rows, 1 cols, 2 eps (f32), 3 gelu (1: apply GELU to the result). in0 x, in1 weight, in2 bias, out0 y. One workgroup per row. @compute @workgroup_size(64) fn layernorm(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: vec3, @builtin(local_invocation_id) li: vec3) { let row = wg.y * nw.x + wg.x; @@ -44,7 +44,40 @@ fn layernorm(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: for (var c = li.x; c < cols; c += 64u) { let d = in0[base + c] - mean; q += d * d; } red[li.x] = q; let inv = inverseSqrt(reduce_sum(li.x) / f32(cols) + f(2u)); - for (var c = li.x; c < cols; c += 64u) { out0[base + c] = (in0[base + c] - mean) * inv * in1[c] + in2[c]; } + for (var c = li.x; c < cols; c += 64u) { + var y = (in0[base + c] - mean) * inv * in1[c] + in2[c]; + if (u(3u) == 1u) { y = gelu_(y); } + out0[base + c] = y; + } +} + +// ---------------------------------------------------------------- convolution over channels-last [H, W, Ci] +// words: 0 n (output elements), 1 H, 2 W, 3 Ci, 4 Co, 5 Ho, 6 Wo, 7 k, 8 stride, 9 pad, 10 groups. +// in0 input, in1 weight [Co, Ci / groups, k, k], in2 bias [Co]; output [Ho, Wo, Co]. +@compute @workgroup_size(256) +fn conv2d(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: vec3, @builtin(local_invocation_id) li: vec3) { + let i = gid1(wg, nw, li); + if (i >= u(0u)) { return; } + let H = u(1u); let W = u(2u); let Ci = u(3u); let Co = u(4u); let Wo = u(6u); + let k = u(7u); let stride = u(8u); let pad = i32(u(9u)); let groups = u(10u); + let co = i % Co; let t = i / Co; + let ox = t % Wo; let oy = t / Wo; + let cig = Ci / groups; + let g0 = co / (Co / groups); + var acc = in2[co]; + for (var c = 0u; c < cig; c++) { + let ci = g0 * cig + c; + for (var ky = 0u; ky < k; ky++) { + let iy = i32(oy * stride + ky) - pad; + if (iy < 0 || iy >= i32(H)) { continue; } + for (var kx = 0u; kx < k; kx++) { + let ix = i32(ox * stride + kx) - pad; + if (ix < 0 || ix >= i32(W)) { continue; } + acc += in0[(u32(iy) * W + u32(ix)) * Ci + ci] * in1[((co * cig + c) * k + ky) * k + kx]; + } + } + } + out0[i] = acc; } // ---------------------------------------------------------------- softmax over rows, in place on out0 diff --git a/crates/sam2/src/gpu/memattn.rs b/crates/sam2/src/gpu/memattn.rs new file mode 100644 index 000000000..38c5139a7 --- /dev/null +++ b/crates/sam2/src/gpu/memattn.rs @@ -0,0 +1,184 @@ +// Original work for EffectCraft; the maths follows the Hugging Face `transformers` SAM 2 video +// implementation (Apache-2.0, see NOTICE). +//! Memory attention on the GPU: four layers of self-attention over the current frame and cross-attention to +//! the memory bank, with axial rotary position encoding. The attention is one head of 256 channels; the +//! score matrix against a full memory bank is 4096 x 28736 floats, which stays on the GPU. + +use std::sync::Arc; + +use candle_core::{Device, Tensor}; + +use super::{Act, Buf, GLinear, GNorm, Gpu, Matmul, Strides}; +use crate::memattn::{GRID, rope_tables}; +use crate::weights::Weights; +use crate::{Error, Result}; + +const DIM: usize = 256; +const LAYERS: usize = 4; + +struct GAttn { + q: GLinear, + k: GLinear, + v: GLinear, + o: GLinear, +} + +struct GLayer { + self_attn: GAttn, + cross: GAttn, + ln1: GNorm, + ln2: GNorm, + ln3: GNorm, + lin1: GLinear, + lin2: GLinear, +} + +pub struct GpuMemoryAttention { + gpu: Arc, + layers: Vec, + final_ln: GNorm, + cos: Buf, + sin: Buf, +} + +fn attn(g: &Gpu, w: &Weights, prefix: &str, kv: usize) -> Result { + Ok(GAttn { + q: GLinear::load(g, w, &format!("{prefix}.q_proj"), DIM, DIM)?, + k: GLinear::load(g, w, &format!("{prefix}.k_proj"), kv, DIM)?, + v: GLinear::load(g, w, &format!("{prefix}.v_proj"), kv, DIM)?, + o: GLinear::load(g, w, &format!("{prefix}.o_proj"), DIM, DIM)?, + }) +} + +impl GpuMemoryAttention { + pub fn load(gpu: Arc, w: &Weights) -> Result { + let layers = (0..LAYERS) + .map(|i| { + let p = format!("memory_attention.layers.{i}"); + Ok(GLayer { + self_attn: attn(&gpu, w, &format!("{p}.self_attn"), DIM)?, + cross: attn(&gpu, w, &format!("{p}.cross_attn_image"), 64)?, + ln1: GNorm::load(&gpu, w, &format!("{p}.layer_norm1"), DIM)?, + ln2: GNorm::load(&gpu, w, &format!("{p}.layer_norm2"), DIM)?, + ln3: GNorm::load(&gpu, w, &format!("{p}.layer_norm3"), DIM)?, + lin1: GLinear::load(&gpu, w, &format!("{p}.linear1"), DIM, 2048)?, + lin2: GLinear::load(&gpu, w, &format!("{p}.linear2"), 2048, DIM)?, + }) + }) + .collect::>>()?; + let (cos, sin) = rope_tables(10000.0)?; + let flat = |t: Tensor| -> Result> { Ok(t.flatten_all()?.to_vec1::()?) }; + Ok(GpuMemoryAttention { + final_ln: GNorm::load(&gpu, w, "memory_attention.layer_norm", DIM)?, + cos: gpu.upload(&flat(cos)?), + sin: gpu.upload(&flat(sin)?), + layers, + gpu, + }) + } + + fn norm(&self, x: &Buf, rows: usize, n: &GNorm) -> Result { + let out = self.gpu.zeros(rows * DIM)?; + self.gpu.layernorm(x, &n.w, &n.b, &out, rows, DIM, 1e-5, false)?; + Ok(out) + } + + fn rope(&self, x: &Buf, rows: usize, rotated: usize) -> Result { + let out = self.gpu.zeros(rows * DIM)?; + self.gpu.rope(x, &self.cos, &self.sin, &out, rows, DIM, GRID * GRID, rotated)?; + Ok(out) + } + + /// `softmax(q . k^T / 16) . v` for `[nq, 256]` queries and `[nk, 256]` keys and values, through `scores` + /// (at least `nq * nk` elements); returns `[nq, 256]`. + fn attention(&self, q: &Buf, k: &Buf, v: &Buf, nq: usize, nk: usize, scores: &Buf) -> Result { + let g = &self.gpu; + g.matmul(&Matmul { + m: nq, + n: nk, + k: DIM, + batch: 1, + a: (q, Strides::row_major(DIM)), + b: (k, Strides { offset: 0, batch: 0, row: 1, col: DIM }), + bias: None, + c: (scores, Strides::row_major(nk)), + residual: None, + alpha: (DIM as f32).powf(-0.5), + act: Act::None, + })?; + g.softmax(scores, nq, nk)?; + let out = g.zeros(nq * DIM)?; + g.matmul(&Matmul { + m: nq, + n: DIM, + k: nk, + batch: 1, + a: (scores, Strides::row_major(nk)), + b: (v, Strides::row_major(DIM)), + bias: None, + c: (&out, Strides::row_major(DIM)), + residual: None, + alpha: 1.0, + act: Act::None, + })?; + Ok(out) + } + + /// As [`crate::memattn::MemoryAttention::forward`]: `current` and `current_pos` `[4096, 256]`, `memory` and + /// `memory_pos` `[M, 64]`, the last `num_pointers` memory rows being object pointers (not rotated). + pub fn forward(&self, current: &Tensor, current_pos: &Tensor, memory: &Tensor, memory_pos: &Tensor, num_pointers: usize) -> Result { + let g = &self.gpu; + let _turn = g.exclusive(); + let n = GRID * GRID; + if current.dims() != [n, DIM] || current_pos.dims() != [n, DIM] { + return Err(Error::Input(format!("current features must be [{n}, {DIM}], got {:?}", current.dims()))); + } + let (m, md) = memory.dims2()?; + if md != 64 || memory_pos.dims() != [m, 64] { + return Err(Error::Input(format!("memory must be [M, 64] with a matching position encoding, got {:?}", memory.dims()))); + } + let rotated = m.checked_sub(num_pointers).ok_or_else(|| Error::Input("more object pointers than memory tokens".into()))?; + if rotated % n != 0 { + return Err(Error::Input(format!("memory length {rotated} is not a multiple of {n}"))); + } + let flat = |t: &Tensor| -> Result> { Ok(t.contiguous()?.flatten_all()?.to_vec1::()?) }; + let mut x = g.upload(&flat(&(current + (current_pos * 0.1)?)?)?); + let mem = g.upload(&flat(memory)?); + let keys_in = g.upload(&flat(&(memory + memory_pos)?)?); + let scores = g.zeros(n * n.max(m))?; + for l in &self.layers { + // self-attention over the frame + let q_in = self.norm(&x, n, &l.ln1)?; + let (q, k, v) = (g.zeros(n * DIM)?, g.zeros(n * DIM)?, g.zeros(n * DIM)?); + g.linear(&q_in, n, &l.self_attn.q, &q, Act::None, None)?; + g.linear(&q_in, n, &l.self_attn.k, &k, Act::None, None)?; + g.linear(&q_in, n, &l.self_attn.v, &v, Act::None, None)?; + let (q, k) = (self.rope(&q, n, n)?, self.rope(&k, n, n)?); + let a = self.attention(&q, &k, &v, n, n, &scores)?; + let x1 = g.zeros(n * DIM)?; + g.linear(&a, n, &l.self_attn.o, &x1, Act::None, Some(&x))?; + // cross-attention to the memory bank + let q_in = self.norm(&x1, n, &l.ln2)?; + let q = g.zeros(n * DIM)?; + g.linear(&q_in, n, &l.cross.q, &q, Act::None, None)?; + let q = self.rope(&q, n, n)?; + let (k, v) = (g.zeros(m * DIM)?, g.zeros(m * DIM)?); + g.linear(&keys_in, m, &l.cross.k, &k, Act::None, None)?; + g.linear(&mem, m, &l.cross.v, &v, Act::None, None)?; + let k = self.rope(&k, m, rotated)?; + let a = self.attention(&q, &k, &v, n, m, &scores)?; + let x2 = g.zeros(n * DIM)?; + g.linear(&a, n, &l.cross.o, &x2, Act::None, Some(&x1))?; + // feed-forward + let f_in = self.norm(&x2, n, &l.ln3)?; + let hidden = g.zeros(n * 2048)?; + g.linear(&f_in, n, &l.lin1, &hidden, Act::Relu, None)?; + let x3 = g.zeros(n * DIM)?; + g.linear(&hidden, n, &l.lin2, &x3, Act::None, Some(&x2))?; + x = x3; + } + let out = self.norm(&x, n, &self.final_ln)?; + g.flush()?; + Ok(Tensor::from_vec(g.download(&out)?, (n, DIM), &Device::Cpu)?) + } +} diff --git a/crates/sam2/src/gpu/memenc.rs b/crates/sam2/src/gpu/memenc.rs new file mode 100644 index 000000000..2cfa693d0 --- /dev/null +++ b/crates/sam2/src/gpu/memenc.rs @@ -0,0 +1,127 @@ +// Original work for EffectCraft; the maths follows the Hugging Face `transformers` SAM 2 video +// implementation (Apache-2.0, see NOTICE). +//! The memory encoder on the GPU: the mask is down-sampled by four strided convolutions, fused with the +//! frame's features by two ConvNeXt-style blocks, and projected to the 64 channels of a memory. + +use std::sync::Arc; + +use candle_core::{Device, Tensor}; + +use super::{Act, Buf, GLinear, GNorm, Gpu}; +use crate::weights::Weights; +use crate::{Error, Result}; + +struct GConv { + w: Buf, + b: Buf, + ci: usize, + co: usize, + k: usize, +} + +struct GFuser { + dw: GConv, + norm: GNorm, + pw1: GLinear, + /// The layer scale is folded into this layer's weights. + pw2: GLinear, +} + +pub struct GpuMemoryEncoder { + gpu: Arc, + down: Vec<(GConv, GNorm)>, + final_conv: GLinear, + feature_projection: GLinear, + fuser: Vec, + projection: GLinear, +} + +fn conv(g: &Gpu, w: &Weights, name: &str, ci: usize, co: usize, k: usize, groups: usize) -> Result { + let flat = |t: Tensor| -> Result> { Ok(t.flatten_all()?.to_vec1::()?) }; + Ok(GConv { + w: g.upload(&flat(w.get(&format!("{name}.weight"), &[co, ci / groups, k, k])?)?), + b: g.upload(&flat(w.get(&format!("{name}.bias"), &[co])?)?), + ci, + co, + k, + }) +} + +impl GpuMemoryEncoder { + pub fn load(gpu: Arc, w: &Weights) -> Result { + let mut down = vec![]; + let mut ch = 1; + for i in 0..4 { + let out = ch * 4; + let p = format!("memory_encoder.mask_downsampler.layers.{i}"); + down.push((conv(&gpu, w, &format!("{p}.conv"), ch, out, 3, 1)?, GNorm::load(&gpu, w, &format!("{p}.layer_norm"), out)?)); + ch = out; + } + let fuser = (0..2) + .map(|i| { + let p = format!("memory_encoder.memory_fuser.layers.{i}"); + let scale = w.get(&format!("{p}.scale"), &[256])?.flatten_all()?.to_vec1::()?; + Ok(GFuser { + dw: conv(&gpu, w, &format!("{p}.depthwise_conv"), 256, 256, 7, 256)?, + norm: GNorm::load(&gpu, w, &format!("{p}.layer_norm"), 256)?, + pw1: GLinear::load(&gpu, w, &format!("{p}.pointwise_conv1"), 256, 1024)?, + pw2: GLinear::load_scaled(&gpu, w, &format!("{p}.pointwise_conv2"), 1024, 256, &scale)?, + }) + }) + .collect::>>()?; + Ok(GpuMemoryEncoder { + final_conv: GLinear::load_conv1x1(&gpu, w, "memory_encoder.mask_downsampler.final_conv", 256, 256)?, + feature_projection: GLinear::load_conv1x1(&gpu, w, "memory_encoder.feature_projection", 256, 256)?, + projection: GLinear::load_conv1x1(&gpu, w, "memory_encoder.projection", 256, 64)?, + down, + fuser, + gpu, + }) + } + + /// `features`: `[1, 256, 64, 64]`; `mask`: `[1, 1, 1024, 1024]` already scaled (`20 m - 10`). Returns the raw + /// memory tokens `[4096, 64]` (the occlusion embedding and bfloat16 rounding are the caller's). + pub fn encode(&self, features: &Tensor, mask: &Tensor) -> Result { + let g = &self.gpu; + let _turn = g.exclusive(); + let flat = |t: &Tensor| -> Result> { Ok(t.contiguous()?.flatten_all()?.to_vec1::()?) }; + if features.dims() != [1, 256, 64, 64] || mask.dims() != [1, 1, 1024, 1024] { + return Err(Error::Input(format!("memory encoder input shapes {:?} / {:?}", features.dims(), mask.dims()))); + } + let tokens = features.permute((0, 2, 3, 1))?.contiguous()?; + let ftok = g.upload(&flat(&tokens)?); + let mut cur = g.upload(&flat(mask)?); + let (mut h, mut w) = (1024usize, 1024usize); + for (c, ln) in &self.down { + let (ho, wo) = (h / 2, w / 2); + let raw = g.zeros(ho * wo * c.co)?; + g.conv2d(&cur, &c.w, &c.b, &raw, (h, w, c.ci, c.co), c.k, 2, 1, 1)?; + let normed = g.zeros(ho * wo * c.co)?; + g.layernorm(&raw, &ln.w, &ln.b, &normed, ho * wo, c.co, 1e-6, true)?; + (cur, h, w) = (normed, ho, wo); + } + let n = h * w; + if n != 4096 { + return Err(Error::Compute(format!("the mask down-sampler ended at {w}x{h}, not 64x64"))); + } + let m = g.zeros(n * 256)?; + g.linear(&cur, n, &self.final_conv, &m, Act::None, None)?; + let mut x = g.zeros(n * 256)?; + g.linear(&ftok, n, &self.feature_projection, &x, Act::None, Some(&m))?; + for b in &self.fuser { + let d = g.zeros(n * 256)?; + g.conv2d(&x, &b.dw.w, &b.dw.b, &d, (h, w, 256, 256), 7, 1, 3, 256)?; + let normed = g.zeros(n * 256)?; + g.layernorm(&d, &b.norm.w, &b.norm.b, &normed, n, 256, 1e-6, false)?; + let hidden = g.zeros(n * 1024)?; + g.linear(&normed, n, &b.pw1, &hidden, Act::Gelu, None)?; + let next = g.zeros(n * 256)?; + g.linear(&hidden, n, &b.pw2, &next, Act::None, Some(&x))?; + x = next; + } + let out = g.zeros(n * 64)?; + g.linear(&x, n, &self.projection, &out, Act::None, None)?; + g.flush()?; + Ok(Tensor::from_vec(g.download(&out)?, (n, 64), &Device::Cpu)?) + } +} diff --git a/crates/sam2/src/gpu/mod.rs b/crates/sam2/src/gpu/mod.rs index a14df4b16..9ca180932 100644 --- a/crates/sam2/src/gpu/mod.rs +++ b/crates/sam2/src/gpu/mod.rs @@ -72,6 +72,7 @@ pub struct Gpu { module: wgpu::ShaderModule, pipelines: Mutex>>, encoder: Mutex>, + turn: Mutex<()>, error: Arc>>, adapter: String, } @@ -89,8 +90,16 @@ fn u32_of(v: usize, what: &str) -> Result { } impl Gpu { + /// The process's one GPU device, created on first use. Several devices at once upset some Vulkan loaders + /// (a crash when two threads create or drop them together), and one is all SAM 2 needs: models come and go, + /// the device stays. An error (no GPU) is remembered too. + pub fn shared() -> Result> { + static DEVICE: std::sync::OnceLock, String>> = std::sync::OnceLock::new(); + DEVICE.get_or_init(|| Gpu::new().map(Arc::new).map_err(|e| e.to_string())).clone().map_err(Error::Compute) + } + /// A device on the best real GPU (a software adapter is refused: the CPU path is faster). - pub fn new() -> Result { + fn new() -> Result { let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle()); let adapter = pollster::block_on( instance.request_adapter(&wgpu::RequestAdapterOptions { power_preference: wgpu::PowerPreference::HighPerformance, ..Default::default() }), @@ -117,7 +126,13 @@ impl Gpu { sink.lock().unwrap_or_else(PoisonError::into_inner).get_or_insert_with(|| format!("GPU device lost ({reason:?}): {message}")); }); let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("sam2 kernels"), source: wgpu::ShaderSource::Wgsl(SHADER.into()) }); - Ok(Gpu { device, queue, module, pipelines: Mutex::default(), encoder: Mutex::new(None), error, adapter: info.name }) + Ok(Gpu { device, queue, module, pipelines: Mutex::default(), encoder: Mutex::new(None), turn: Mutex::new(()), error, adapter: info.name }) + } + + /// One GPU pass at a time: a caller holds this while it queues, flushes and reads back a whole job, so + /// two threads never interleave their dispatches in one command stream. + pub fn exclusive(&self) -> std::sync::MutexGuard<'_, ()> { + self.turn.lock().unwrap_or_else(PoisonError::into_inner) } pub fn adapter_name(&self) -> &str { @@ -304,6 +319,8 @@ impl Gpu { } pub mod encoder; +pub mod memattn; +pub mod memenc; mod ops; /// A linear layer's weight `[out, in]` and bias `[out]` on the GPU. @@ -335,6 +352,24 @@ impl GLinear { }) } + /// As [`GLinear::load`], with output channel `j` scaled by `scale[j]` (the memory fuser's layer scale). + pub fn load_scaled(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize, scale: &[f32]) -> Result { + if scale.len() != out { + return Err(Error::Weights(format!("`{name}`: {} scales for {out} outputs", scale.len()))); + } + let mut wt = flat(w.get(&format!("{name}.weight"), &[out, inp])?)?; + let mut bias = flat(w.get(&format!("{name}.bias"), &[out])?)?; + for (j, s) in scale.iter().enumerate() { + if let Some(row) = wt.get_mut(j * inp..(j + 1) * inp) { + row.iter_mut().for_each(|v| *v *= s); + } + if let Some(b) = bias.get_mut(j) { + *b *= s; + } + } + Ok(GLinear { w: g.upload(&wt), b: g.upload(&bias), inp, out }) + } + /// A 1 x 1 convolution (`[out, in, 1, 1]`) as a linear layer over channels-last tokens. pub fn load_conv1x1(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize) -> Result { Ok(GLinear { diff --git a/crates/sam2/src/gpu/ops.rs b/crates/sam2/src/gpu/ops.rs index 7bb99c495..3d9b8f184 100644 --- a/crates/sam2/src/gpu/ops.rs +++ b/crates/sam2/src/gpu/ops.rs @@ -27,12 +27,18 @@ fn need(buf: &Buf, len: usize, what: &str) -> Result<()> { impl Gpu { /// `y = layer_norm(x) * weight + bias` over the last dimension of a `[rows, cols]` buffer. - pub fn layernorm(&self, x: &Buf, weight: &Buf, bias: &Buf, y: &Buf, rows: usize, cols: usize, eps: f32) -> Result<()> { + #[allow(clippy::too_many_arguments)] + pub fn layernorm(&self, x: &Buf, weight: &Buf, bias: &Buf, y: &Buf, rows: usize, cols: usize, eps: f32, gelu: bool) -> Result<()> { need(x, rows * cols, "layernorm x")?; need(y, rows * cols, "layernorm y")?; need(weight, cols, "layernorm weight")?; need(bias, cols, "layernorm bias")?; - self.run("layernorm", &[(0, x), (1, weight), (2, bias), (3, y)], &[n32(rows, "rows")?, n32(cols, "cols")?, eps.to_bits()], grid(rows)?); + self.run( + "layernorm", + &[(0, x), (1, weight), (2, bias), (3, y)], + &[n32(rows, "rows")?, n32(cols, "cols")?, eps.to_bits(), u32::from(gelu)], + grid(rows)?, + ); Ok(()) } @@ -43,6 +49,52 @@ impl Gpu { Ok(()) } + /// A convolution over channels-last `[h, w, ci]` with weight `[co, ci / groups, k, k]` and bias `[co]`, + /// output `[ho, wo, co]` (`groups == ci == co` is depthwise). + #[allow(clippy::too_many_arguments)] + pub fn conv2d( + &self, + x: &Buf, + weight: &Buf, + bias: &Buf, + out: &Buf, + dims: (usize, usize, usize, usize), + k: usize, + stride: usize, + pad: usize, + groups: usize, + ) -> Result<()> { + let (h, w, ci, co) = dims; + if k == 0 || stride == 0 || groups == 0 || ci % groups != 0 || co % groups != 0 || h + 2 * pad < k || w + 2 * pad < k { + return Err(Error::Input(format!("a {k}x{k} convolution (stride {stride}, groups {groups}) does not fit {w}x{h}x{ci} -> {co}"))); + } + let (ho, wo) = ((h + 2 * pad - k) / stride + 1, (w + 2 * pad - k) / stride + 1); + need(x, h * w * ci, "conv2d input")?; + need(weight, co * (ci / groups) * k * k, "conv2d weight")?; + need(bias, co, "conv2d bias")?; + need(out, ho * wo * co, "conv2d output")?; + let n = ho * wo * co; + self.run( + "conv2d", + &[(0, x), (1, weight), (2, bias), (3, out)], + &[ + n32(n, "n")?, + n32(h, "h")?, + n32(w, "w")?, + n32(ci, "ci")?, + n32(co, "co")?, + n32(ho, "ho")?, + n32(wo, "wo")?, + n32(k, "k")?, + n32(stride, "stride")?, + n32(pad, "pad")?, + n32(groups, "groups")?, + ], + items(n)?, + ); + Ok(()) + } + /// `out = a + b` over `n` elements. pub fn add(&self, a: &Buf, b: &Buf, out: &Buf, n: usize) -> Result<()> { need(a, n, "add a")?; diff --git a/crates/sam2/src/gpu/tests.rs b/crates/sam2/src/gpu/tests.rs index 92ad20f41..ed5b3528a 100644 --- a/crates/sam2/src/gpu/tests.rs +++ b/crates/sam2/src/gpu/tests.rs @@ -6,27 +6,25 @@ use super::*; /// The one device the tests share (several devices at once upset the Vulkan loader), held for the whole /// test so the tests take turns. -pub(crate) struct Shared(#[allow(dead_code)] std::sync::MutexGuard<'static, ()>, &'static Gpu); +pub(crate) struct Shared(#[allow(dead_code)] std::sync::MutexGuard<'static, ()>, Arc); impl std::ops::Deref for Shared { type Target = Gpu; fn deref(&self) -> &Gpu { - self.1 + &self.1 } } pub(crate) fn gpu() -> Option { static TURN: Mutex<()> = Mutex::new(()); - static DEVICE: std::sync::OnceLock> = std::sync::OnceLock::new(); let guard = TURN.lock().unwrap_or_else(PoisonError::into_inner); - let device = DEVICE.get_or_init(|| match Gpu::new() { - Ok(g) => Some(g), + match Gpu::shared() { + Ok(g) => Some(Shared(guard, g)), Err(e) => { eprintln!("skipped: {e}"); None } - }); - device.as_ref().map(|g| Shared(guard, g)) + } } /// Deterministic pseudo-random values in -1..1. @@ -176,7 +174,7 @@ fn layernorm_and_softmax_match_the_definition() { let x = noise(rows * cols, 7); let (w, b) = (noise(cols, 8), noise(cols, 9)); let (gx, gw, gb, gy) = (g.upload(&x), g.upload(&w), g.upload(&b), g.zeros(rows * cols).unwrap()); - g.layernorm(&gx, &gw, &gb, &gy, rows, cols, 1e-6).unwrap(); + g.layernorm(&gx, &gw, &gb, &gy, rows, cols, 1e-6, false).unwrap(); let want: Vec = x .chunks(cols) .flat_map(|r| { @@ -299,3 +297,50 @@ fn patch_embed_rope_and_upsample_match_the_cpu_definitions() { let want: Vec = (0..h * w * c).map(|i| lat[i] + prev[((i / c / w / 2) * (w / 2) + (i / c % w) / 2) * c + i % c]).collect(); assert_close(&run_and_read(&g, &go), &want, 0.0, "upsample_add"); } + +#[test] +fn convolutions_and_the_layernorm_gelu_epilogue_match_candle() { + use candle_core::{Device, Tensor}; + let Some(g) = gpu() else { return }; + // (h, w, ci, co, k, stride, pad, groups): a strided 3 x 3, a depthwise 7 x 7, and a 1-channel input + for (h, w, ci, co, k, stride, pad, groups) in + [(9usize, 13usize, 3usize, 5usize, 3usize, 2usize, 1usize, 1usize), (10, 10, 4, 4, 7, 1, 3, 4), (16, 16, 1, 4, 3, 2, 1, 1)] + { + let (x, wt, b) = (noise(h * w * ci, 30), noise(co * (ci / groups) * k * k, 31), noise(co, 32)); + let (gx, gw, gb) = (g.upload(&x), g.upload(&wt), g.upload(&b)); + let (ho, wo) = ((h + 2 * pad - k) / stride + 1, (w + 2 * pad - k) / stride + 1); + let out = g.zeros(ho * wo * co).unwrap(); + g.conv2d(&gx, &gw, &gb, &out, (h, w, ci, co), k, stride, pad, groups).unwrap(); + // candle works on NCHW: bring the channels-last input over, convolve, bring the result back + let t = Tensor::from_vec(x, (1, h, w, ci), &Device::Cpu).unwrap().permute((0, 3, 1, 2)).unwrap().contiguous().unwrap(); + let kern = Tensor::from_vec(wt, (co, ci / groups, k, k), &Device::Cpu).unwrap(); + let want = t + .conv2d(&kern, pad, stride, 1, groups) + .unwrap() + .broadcast_add(&Tensor::from_vec(b, (1, co, 1, 1), &Device::Cpu).unwrap()) + .unwrap() + .permute((0, 2, 3, 1)) + .unwrap(); + assert_close( + &run_and_read(&g, &out), + &want.flatten_all().unwrap().to_vec1::().unwrap(), + 2e-4, + &format!("conv {k}x{k} stride {stride} groups {groups}"), + ); + } + assert!(g.conv2d(&g.upload(&[0.0; 4]), &g.upload(&[0.0; 4]), &g.upload(&[0.0; 1]), &g.zeros(4).unwrap(), (2, 2, 1, 1), 5, 1, 0, 1).is_err()); + // layer norm followed by GELU in the same pass + let (rows, cols) = (50usize, 96usize); + let (x, wt, b) = (noise(rows * cols, 33), noise(cols, 34), noise(cols, 35)); + let (gx, gw, gb, gy) = (g.upload(&x), g.upload(&wt), g.upload(&b), g.zeros(rows * cols).unwrap()); + g.layernorm(&gx, &gw, &gb, &gy, rows, cols, 1e-6, true).unwrap(); + let want: Vec = x + .chunks(cols) + .flat_map(|r| { + let mean = r.iter().sum::() / cols as f32; + let var = r.iter().map(|v| (v - mean).powi(2)).sum::() / cols as f32; + r.iter().enumerate().map(|(i, v)| gelu((v - mean) / (var + 1e-6).sqrt() * wt[i] + b[i])).collect::>() + }) + .collect(); + assert_close(&run_and_read(&g, &gy), &want, 1e-3, "layernorm + gelu"); +} diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index f917e29ef..94cf8b40a 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -79,7 +79,7 @@ impl Sam2Mask { { return Ok(e.clone()); } - let enc = Arc::new(self.model.image.encode(&preprocess(rgb, w, h)?)?); + let enc = Arc::new(self.model.encode(&preprocess(rgb, w, h)?)?); *self.last.lock().unwrap_or_else(PoisonError::into_inner) = Some((key, enc.clone())); Ok(enc) } @@ -117,6 +117,10 @@ impl MaskModel for Sam2Mask { Some(self) } + fn device_name(&self) -> Option { + Some(self.model.device_name()) + } + fn idle_secs(&self) -> Option { Some(self.used.lock().unwrap_or_else(PoisonError::into_inner).elapsed().as_secs_f64()) } @@ -129,7 +133,7 @@ pub fn load(bytes: &[u8]) -> std::result::Result, String> { 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())?; + let model = VideoModel::load_auto(&w).map_err(|e| e.to_string())?; Ok(Arc::new(Sam2Mask::new(Arc::new(model)))) } diff --git a/crates/sam2/src/memattn.rs b/crates/sam2/src/memattn.rs index c3bfc0893..d455f1ed1 100644 --- a/crates/sam2/src/memattn.rs +++ b/crates/sam2/src/memattn.rs @@ -18,7 +18,7 @@ const LAYERS: usize = 4; pub const GRID: usize = 64; /// `cos` and `sin` of the axial rotary encoding for the `GRID x GRID` grid, each `[GRID², 256]`. -fn rope_tables(theta: f32) -> Result<(Tensor, Tensor)> { +pub(crate) fn rope_tables(theta: f32) -> Result<(Tensor, Tensor)> { let freqs = DIM / 2 / 2; // per axis: head_dim / 4 let inv: Vec = (0..freqs).map(|i| 1.0 / theta.powf(i as f32 / freqs as f32)).collect(); let n = GRID * GRID; diff --git a/crates/sam2/src/memory.rs b/crates/sam2/src/memory.rs index 45b470a3f..7a73fd5a3 100644 --- a/crates/sam2/src/memory.rs +++ b/crates/sam2/src/memory.rs @@ -123,10 +123,16 @@ impl MemoryEncoder { /// The memory kept for a frame: the encoder output, plus the occlusion embedding when the /// object is absent, rounded through bfloat16 like the reference. pub fn memory(&self, features: &Tensor, mask: &Tensor, object_present: bool) -> Result { - let (x, pos) = self.encode(features, mask)?; - let x = if object_present { x } else { x.broadcast_add(&self.occlusion.reshape((1, 64, 1, 1))?)? }; - let flat = |t: &Tensor| -> Result { Ok(t.flatten_from(2)?.squeeze(0)?.t()?.contiguous()?) }; - let features = flat(&x)?.to_dtype(candle_core::DType::BF16)?.to_dtype(candle_core::DType::F32)?; - Ok(Memory { features, pos: flat(&pos)? }) + let (x, _) = self.encode(features, mask)?; + self.finish(&x.flatten_from(2)?.squeeze(0)?.t()?.contiguous()?, object_present) + } + + /// A memory from the encoder's raw tokens `[HW, 64]`: the occlusion embedding when the object is absent, + /// rounded through bfloat16, with the position encoding (shared by the CPU and GPU encoders). + pub fn finish(&self, tokens: &Tensor, object_present: bool) -> Result { + let tokens = if object_present { tokens.clone() } else { tokens.broadcast_add(&self.occlusion)? }; + let features = tokens.to_dtype(candle_core::DType::BF16)?.to_dtype(candle_core::DType::F32)?; + let pos = sine_position(32, 64, 64)?; + Ok(Memory { features, pos: pos.flatten_from(2)?.squeeze(0)?.t()?.contiguous()? }) } } diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 2e45cc407..9de8fedc2 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -14,9 +14,14 @@ use candle_core::{DType, Tensor}; use candle_nn::{Linear, ops::sigmoid}; use crate::decoder::FeedForward; +use crate::gpu::Gpu; +use crate::gpu::encoder::GpuEncoder; +use crate::gpu::memattn::GpuMemoryAttention; +use crate::gpu::memenc::GpuMemoryEncoder; use crate::memattn::MemoryAttention; use crate::memory::{Memory, MemoryEncoder}; -use crate::model::ImageModel; +use crate::model::{Encoded, ImageModel}; +use crate::neck::ImageFeatures; use crate::nn::{apply, linear}; use crate::prompt::{INPUT_SIZE, Prompt}; use crate::resize::antialias_bilinear; @@ -79,8 +84,31 @@ pub struct FrameMask { pub occluded: bool, } +/// Which device runs the heavy parts (the backbone, the neck and the memory attention). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Backend { + /// Candle on the CPU: the reference path. + Cpu, + /// wgpu compute on the GPU (any vendor); an error if there is none. + Gpu, + /// The GPU when there is a real one and it loads, else the CPU. The environment variable + /// `EFFECTCRAFT_SAM2_DEVICE` (`cpu`, `gpu` or `auto`) overrides the choice. + Auto, +} + +/// The GPU path: kernels for the backbone and the memory attention. Any failure switches it off for good +/// and the CPU path takes over. +struct GpuParts { + encoder: GpuEncoder, + attention: GpuMemoryAttention, + memory_encoder: GpuMemoryEncoder, + ok: AtomicBool, + adapter: String, +} + /// The weights of everything the video tracker needs. pub struct VideoModel { + gpu: Option, pub image: ImageModel, memory_encoder: MemoryEncoder, memory_attention: MemoryAttention, @@ -93,9 +121,37 @@ pub struct VideoModel { } impl VideoModel { + /// On the CPU: the reference path every test compares against. pub fn load(w: &Weights) -> Result { + VideoModel::load_with(w, Backend::Cpu) + } + + /// On the GPU when there is one (see [`Backend::Auto`]). + pub fn load_auto(w: &Weights) -> Result { + VideoModel::load_with(w, Backend::Auto) + } + + pub fn load_with(w: &Weights, backend: Backend) -> Result { + let backend = match (backend, std::env::var("EFFECTCRAFT_SAM2_DEVICE").ok().as_deref()) { + (Backend::Auto, Some("cpu")) => Backend::Cpu, + (Backend::Auto, Some("gpu")) => Backend::Gpu, + (b, _) => b, + }; + let image = ImageModel::load(w)?; + let gpu = match backend { + Backend::Cpu => None, + Backend::Gpu => Some(Self::load_gpu(w, &image)?), + Backend::Auto => match Self::load_gpu(w, &image) { + Ok(g) => Some(g), + Err(e) => { + log::info!("SAM 2 runs on the CPU: {e}"); + None + } + }, + }; Ok(VideoModel { - image: ImageModel::load(w)?, + gpu, + image, memory_encoder: MemoryEncoder::load(w)?, memory_attention: MemoryAttention::load(w)?, temporal_pos: w.get("memory_temporal_positional_encoding", &[MEMORY_FRAMES, 1, 1, 64])?.flatten_all()?.to_vec1::()?, @@ -106,6 +162,79 @@ impl VideoModel { }) } + fn load_gpu(w: &Weights, image: &ImageModel) -> Result { + let gpu = Gpu::shared()?; + Ok(GpuParts { + encoder: GpuEncoder::load(gpu.clone(), w, image.vision.hiera.clone())?, + attention: GpuMemoryAttention::load(gpu.clone(), w)?, + memory_encoder: GpuMemoryEncoder::load(gpu.clone(), w)?, + ok: AtomicBool::new(true), + adapter: gpu.adapter_name().to_string(), + }) + } + + /// The device the heavy parts run on now: the GPU's name, or "CPU". + pub fn device_name(&self) -> String { + match &self.gpu { + Some(g) if g.ok.load(Ordering::Relaxed) => g.adapter.clone(), + _ => "CPU".to_string(), + } + } + + /// Run on the CPU from now on (a GPU that misbehaves, or the user's choice). Weights and results are + /// unchanged; only the device differs. + pub fn disable_gpu(&self) { + if let Some(g) = &self.gpu { + g.ok.store(false, Ordering::Relaxed); + } + } + + fn gpu_failed(&self, what: &str, e: &Error) { + if let Some(g) = &self.gpu + && g.ok.swap(false, Ordering::Relaxed) + { + log::warn!("SAM 2 {what} failed on the GPU ({e}); the CPU takes over"); + } + } + + /// Encode a normalised `[1, 3, 1024, 1024]` frame on the GPU if it is working, else on the CPU. + pub fn encode(&self, pixels: &Tensor) -> Result { + if let Some(g) = self.gpu.as_ref().filter(|g| g.ok.load(Ordering::Relaxed)) { + 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 features = ImageFeatures { features: vec![none()?, none()?, f.low], positions: vec![none()?, none()?, pos] }; + return Ok(Encoded { features, high: f.high }); + } + Err(e) => self.gpu_failed("encoding", &e), + } + } + self.image.encode(pixels) + } + + /// The memory of a frame (its features and the scaled mask), on the GPU if it is working, else on the CPU. + fn memorize(&self, features: &Tensor, mask: &Tensor, present: bool) -> Result { + if let Some(g) = self.gpu.as_ref().filter(|g| g.ok.load(Ordering::Relaxed)) { + match g.memory_encoder.encode(features, mask).and_then(|t| self.memory_encoder.finish(&t, present)) { + Ok(m) => return Ok(m), + Err(e) => self.gpu_failed("memory encoding", &e), + } + } + self.memory_encoder.memory(features, mask, present) + } + + /// Memory attention on the GPU if it is working, else on the CPU. + fn attend(&self, current: &Tensor, current_pos: &Tensor, memory: &Tensor, memory_pos: &Tensor, pointers: usize) -> Result { + if let Some(g) = self.gpu.as_ref().filter(|g| g.ok.load(Ordering::Relaxed)) { + match g.attention.forward(current, current_pos, memory, memory_pos, pointers) { + Ok(t) => return Ok(t), + Err(e) => self.gpu_failed("memory attention", &e), + } + } + self.memory_attention.forward(current, current_pos, memory, memory_pos, pointers) + } + pub fn open(path: &Path) -> Result { VideoModel::load(&Weights::open(path)?) } @@ -221,7 +350,7 @@ impl<'a> VideoTracker<'a> { return Ok(f.clone()); } let pixels = self.frames.frame(frame)?; - let enc = self.model.image.encode(&pixels)?; + let enc = self.model.encode(&pixels)?; let (low, pos) = match (enc.features.features.get(2), enc.features.positions.get(2)) { (Some(l), Some(p)) => (l.clone(), p.clone()), _ => return Err(Error::Compute("vision encoder returned too few levels".into())), @@ -295,7 +424,17 @@ impl<'a> VideoTracker<'a> { /// Run the model on one frame for one object and store the result. fn step(&mut self, id: ObjectId, frame: usize, prompt: Option<&Prompt>, reverse: bool, conditioning: bool) -> Result { + // EFFECTCRAFT_SAM2_PROFILE=1 prints where a frame's time goes (stderr). + let profile = std::env::var_os("EFFECTCRAFT_SAM2_PROFILE").is_some(); + let mut lap = std::time::Instant::now(); + let mut mark = |what: &str| { + if profile { + eprintln!("sam2 frame {frame}: {what} {:.0} ms", lap.elapsed().as_secs_f64() * 1e3); + } + lap = std::time::Instant::now(); + }; let feats = self.features(frame)?; + mark("encode"); let model = self.model.clone(); let tokens = |t: &Tensor| -> Result { Ok(t.flatten_from(2)?.squeeze(0)?.t()?.contiguous()?) }; let conditioned = if conditioning { @@ -303,9 +442,10 @@ impl<'a> VideoTracker<'a> { } else { let obj = self.objects.get(&id).ok_or_else(|| Error::Input("unknown object".into()))?; let (mem, mem_pos, n_ptr) = self.memory_bank(obj, frame, reverse)?; - let out = model.memory_attention.forward(&tokens(&feats.low)?, &tokens(&feats.pos)?, &mem, &mem_pos, n_ptr)?; + let out = model.attend(&tokens(&feats.low)?, &tokens(&feats.pos)?, &mem, &mem_pos, n_ptr)?; out.t()?.contiguous()?.reshape((1, 256, GRID, GRID))? }; + mark("memory attention"); let num_points = prompt.map_or(0, |p| p.points.len() + 2 * usize::from(p.bbox.is_some())); let multimask = num_points <= 1; let sparse = match prompt { @@ -315,6 +455,7 @@ impl<'a> VideoTracker<'a> { let dense = model.image.prompt.dense_empty()?; let out = model.image.decode_embedded(&conditioned, &feats.high, &sparse, &dense, multimask)?; + mark("decoder"); let score = scalar(&out.object_score)?; let present = score > 0.0; let ious = out.iou.flatten_all()?.to_vec1::()?; @@ -327,7 +468,9 @@ impl<'a> VideoTracker<'a> { let high = low.upsample_bilinear2d(INPUT_SIZE, INPUT_SIZE, false)?; let for_memory = if prompt.is_some() { high.gt(0f32)?.to_dtype(DType::F32)? } else { sigmoid(&high)? }; let for_memory = ((for_memory * 20.0)? - 10.0)?; - let memory = model.memory_encoder.memory(&feats.low, &for_memory, present)?; + mark("upsample and mask"); + let memory = model.memorize(&feats.low, &for_memory, present)?; + mark("memory encoder"); let iou = ious.get(best).copied().unwrap_or(0.0); let result = FrameMask { frame, object: id, logits: low.flatten_all()?.to_vec1::()?, iou, object_score: score, occluded: !present }; @@ -414,7 +557,7 @@ impl<'a> VideoTracker<'a> { let decoder_present = scalar(&out.object_score)? > 0.0; let pointer = if decoder_present { model.pointer_proj.forward(&token)? } else { model.no_object_pointer.clone() }; let pointer = if present { pointer } else { model.no_object_pointer.clone() }; - let memory = model.memory_encoder.memory(&feats.low, &high, present)?; + let memory = model.memorize(&feats.low, &high, present)?; let score = if present { 10.0 } else { -10.0 }; let result = FrameMask { frame, object: id, logits: low, iou: 1.0, object_score: score, occluded: !present }; let state = self.objects.entry(id).or_default(); diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index 59b09b7f1..16f5cfd24 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -223,12 +223,17 @@ fn iou(a: &[f32], b: &[f32]) -> f64 { const VIDEO_LOGIT_TOL: f32 = 0.02; const VIDEO_MIN_IOU: f64 = 0.995; -#[test] -fn video_tracking_matches_sam2_video_model() { +fn check_video_tracking(backend: effectcraft_sam2::session::Backend) { use effectcraft_sam2::prompt::Prompt; use effectcraft_sam2::session::{Cancel, Direction, ObjectId, VideoModel, VideoTracker}; let Some((w, r)) = load_video() else { return }; - let model = std::sync::Arc::new(VideoModel::load(&w).unwrap()); + let model = std::sync::Arc::new(match VideoModel::load_with(&w, backend) { + Ok(m) => m, + Err(e) => { + eprintln!("skipped: {e}"); + return; + } + }); let clip = r.get("clip.pixel_values").unwrap().clone(); let mut tracker = VideoTracker::new(model, Box::new(Clip(clip))).unwrap(); let obj = ObjectId(1); @@ -261,6 +266,19 @@ fn video_tracking_matches_sam2_video_model() { } } +#[test] +fn video_tracking_matches_sam2_video_model() { + check_video_tracking(effectcraft_sam2::session::Backend::Cpu); +} + +/// The same clip on the GPU: every frame's mask against `Sam2VideoModel`, forwards and backwards. The GPU path +/// is within 1e-5 of PyTorch stage by stage; through the clip the bfloat16 memories can round the other way +/// on a few values, so the bound on logits is the video one, and the masks must agree. +#[test] +fn video_tracking_on_the_gpu_matches_sam2_video_model() { + check_video_tracking(effectcraft_sam2::session::Backend::Gpu); +} + #[test] fn tracker_rejects_hostile_input_and_honours_cancel() { use effectcraft_sam2::Error; @@ -603,8 +621,8 @@ fn gpu_encoder_matches_the_cpu_path_and_pytorch() { use effectcraft_sam2::gpu::{Gpu, encoder::GpuEncoder}; use effectcraft_sam2::model::ImageModel; let Some((w, r)) = load() else { return }; - let gpu = match Gpu::new() { - Ok(g) => std::sync::Arc::new(g), + let gpu = match Gpu::shared() { + Ok(g) => g, Err(e) => { eprintln!("skipped: {e}"); return; @@ -632,3 +650,105 @@ fn gpu_encoder_matches_the_cpu_path_and_pytorch() { eprintln!("skip features vs CPU: {e0:e} / {e1:e}"); assert!(e0 < GPU_TOL && e1 < GPU_TOL, "skip feature errors {e0} / {e1}"); } + +/// The GPU memory attention against PyTorch (the fixture's frame: one memory frame and four pointer tokens). +/// Measured on the RX 7900 GRE, 2026-10-15; the CPU path is at 2.6e-5. +const GPU_ATTN_TOL: f32 = 1e-3; + +#[test] +fn gpu_memory_attention_matches_pytorch() { + use effectcraft_sam2::gpu::{Gpu, memattn::GpuMemoryAttention}; + let Some((w, r)) = load_video() else { return }; + let gpu = match Gpu::shared() { + Ok(g) => g, + Err(e) => { + eprintln!("skipped: {e}"); + return; + } + }; + let att = GpuMemoryAttention::load(gpu, &w).unwrap(); + let get = |k: &str| r.get(k).unwrap().squeeze(1).unwrap(); + let pointers = r.get("mem.attn.num_pointers").unwrap().to_vec1::().unwrap()[0] as usize; + let args = ( + get("mem.attn.current_vision_features"), + get("mem.attn.current_vision_position_embeddings"), + get("mem.attn.memory"), + get("mem.attn.memory_posision_embeddings"), + ); + let out = att.forward(&args.0, &args.1, &args.2, &args.3, pointers).unwrap(); + let want = r.get("mem.attn.out").unwrap().squeeze(0).unwrap().squeeze(0).unwrap(); + let e = max_abs(&out, &want); + eprintln!("GPU memory attention vs PyTorch: max abs error {e:e}"); + assert!(e < GPU_ATTN_TOL, "GPU memory attention error {e}"); + // a full bank (7 memory frames and 16 pointers of 4 tokens): the same maths on a bigger problem, against the CPU path + let cpu = effectcraft_sam2::memattn::MemoryAttention::load(&w).unwrap(); + let m = 7 * 4096 + 64; + let mem = Tensor::randn(0f32, 1.0, (m, 64), &Device::Cpu).unwrap(); + let mem_pos = Tensor::randn(0f32, 1.0, (m, 64), &Device::Cpu).unwrap(); + let t = std::time::Instant::now(); + let big = att.forward(&args.0, &args.1, &mem, &mem_pos, 64).unwrap(); + eprintln!("GPU memory attention, full bank: {:.0} ms", t.elapsed().as_secs_f64() * 1e3); + let want = cpu.forward(&args.0, &args.1, &mem, &mem_pos, 64).unwrap(); + let e = max_abs(&big, &want); + eprintln!("full bank vs CPU: max abs error {e:e}"); + assert!(e < GPU_ATTN_TOL, "full-bank error {e}"); + // bad shapes are errors + assert!(att.forward(&args.0, &args.1, &mem, &mem_pos, m + 1).is_err()); + assert!(att.forward(&args.0.narrow(0, 0, 10).unwrap(), &args.1, &mem, &mem_pos, 64).is_err()); +} + +/// The GPU can be switched off in the middle of a clip: the CPU takes over, the device name changes, and +/// the masks still match `Sam2VideoModel` (the memory banks built on the GPU are read by the CPU path). +#[test] +fn switching_the_gpu_off_mid_clip_keeps_the_tracking_right() { + use effectcraft_sam2::prompt::Prompt; + use effectcraft_sam2::session::{Backend, Cancel, Direction, ObjectId, VideoModel, VideoTracker}; + let Some((w, r)) = load_video() else { return }; + let Ok(model) = VideoModel::load_with(&w, Backend::Gpu) else { + eprintln!("skipped: no GPU"); + return; + }; + let model = std::sync::Arc::new(model); + assert_ne!(model.device_name(), "CPU"); + let mut t = VideoTracker::new(model.clone(), Box::new(Clip(r.get("clip.pixel_values").unwrap().clone()))).unwrap(); + t.prompt(2, ObjectId(1), &Prompt { points: vec![([416.0, 512.0], true)], bbox: None }).unwrap(); + let mut seen = vec![]; + t.propagate(2, Direction::Forward, Some(1), &Cancel::new(), &mut |m| seen.push(m)).unwrap(); + model.disable_gpu(); + assert_eq!(model.device_name(), "CPU"); + t.propagate(2, Direction::Forward, None, &Cancel::new(), &mut |m| seen.push(m)).unwrap(); + let last = seen.last().unwrap(); + assert_eq!(last.frame, 6); + let want = r.get("video.fwd.6.pred_masks").unwrap().flatten_all().unwrap().to_vec1::().unwrap(); + let j = iou(&last.logits, &want); + assert!(j >= VIDEO_MIN_IOU, "frame 6 after the switch: IoU {j}"); +} + +/// The GPU memory encoder against the CPU one on the fixture's frame (a mask, the frame's features). The CPU +/// path matches PyTorch to 3.5e-6; the GPU is held to the same output. +#[test] +fn gpu_memory_encoder_matches_the_cpu_one() { + use effectcraft_sam2::gpu::{Gpu, memenc::GpuMemoryEncoder}; + let Some((w, r)) = load_video() else { return }; + let Ok(gpu) = Gpu::shared() else { + eprintln!("skipped: no GPU"); + return; + }; + let cpu = effectcraft_sam2::memory::MemoryEncoder::load(&w).unwrap(); + let enc = GpuMemoryEncoder::load(gpu, &w).unwrap(); + let (features, mask) = (r.get("mem.enc.vision_features").unwrap(), r.get("mem.enc.masks").unwrap()); + let t = std::time::Instant::now(); + let got = enc.encode(features, mask).unwrap(); + eprintln!("GPU memory encoder: {:.0} ms", t.elapsed().as_secs_f64() * 1e3); + let want = r.get("mem.enc.out_features").unwrap().flatten_from(2).unwrap().squeeze(0).unwrap().t().unwrap().contiguous().unwrap(); + let e = max_abs(&got, &want); + eprintln!("GPU memory encoder vs PyTorch: max abs error {e:e}"); + assert!(e < 1e-3, "memory encoder error {e}"); + // the whole memory (with the occlusion embedding and bfloat16 rounding) equals the CPU one + let (a, b) = (cpu.finish(&got, false).unwrap(), cpu.memory(features, mask, false).unwrap()); + let e = max_abs(&a.features, &b.features); + assert!(e < 0.05, "memory features differ by {e} (bfloat16 steps are 2^-8 of the value)"); + assert_eq!(max_abs(&a.pos, &b.pos), 0.0); + // wrong shapes are errors + assert!(enc.encode(&features.narrow(1, 0, 8).unwrap(), mask).is_err()); +} diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index bd0b1e672..3fff1eec4 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -49,6 +49,10 @@ pub trait MaskModel: Send + Sync { fn tracker(&self) -> Option<&dyn ObjectTracker> { None } + /// What runs the model: the GPU's name or "CPU" (`None`: it does not say). + fn device_name(&self) -> Option { + None + } /// Seconds since the model was last used (`None`: it does not track use, so it is never unloaded /// for being idle). Large models free their memory after a quiet spell and load again on use. fn idle_secs(&self) -> Option { diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index 29a9e83a7..b3ceba2ad 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -64,20 +64,28 @@ the CPU. SAM 2 is prompted with at most three foreground and three background points per frame: it was trained on a handful of clicks, and measured confidence collapses with many (0.99 with 1–7 points, 0.25 with 16, 0.00 with 32). -## Speed (this PC: 16 threads, CPU only, release build) -| Step | Time | -|---|---| -| Encode one 1024 × 1024 frame | about **2.1 s** (3.4 s before the speed work: a parallel GELU, cached position encodings, pooling without permutes, a parallel bias add) | -| A click on an encoded frame | about 100 ms | -| Memory attention on a tracked frame (full bank of 7 memory frames) | about **2.3 s** | -| A tracked frame in all | about **4.6 s** (the first six have fewer memories) | +## Speed (this PC: AMD RX 7900 GRE and a 16-thread CPU, release build) +SAM 2 runs on the GPU when there is one (wgpu compute kernels, so AMD, NVIDIA and Intel alike; candle itself has no AMD +backend) and on the CPU otherwise. The model loads on the GPU automatically; `EFFECTCRAFT_SAM2_DEVICE=cpu` (or `gpu`) +overrides it. A GPU that fails (no adapter, out of memory, a lost device) makes the CPU take over for the rest of the +session, and `roto.status` reports `tracker.device` (the GPU's name, or `CPU`). -A 300-frame clip therefore takes on the order of 23 minutes on this CPU. (An earlier version of this page said one encode -per frame; the memory attention is as large as the encode.) What was tried and did not help: chunking the queries of the -global and memory attention to keep scores in cache. Candle has no AMD/Vulkan backend, so a GPU path is future work. A spike (`--example gpu_matmul`) measured a plain tiled -WGSL matmul on the RX 7900 GRE at 1,900–3,600 GFLOPS against 160–470 for candle's CPU matmul on the same shapes (6–13×), -so hand-written wgpu kernels could bring a tracked frame to a second or less; `cargo run --release -p effectcraft-sam2 --example sam2_bench` and `--example op_bench` time the -pieces. +| Step | CPU | GPU | +|---|---|---| +| Encode one 1024 × 1024 frame | 2.05 s | **0.24 s** | +| Memory attention (full bank of 7 memory frames) | 2.3 s | **0.29 s** | +| Memory encoder | 0.24 s | 0.05 s | +| Decoder (CPU on both) | 0.11 s | 0.11 s | +| **A tracked frame in all** | **4.3 s** | **0.7 s** | +| A click on an encoded frame | 0.10 s | 0.10 s | + +A 300-frame clip takes about 21 minutes on the CPU and about 3.5 minutes on this GPU; the decoder, still on the CPU, is the +next 16 % of a GPU frame. The CPU numbers came down from 3.4 s per encode at the start of the speed +work (parallel GELU and bias add, cached position encodings, pooling without permutes); chunking the attention queries was +tried and did not help. The GPU kernels (`crates/sam2/src/gpu`) are tested against CPU definitions and PyTorch: the +encoder's 64 × 64 level matches PyTorch to 7e-6, the memory attention to 1.8e-5, and a clip tracked on the GPU matches +`Sam2VideoModel` as the CPU does (worst logit error 0.0032). `cargo run --release -p effectcraft-sam2 --example sam2_bench` +times the pieces on both devices (`EFFECTCRAFT_SAM2_PROFILE=1` prints a frame's breakdown). ## Benchmark `effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt]` runs synthetic clips From 7f56fd87097bf1ae01e633c1d68106a482ca2ec4 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 21:53:57 +0200 Subject: [PATCH 38/77] S2.52: refreshed benchmark (24 frames, SAM 2 on the GPU) Occlusion 0.491 classical / 0.494 MobileSAM / 0.989 SAM 2; re-entry 0.618 / 0.620 / 0.982. Co-Authored-By: Claude Sonnet 5.5 --- docs/roto-brush-sam2.md | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index b3ceba2ad..d36d2086c 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -88,18 +88,18 @@ encoder's 64 × 64 level matches PyTorch to 7e-6, the memory attention to 1.8e-5 times the pieces on both devices (`EFFECTCRAFT_SAM2_PROFILE=1` prints a frame's breakdown). ## Benchmark -`effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt]` runs synthetic clips -with known mattes through all three. 12 frames, debug build, this PC; mean IoU (mean IoU after the object returns), seconds +`effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt]` runs synthetic clips with +known mattes through all three. 24 frames, debug build, SAM 2 on the GPU; mean IoU (mean IoU after the object returns), seconds per frame: | Clip | Classical | MobileSAM | SAM 2.1 Tiny | |---|---|---|---| -| move | 0.982 | 0.992 | 0.993 | -| light | 0.981 | 0.992 | 0.992 | -| occlusion | 0.978 | 0.881 | 0.988 | -| reenter | 0.632 (0.000) | 0.499 (0.725) | 0.975 (0.967) | -| s/frame | 0.01 | 1.1 | 8.2 | +| move | 0.980 | 0.989 | 0.992 | +| light | 0.979 | 0.990 | 0.992 | +| occlusion (a bar hides the disk, then clears it) | 0.491 (0.000) | 0.494 (0.000) | 0.989 (0.988) | +| reenter (the disk leaves and comes back) | 0.618 (0.000) | 0.620 (0.000) | 0.982 (0.976) | +| s/frame | 0.01 | 0.5–1.3 | 0.8 | MobileSAM prompts every frame independently from the flow-warped matte, so it follows the object while the flow does and -loses it when it is hidden; SAM 2's streaming memory is what lets it re-acquire the object. The synthetic clips are easy -(one textured disc), so the "move" and "light" rows mostly show all three working; the re-entry row is the difference. +loses it when it is hidden; SAM 2's streaming memory is what lets it re-acquire the object. The synthetic clips are easy (one +textured disc), so the "move" and "light" rows mostly show all three working; the last two rows are the difference. From a39cee8104248d368ca1146060d120b946cdae6a Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 22:03:29 +0200 Subject: [PATCH 39/77] S2.51: bench --roto --fixes counts the corrections each tracker needs Simulated user paints a foreground stroke through the true object and a background stroke on the biggest false positive on the first frame under IoU 0.9, up to 6 rounds. SAM 2 needs none on every clip; classical and MobileSAM need more than 6 on occlusion and re-entry. All trackers now take strokes on any frame. Co-Authored-By: Claude Sonnet 5.5 --- apps/effectcraft-cli/src/main.rs | 16 +++-- crates/track/src/roto/bench.rs | 108 ++++++++++++++++++++++++++++--- docs/roto-brush-sam2.md | 15 +++++ 3 files changed, 124 insertions(+), 15 deletions(-) diff --git a/apps/effectcraft-cli/src/main.rs b/apps/effectcraft-cli/src/main.rs index 27710a4a2..be066ecd1 100644 --- a/apps/effectcraft-cli/src/main.rs +++ b/apps/effectcraft-cli/src/main.rs @@ -27,7 +27,7 @@ //! (--gpu: CPU vs GPU ms/frame for every comp at Full and Half) //! effectcraft-cli bench --ops [--small] [--layers N] [--comps N] [--footage N] everyday-operation timings //! on a large generated project (open, save, auto-save, undo/redo, timeline, Project panel) -//! effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt] [--json] Roto Brush trackers on synthetic clips with known mattes +//! effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt] [--fixes] [--json] Roto Brush trackers on synthetic clips with known mattes //! (classical engine, MobileSAM and SAM 2.1 Tiny when installed; IoU, flicker, re-appearance, seconds per frame) //! effectcraft-cli script FILE.jsx [F.ecproj] | --eval CODE run an After Effects-style script //! effectcraft-cli mcp [--autosave | --bridge PORT] MCP server on stdio @@ -126,7 +126,7 @@ const USAGE: &str = "usage: effectcraft-cli Result<(), Failure> { if sam2.is_none() { note!("bench --roto: SAM 2.1 Tiny is not installed; no SAM 2 column"); } + // `--fixes`: also count the correction strokes each needs before every frame reaches IoU 0.9 (up to 6). + let fixes = args.flag("--fixes"); + let measure = |sc: &bench::Scenario, run: &dyn Fn(&bench::Scenario) -> Option| if fixes { bench::with_fixes(sc, 0.9, 6, run) } else { run(sc) }; let mut runs = vec![]; for sc in bench::scenarios(frames) { - runs.push(bench::run_flow(&sc, "classical", None)); + runs.extend(measure(&sc, &|s| Some(bench::run_flow(s, "classical", None)))); if let Some(m) = &mobilesam { - runs.push(bench::run_flow(&sc, "mobilesam", Some(m.as_ref()))); + runs.extend(measure(&sc, &|s| Some(bench::run_flow(s, "mobilesam", Some(m.as_ref()))))); } - if let Some(run) = sam2.as_ref().and_then(|m| bench::run_tracker(&sc, "sam2-tiny", m.as_ref())) { - runs.push(run); + if let Some(m) = &sam2 { + runs.extend(measure(&sc, &|s| bench::run_tracker(s, "sam2-tiny", m.as_ref()))); } } if args.flag("--json") { diff --git a/crates/track/src/roto/bench.rs b/crates/track/src/roto/bench.rs index 2ce0663df..58aea2385 100644 --- a/crates/track/src/roto/bench.rs +++ b/crates/track/src/roto/bench.rs @@ -11,13 +11,14 @@ use effectcraft_segment::{ClipFrame, MaskModel, TrackRequest}; use serde::Serialize; use web_time::Instant; -use super::{FrameSeg, SegOpts, Stroke, StrokeKind, iou, propagate_with, rgb_of, segment_tracked, segment_with, tracker_prompt}; +use super::{FrameSeg, SegOpts, Stroke, StrokeKind, iou, propagate_with, rgb_of, segment_tracked, segment_with, tracker_prompt_boxed}; pub const W: u32 = 200; pub const H: u32 = 150; const R: f64 = 32.0; /// A clip with the true matte of the visible object on every frame. +#[derive(Clone)] pub struct Scenario { pub name: &'static str, pub what: &'static str, @@ -42,6 +43,12 @@ pub struct Run { /// Mean IoU over the frames after the object's last absence (re-appearance), if it had one. pub after_return: Option, pub per_frame: Vec, + /// Correction strokes needed before every frame reaches the target IoU (`>N`: not within N); set by + /// [`with_fixes`], `None` when not measured. + pub fixes: Option, + /// The matte of every frame (not serialised). + #[serde(skip)] + pub mattes: Vec>, } fn noise(x: f64, y: f64, seed: u32) -> f64 { @@ -182,6 +189,8 @@ fn metrics(sc: &Scenario, tracker: &str, mattes: &[Vec], secs: f64) -> Run { secs_per_frame: secs / n, after_return, per_frame, + fixes: None, + mattes: mattes.to_vec(), } } @@ -190,7 +199,7 @@ fn metrics(sc: &Scenario, tracker: &str, mattes: &[Vec], secs: f64) -> Run { pub fn run_flow(sc: &Scenario, tracker: &str, model: Option<&dyn MaskModel>) -> Run { let opts = SegOpts::default(); let start = Instant::now(); - let refs: Vec<&Stroke> = sc.strokes.iter().collect(); + let refs: Vec<&Stroke> = sc.strokes.iter().filter(|s| s.frame == sc.base as i64).collect(); let mut segs: Vec> = vec![None; sc.frames.len()]; if let Some(first) = sc.frames.get(sc.base) { let base = segment_with(first, &refs, None, &opts, 0.0, model); @@ -201,7 +210,8 @@ pub fn run_flow(sc: &Scenario, tracker: &str, model: Option<&dyn MaskModel>) -> for f in sc.base + 1..sc.frames.len() { let (Some(prev_img), Some(img)) = (sc.frames.get(f - 1), sc.frames.get(f)) else { break }; let Some(Some(prev)) = segs.get(f - 1) else { break }; - let next = propagate_with(prev_img, prev, img, &[], &opts, 0.0, model); + let here: Vec<&Stroke> = sc.strokes.iter().filter(|s| s.frame == f as i64).collect(); + let next = propagate_with(prev_img, prev, img, &here, &opts, 0.0, model); if let Some(slot) = segs.get_mut(f) { *slot = Some(next); } @@ -218,16 +228,23 @@ pub fn run_flow(sc: &Scenario, tracker: &str, model: Option<&dyn MaskModel>) -> pub fn run_tracker(sc: &Scenario, tracker: &str, model: &dyn MaskModel) -> Option { let t = model.tracker()?; let opts = SegOpts::default(); - let refs: Vec<&Stroke> = sc.strokes.iter().collect(); + let at = |f: usize| -> Vec<&Stroke> { sc.strokes.iter().filter(|s| s.frame == f as i64).collect() }; let clip: Vec> = sc.frames.iter().map(|i| Arc::new(ClipFrame { rgb: rgb_of(i), w: i.width as usize, h: i.height as usize })).collect(); - let req = TrackRequest::single(0, sc.frames.len(), vec![(sc.base as i64, tracker_prompt(&refs, 1.0))]); + // a prompt per stroke frame: a few clicks and the box of the classic matte (as Chain::track does) + let mut prompts = vec![]; + for f in 0..sc.frames.len() { + let strokes = at(f); + let Some(img) = sc.frames.get(f).filter(|_| strokes.iter().any(|s| s.kind == StrokeKind::Fg)) else { continue }; + prompts.push((f as i64, tracker_prompt_boxed(&strokes, 1.0, &super::segment(img, &strokes, None, &opts, 0.0)))); + } + let req = TrackRequest::single(0, sc.frames.len(), prompts); let start = Instant::now(); let mut mattes: Vec> = vec![vec![0u8; (W * H) as usize]; sc.frames.len()]; let mut get = |f: i64| usize::try_from(f).ok().and_then(|i| clip.get(i).cloned()); let r = t.track(&req, &mut get, &|| false, &mut |_, f, tf| { let (Ok(i), true) = (usize::try_from(f), tf.prob.len() == (W * H) as usize) else { return }; let Some(img) = sc.frames.get(i) else { return }; - let strokes: Vec<&Stroke> = if i == sc.base { refs.clone() } else { vec![] }; + let strokes = at(i); if let Some(slot) = mattes.get_mut(i) { *slot = segment_tracked(img, &strokes, tf.prob, tf.occluded, &opts, 0.0).matte; } @@ -236,14 +253,74 @@ pub fn run_tracker(sc: &Scenario, tracker: &str, model: &dyn MaskModel) -> Optio Some(metrics(sc, tracker, &mattes, start.elapsed().as_secs_f64())) } +/// The correction a user would paint on frame `f` where `pred` is wrong: a short foreground stroke through +/// the true object's centre, and a background stroke on the biggest false-positive blob's centre. +fn correction(sc: &Scenario, f: usize, pred: &[u8]) -> Vec { + let truth = sc.truth.get(f).map(Vec::as_slice).unwrap_or(&[]); + let centre = |keep: &dyn Fn(usize) -> bool| -> Option<[f64; 2]> { + let (mut sx, mut sy, mut n) = (0.0, 0.0, 0.0); + for i in (0..truth.len().min(pred.len())).filter(|i| keep(*i)) { + sx += (i % W as usize) as f64 + 0.5; + sy += (i / W as usize) as f64 + 0.5; + n += 1.0; + } + (n > 0.0).then(|| [sx / n, sy / n]) + }; + let mut out = vec![]; + if let Some(c) = centre(&|i| truth.get(i).is_some_and(|v| *v != 0)) { + out.push(line(StrokeKind::Fg, f as i64, [c[0] - 12.0, c[1]], [c[0] + 12.0, c[1]], 5.0)); + } + if let Some(c) = centre(&|i| pred.get(i).is_some_and(|v| *v != 0) && truth.get(i).is_some_and(|v| *v == 0)) { + out.push(line(StrokeKind::Bg, f as i64, [c[0] - 6.0, c[1]], [c[0] + 6.0, c[1]], 6.0)); + } + out +} + +/// Repeat `run` on `sc`, each time painting the correction a user would on the first frame under `target` IoU, +/// until every frame reaches it or `max` corrections were needed. Sets `fixes` on the first run's result: the +/// number of corrections, or `>max`. +pub fn with_fixes(sc: &Scenario, target: f64, max: usize, run: &dyn Fn(&Scenario) -> Option) -> Option { + let mut cur = sc.clone(); + let mut first: Option = None; + for k in 0..=max { + let r = run(&cur)?; + let bad = r.per_frame.iter().position(|v| *v < target); + let first_run = first.get_or_insert_with(|| r.clone()); + let Some(f) = bad else { + first_run.fixes = Some(k.to_string()); + break; + }; + if k == max { + first_run.fixes = Some(format!(">{max}")); + break; + } + let pred = r.mattes.get(f).cloned().unwrap_or_default(); + let fix = correction(&cur, f, &pred); + if fix.is_empty() { + first_run.fixes = Some(format!(">{max}")); + break; + } + cur.strokes.extend(fix); + } + first +} + /// A fixed-width table of the runs, grouped by scenario. pub fn table(runs: &[Run]) -> String { - let mut out = format!("{:<10} {:<12} {:>8} {:>8} {:>8} {:>10} {:>9}\n", "clip", "tracker", "meanIoU", "minIoU", "flicker", "after ret.", "s/frame"); + let mut out = + format!("{:<10} {:<12} {:>8} {:>8} {:>8} {:>10} {:>9} {:>6}\n", "clip", "tracker", "meanIoU", "minIoU", "flicker", "after ret.", "s/frame", "fixes"); for r in runs { let back = r.after_return.map_or_else(|| "-".to_string(), |v| format!("{v:.3}")); out.push_str(&format!( - "{:<10} {:<12} {:>8.3} {:>8.3} {:>8.3} {:>10} {:>9.3}\n", - r.scenario, r.tracker, r.mean_iou, r.min_iou, r.flicker, back, r.secs_per_frame + "{:<10} {:<12} {:>8.3} {:>8.3} {:>8.3} {:>10} {:>9.3} {:>6}\n", + r.scenario, + r.tracker, + r.mean_iou, + r.min_iou, + r.flicker, + back, + r.secs_per_frame, + r.fixes.as_deref().unwrap_or("-") )); } out @@ -283,6 +360,19 @@ mod tests { assert!(table(&[run]).contains("classical")); } + #[test] + fn corrections_are_counted_until_every_frame_reaches_the_target() { + let all = scenarios(24); + let classical = |s: &Scenario| Some(run_flow(s, "classical", None)); + // plain motion needs none + let easy = all.iter().find(|s| s.name == "move").and_then(|s| with_fixes(s, 0.9, 6, &classical)).unwrap(); + assert_eq!(easy.fixes.as_deref(), Some("0")); + // after the disk has left and returned the classic engine has lost it: it needs a correction (or never gets there) + let hard = all.iter().find(|s| s.name == "reenter").and_then(|s| with_fixes(s, 0.9, 6, &classical)).unwrap(); + assert_ne!(hard.fixes.as_deref(), Some("0"), "{:?}", hard.fixes); + assert!(table(&[easy, hard]).contains("fixes")); + } + #[test] fn tiny_clip_requests_are_clamped() { assert!(scenarios(0).iter().all(|s| s.frames.len() == 12)); diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index d36d2086c..a08c8a5e7 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -103,3 +103,18 @@ per frame: MobileSAM prompts every frame independently from the flow-warped matte, so it follows the object while the flow does and loses it when it is hidden; SAM 2's streaming memory is what lets it re-acquire the object. The synthetic clips are easy (one textured disc), so the "move" and "light" rows mostly show all three working; the last two rows are the difference. + +### Corrections needed (`--fixes`) +`bench --roto --fixes` also counts the correction strokes each tracker needs before every frame reaches IoU 0.9: after each +run it paints what a user would on the first frame under the target (a short foreground stroke through the true object's +centre, and a background stroke on the biggest false-positive blob), up to 6. 16 frames, debug build: + +| Clip | Classical | MobileSAM | SAM 2.1 Tiny | +|---|---|---|---| +| move | 0 | 0 | 0 | +| light | 0 | 0 | 0 | +| occlusion | more than 6 | more than 6 | 0 | +| reenter | more than 6 | more than 6 | 0 | + +The simulated user is simple (one stroke pair per round, always on the first bad frame), so "more than 6" means this +correction model never recovers the flow-based engines on those clips, not that no human could. Peak memory is not measured. From 65c5eda39a8dde77c3c6dfbd4b28420b44a149b0 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 16 Oct 2026 04:02:53 +0200 Subject: [PATCH 40/77] S2.53: bench --roto takes the user's own clips and reports peak memory --clip FRAMES_DIR --truth MASKS_DIR runs the trackers on image sequences with true mattes (strokes made from the truth; nothing is bundled or copied), tested end to end. Trackers report the most memory their banks held (peakMB column). Co-Authored-By: Claude Sonnet 5.5 --- apps/effectcraft-cli/src/main.rs | 48 +++++++++++++- apps/effectcraft-cli/tests/it/cli.rs | 41 ++++++++++++ crates/sam2/src/mask.rs | 7 +- crates/sam2/src/session.rs | 11 +++- crates/sam2/src/track.rs | 1 + crates/sam2/tests/reference.rs | 22 +++++++ crates/segment/src/lib.rs | 4 ++ crates/track/src/roto/bench.rs | 99 +++++++++++++++++++++++----- 8 files changed, 211 insertions(+), 22 deletions(-) diff --git a/apps/effectcraft-cli/src/main.rs b/apps/effectcraft-cli/src/main.rs index be066ecd1..e6a3352dc 100644 --- a/apps/effectcraft-cli/src/main.rs +++ b/apps/effectcraft-cli/src/main.rs @@ -27,7 +27,7 @@ //! (--gpu: CPU vs GPU ms/frame for every comp at Full and Half) //! effectcraft-cli bench --ops [--small] [--layers N] [--comps N] [--footage N] everyday-operation timings //! on a large generated project (open, save, auto-save, undo/redo, timeline, Project panel) -//! effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt] [--fixes] [--json] Roto Brush trackers on synthetic clips with known mattes +//! effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt] [--fixes] [--clip FRAMES --truth MASKS] [--json] Roto Brush trackers on synthetic clips with known mattes //! (classical engine, MobileSAM and SAM 2.1 Tiny when installed; IoU, flicker, re-appearance, seconds per frame) //! effectcraft-cli script FILE.jsx [F.ecproj] | --eval CODE run an After Effects-style script //! effectcraft-cli mcp [--autosave | --bridge PORT] MCP server on stdio @@ -126,7 +126,7 @@ const USAGE: &str = "usage: effectcraft-cli Result<(), Failure> { let fixes = args.flag("--fixes"); let measure = |sc: &bench::Scenario, run: &dyn Fn(&bench::Scenario) -> Option| if fixes { bench::with_fixes(sc, 0.9, 6, run) } else { run(sc) }; let mut runs = vec![]; - for sc in bench::scenarios(frames) { + let mut scenarios = bench::scenarios(frames); + // `--clip FRAMES_DIR --truth MASKS_DIR`: the user's own clip (image sequences, same file order) with its true mattes. + if let Some(dir) = args.opt("--clip") { + let truth = args.opt("--truth").ok_or_else(|| Failure::Error("--clip needs --truth MASKS_DIR".into()))?; + let name = std::path::Path::new(dir).file_name().map_or_else(|| "clip".to_string(), |n| n.to_string_lossy().chars().take(10).collect()); + scenarios = vec![load_bench_clip(&name, dir, truth, frames).map_err(Failure::Error)?]; + } + for sc in scenarios { runs.extend(measure(&sc, &|s| Some(bench::run_flow(s, "classical", None)))); if let Some(m) = &mobilesam { runs.extend(measure(&sc, &|s| Some(bench::run_flow(s, "mobilesam", Some(m.as_ref()))))); @@ -1055,6 +1064,39 @@ fn bench_roto(args: &Args) -> Result<(), Failure> { Ok(()) } +/// The image files of `dir` (png, jpg, jpeg), in file-name order. +fn image_files(dir: &str) -> Result, String> { + let mut files: Vec = std::fs::read_dir(dir) + .map_err(|e| format!("cannot read {dir}: {e}"))? + .filter_map(|e| e.ok().map(|e| e.path())) + .filter(|p| p.extension().and_then(|e| e.to_str()).is_some_and(|e| ["png", "jpg", "jpeg"].contains(&e.to_ascii_lowercase().as_str()))) + .collect(); + files.sort(); + Ok(files) +} + +/// A benchmark scenario from the user's frames and truth masks (at most `limit` frames, frames at most 480 px +/// wide; a mask pixel above 127 is the object). Nothing is copied or kept. +fn load_bench_clip(name: &str, frames_dir: &str, truth_dir: &str, limit: usize) -> Result { + let (frames, masks) = (image_files(frames_dir)?, image_files(truth_dir)?); + if frames.is_empty() || frames.len() != masks.len() { + return Err(format!("{} frames in {frames_dir} but {} masks in {truth_dir}: they must match one to one", frames.len(), masks.len())); + } + let n = frames.len().min(limit.clamp(1, 240)); + let (mut imgs, mut truth) = (vec![], vec![]); + for (f, m) in frames.iter().zip(&masks).take(n) { + let rgba = image::open(f).map_err(|e| format!("{}: {e}", f.display()))?.to_rgba8(); + let scale = (480.0 / f64::from(rgba.width().max(1))).min(1.0); + let (w, h) = (((f64::from(rgba.width()) * scale).round() as u32).max(1), ((f64::from(rgba.height()) * scale).round() as u32).max(1)); + let rgba = if (w, h) == rgba.dimensions() { rgba } else { image::imageops::resize(&rgba, w, h, image::imageops::FilterType::Triangle) }; + let mask = image::open(m).map_err(|e| format!("{}: {e}", m.display()))?.to_luma8(); + let mask = image::imageops::resize(&mask, w, h, image::imageops::FilterType::Nearest); + imgs.push(effectcraft_render::Image::from_rgba8(w, h, rgba.as_raw())); + truth.push(mask.as_raw().iter().map(|v| u8::from(*v > 127)).collect()); + } + effectcraft_track::roto::bench::scenario_from(name, imgs, truth) +} + /// `bench --ops`: everyday operations on a large generated project (200 comps, 5,000 layers in /// the main comp, 300 footage items, 20-deep precomp nesting, expressions): startup to first /// frame, open, save, auto-save, small and large edits with undo/redo, first frame, timeline diff --git a/apps/effectcraft-cli/tests/it/cli.rs b/apps/effectcraft-cli/tests/it/cli.rs index 7a770fa87..5348f2a3a 100644 --- a/apps/effectcraft-cli/tests/it/cli.rs +++ b/apps/effectcraft-cli/tests/it/cli.rs @@ -805,3 +805,44 @@ fn render_without_progress_writes_no_progress_lines() { let stderr = render_demo_stderr("no-progress.gif", &[]); assert!(!stderr.lines().any(|l| l.starts_with("progress")), "{stderr}"); } + +/// `bench --roto --clip FRAMES --truth MASKS` runs the trackers on the user's own image sequences: here a bright disc +/// drifting over a dark textured frame, with its true masks. Mismatched sequences are refused with a message. +#[test] +fn bench_roto_runs_on_a_users_clip() { + let (frames, masks) = (tmp("clip-frames"), tmp("clip-masks")); + for d in [&frames, &masks] { + let _ = std::fs::remove_dir_all(d); + std::fs::create_dir_all(d).unwrap(); + } + let (w, h) = (96u32, 72u32); + for f in 0..6u32 { + let (cx, cy) = (30.0 + 4.0 * f as f32, 36.0f32); + let (mut img, mut mask) = (image::RgbImage::new(w, h), image::GrayImage::new(w, h)); + for y in 0..h { + for x in 0..w { + let inside = ((x as f32 - cx).powi(2) + (y as f32 - cy).powi(2)).sqrt() < 14.0; + let tex = 40 + ((x * 7 + y * 13) % 30) as u8; + img.put_pixel(x, y, if inside { image::Rgb([230, 120, 40 + (x % 20) as u8]) } else { image::Rgb([tex / 2, tex, tex + 20]) }); + mask.put_pixel(x, y, image::Luma([if inside { 255 } else { 0 }])); + } + } + img.save(frames.join(format!("f{f:03}.png"))).unwrap(); + mask.save(masks.join(format!("f{f:03}.png"))).unwrap(); + } + let (f, m) = (frames.to_str().unwrap(), masks.to_str().unwrap()); + let out = bin().args(["bench", "--roto", "--clip", f, "--truth", m, "--frames", "6"]).output().unwrap(); + let text = String::from_utf8_lossy(&out.stdout).to_string(); + assert!(out.status.success(), "{text}\n{}", String::from_utf8_lossy(&out.stderr)); + let row = text.lines().find(|l| l.contains("classical")).unwrap_or_else(|| panic!("no classical row: {text}")); + let mean: f64 = row.split_whitespace().nth(2).unwrap().parse().unwrap(); + assert!(mean > 0.8, "classical mean IoU on the user's clip: {mean}\n{text}"); + // one mask short: refused + std::fs::remove_file(masks.join("f005.png")).unwrap(); + let out = bin().args(["bench", "--roto", "--clip", f, "--truth", m]).output().unwrap(); + assert!(!out.status.success()); + assert!(String::from_utf8_lossy(&out.stderr).contains("must match"), "{}", String::from_utf8_lossy(&out.stderr)); + // --clip without --truth + let out = bin().args(["bench", "--roto", "--clip", f]).output().unwrap(); + assert!(!out.status.success()); +} diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index 94cf8b40a..4578af0f9 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -55,11 +55,12 @@ pub struct Sam2Mask { model: Arc, last: Mutex)>>, used: Mutex, + pub(crate) peak: std::sync::atomic::AtomicU64, } impl Sam2Mask { pub fn new(model: Arc) -> Sam2Mask { - Sam2Mask { model, last: Mutex::new(None), used: Mutex::new(std::time::Instant::now()) } + Sam2Mask { model, last: Mutex::new(None), used: Mutex::new(std::time::Instant::now()), peak: std::sync::atomic::AtomicU64::new(0) } } /// Note that the model is being used (for [`MaskModel::idle_secs`]). @@ -121,6 +122,10 @@ impl MaskModel for Sam2Mask { Some(self.model.device_name()) } + fn peak_state_bytes(&self) -> Option { + Some(self.peak.load(std::sync::atomic::Ordering::Relaxed)) + } + fn idle_secs(&self) -> Option { Some(self.used.lock().unwrap_or_else(PoisonError::into_inner).elapsed().as_secs_f64()) } diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 9de8fedc2..336518bd5 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -290,6 +290,7 @@ pub struct VideoTracker<'a> { objects: BTreeMap, cache: VecDeque<(usize, Arc)>, budget: Budget, + peak: usize, } fn sine_embedding_1d(positions: &[f32], dim: usize) -> Result { @@ -313,7 +314,7 @@ impl<'a> VideoTracker<'a> { if n == 0 || n > MAX_FRAMES { return Err(Error::Input(format!("a clip needs between 1 and {MAX_FRAMES} frames, got {n}"))); } - Ok(VideoTracker { model, frames, objects: BTreeMap::new(), cache: VecDeque::new(), budget: Budget::default() }) + Ok(VideoTracker { model, frames, objects: BTreeMap::new(), cache: VecDeque::new(), budget: Budget::default(), peak: 0 }) } pub fn num_frames(&self) -> usize { @@ -331,8 +332,14 @@ impl<'a> VideoTracker<'a> { frames + self.cache.len() * FEATURE_BYTES } - fn check_budget(&self) -> Result<()> { + /// The most [`VideoTracker::state_bytes`] reached so far. + pub fn peak_state_bytes(&self) -> usize { + self.peak + } + + fn check_budget(&mut self) -> Result<()> { let used = self.state_bytes(); + self.peak = self.peak.max(used); if used > self.budget.max_state_bytes { return Err(Error::Budget(format!("the tracker holds {} MB, over its {} MB limit", used >> 20, self.budget.max_state_bytes >> 20))); } diff --git a/crates/sam2/src/track.rs b/crates/sam2/src/track.rs index 2476ddb74..1aff06d9d 100644 --- a/crates/sam2/src/track.rs +++ b/crates/sam2/src/track.rs @@ -139,6 +139,7 @@ impl Sam2Mask { })?; } } + self.peak.store(tracker.peak_state_bytes() as u64, std::sync::atomic::Ordering::Relaxed); failure.map_or(Ok(()), Err) } } diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index 16f5cfd24..d121df840 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -752,3 +752,25 @@ fn gpu_memory_encoder_matches_the_cpu_one() { // wrong shapes are errors assert!(enc.encode(&features.narrow(1, 0, 8).unwrap(), mask).is_err()); } + +/// The tracker reports the most memory it held (memory banks, kept results, cached frames), and the model +/// reports it for the last tracking run. +#[test] +fn the_tracker_reports_its_peak_memory() { + use effectcraft_segment::{ClipFrame, MaskModel, Prompt, TrackRequest}; + let Some((w, _)) = load() else { return }; + let model = effectcraft_sam2::mask::Sam2Mask::new(std::sync::Arc::new(effectcraft_sam2::session::VideoModel::load(&w).unwrap())); + assert_eq!(model.peak_state_bytes(), Some(0)); + let (fw, fh) = (64usize, 48usize); + let frame = std::sync::Arc::new(ClipFrame { + rgb: (0..fw * fh).map(|i| if (i % fw).abs_diff(32) < 12 && (i / fw).abs_diff(24) < 12 { [0.9, 0.2, 0.1] } else { [0.2, 0.3, 0.5] }).collect(), + w: fw, + h: fh, + }); + let mut get = |_f: i64| Some(frame.clone()); + let req = TrackRequest::single(0, 3, vec![(0, Prompt { points: vec![([32.0, 24.0], true)], bbox: None, mask: None })]); + model.tracker().unwrap().track(&req, &mut get, &|| false, &mut |_, _, _| {}).unwrap(); + let peak = model.peak_state_bytes().unwrap(); + // three stored frames (about 2 MB each) and the cached frame features + assert!(peak > 6 << 20 && peak < 400 << 20, "peak {peak}"); +} diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index 3fff1eec4..3ebf964f9 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -49,6 +49,10 @@ pub trait MaskModel: Send + Sync { fn tracker(&self) -> Option<&dyn ObjectTracker> { None } + /// The most memory (bytes) the last tracking run held for its memory banks and cached frames (`None`: not tracked). + fn peak_state_bytes(&self) -> Option { + None + } /// What runs the model: the GPU's name or "CPU" (`None`: it does not say). fn device_name(&self) -> Option { None diff --git a/crates/track/src/roto/bench.rs b/crates/track/src/roto/bench.rs index 58aea2385..2dafb9f3f 100644 --- a/crates/track/src/roto/bench.rs +++ b/crates/track/src/roto/bench.rs @@ -20,10 +20,13 @@ const R: f64 = 32.0; /// A clip with the true matte of the visible object on every frame. #[derive(Clone)] pub struct Scenario { - pub name: &'static str, - pub what: &'static str, + pub name: String, + pub what: String, + /// Frame size in pixels (the synthetic clips are `W × H`; a real clip has its own). + pub w: usize, + pub h: usize, pub frames: Vec, - /// Visible-object matte per frame (`W×H`, 0/1); empty when the object is out of view. + /// Visible-object matte per frame (`w × h`, 0/1); empty when the object is out of view. pub truth: Vec>, /// Frame the strokes are painted on. pub base: usize, @@ -33,7 +36,7 @@ pub struct Scenario { /// One tracker's result on one scenario. #[derive(Clone, Debug, Serialize)] pub struct Run { - pub scenario: &'static str, + pub scenario: String, pub tracker: String, pub mean_iou: f64, pub min_iou: f64, @@ -46,6 +49,8 @@ pub struct Run { /// Correction strokes needed before every frame reaches the target IoU (`>N`: not within N); set by /// [`with_fixes`], `None` when not measured. pub fixes: Option, + /// The most memory (MB) the tracker held for its memory banks and cached frames (trackers that have none: `None`). + pub peak_mb: Option, /// The matte of every frame (not serialised). #[serde(skip)] pub mattes: Vec>, @@ -136,7 +141,46 @@ fn scenario(name: &'static str, what: &'static str, n: usize, look: impl Fn(usiz frames.push(img); truth.push(gt); } - Scenario { name, what, frames, truth, base: 0, strokes: strokes_for(look(0).c) } + Scenario { name: name.to_string(), what: what.to_string(), w: W as usize, h: H as usize, frames, truth, base: 0, strokes: strokes_for(look(0).c) } +} + +/// A scenario from the user's own clip: frames and their true mattes (0/1, same size as the frames). The strokes +/// are made from the first frame with an object: a short foreground stroke through its centre and background +/// strokes along the top and bottom edges where the truth is empty. Nothing is bundled: the clip stays with the +/// user (a licensed dataset, say), and so does its licence. +pub fn scenario_from(name: &str, frames: Vec, truth: Vec>) -> std::result::Result { + let first = frames.first().ok_or("a clip needs at least one frame")?; + let (w, h) = (first.width as usize, first.height as usize); + if w == 0 + || h == 0 + || frames.len() != truth.len() + || frames.iter().any(|f| (f.width as usize, f.height as usize) != (w, h)) + || truth.iter().any(|t| t.len() != w * h) + { + return Err("frames and truth masks must be the same size and count".into()); + } + let base = truth.iter().position(|t| t.iter().any(|v| *v != 0)).ok_or("no frame has the object in its truth mask")?; + let t = truth.get(base).map(Vec::as_slice).unwrap_or(&[]); + let (mut x0, mut x1, mut sx, mut sy, mut n) = (usize::MAX, 0usize, 0.0f64, 0.0f64, 0.0f64); + for i in t.iter().enumerate().filter(|(_, v)| **v != 0).map(|(i, _)| i) { + let (x, y) = (i % w, i / w); + (x0, x1) = (x0.min(x), x1.max(x)); + sx += x as f64 + 0.5; + sy += y as f64 + 0.5; + n += 1.0; + } + let (cx, cy) = (sx / n, sy / n); + let reach = (((x1 - x0) as f64) * 0.15).max(2.0); + let radius = ((w.min(h) as f64) * 0.03).max(3.0); + let mut strokes = vec![line(StrokeKind::Fg, base as i64, [cx - reach, cy], [cx + reach, cy], radius)]; + for y in [h as f64 * 0.04, h as f64 * 0.96] { + let row = y as usize; + let free = t.get(row * w..(row + 1) * w).is_some_and(|r| r.iter().all(|v| *v == 0)); + if free { + strokes.push(line(StrokeKind::Bg, base as i64, [w as f64 * 0.05, y], [w as f64 * 0.95, y], radius)); + } + } + Ok(Scenario { name: name.to_string(), what: "a clip supplied by the user".to_string(), w, h, frames, truth, base, strokes }) } /// The benchmark clips, `n` frames each (at least 12, so every event happens). @@ -181,7 +225,7 @@ fn metrics(sc: &Scenario, tracker: &str, mattes: &[Vec], secs: f64) -> Run { (!tail.is_empty()).then(|| tail.iter().sum::() / tail.len() as f64) }); Run { - scenario: sc.name, + scenario: sc.name.clone(), tracker: tracker.to_string(), mean_iou: per_frame.iter().sum::() / n, min_iou: per_frame.iter().copied().fold(1.0, f64::min), @@ -190,6 +234,7 @@ fn metrics(sc: &Scenario, tracker: &str, mattes: &[Vec], secs: f64) -> Run { after_return, per_frame, fixes: None, + peak_mb: None, mattes: mattes.to_vec(), } } @@ -217,7 +262,7 @@ pub fn run_flow(sc: &Scenario, tracker: &str, model: Option<&dyn MaskModel>) -> } } let secs = start.elapsed().as_secs_f64(); - let empty = vec![0u8; (W * H) as usize]; + let empty = vec![0u8; sc.w * sc.h]; let mattes: Vec> = segs.into_iter().map(|s| s.map_or_else(|| empty.clone(), |s| s.matte)).collect(); metrics(sc, tracker, &mattes, secs) } @@ -239,10 +284,10 @@ pub fn run_tracker(sc: &Scenario, tracker: &str, model: &dyn MaskModel) -> Optio } let req = TrackRequest::single(0, sc.frames.len(), prompts); let start = Instant::now(); - let mut mattes: Vec> = vec![vec![0u8; (W * H) as usize]; sc.frames.len()]; + let mut mattes: Vec> = vec![vec![0u8; sc.w * sc.h]; sc.frames.len()]; let mut get = |f: i64| usize::try_from(f).ok().and_then(|i| clip.get(i).cloned()); let r = t.track(&req, &mut get, &|| false, &mut |_, f, tf| { - let (Ok(i), true) = (usize::try_from(f), tf.prob.len() == (W * H) as usize) else { return }; + let (Ok(i), true) = (usize::try_from(f), tf.prob.len() == sc.w * sc.h) else { return }; let Some(img) = sc.frames.get(i) else { return }; let strokes = at(i); if let Some(slot) = mattes.get_mut(i) { @@ -250,7 +295,9 @@ pub fn run_tracker(sc: &Scenario, tracker: &str, model: &dyn MaskModel) -> Optio } }); r.ok()?; - Some(metrics(sc, tracker, &mattes, start.elapsed().as_secs_f64())) + let mut run = metrics(sc, tracker, &mattes, start.elapsed().as_secs_f64()); + run.peak_mb = model.peak_state_bytes().map(|b| b as f64 / (1 << 20) as f64); + Some(run) } /// The correction a user would paint on frame `f` where `pred` is wrong: a short foreground stroke through @@ -260,8 +307,8 @@ fn correction(sc: &Scenario, f: usize, pred: &[u8]) -> Vec { let centre = |keep: &dyn Fn(usize) -> bool| -> Option<[f64; 2]> { let (mut sx, mut sy, mut n) = (0.0, 0.0, 0.0); for i in (0..truth.len().min(pred.len())).filter(|i| keep(*i)) { - sx += (i % W as usize) as f64 + 0.5; - sy += (i / W as usize) as f64 + 0.5; + sx += (i % sc.w.max(1)) as f64 + 0.5; + sy += (i / sc.w.max(1)) as f64 + 0.5; n += 1.0; } (n > 0.0).then(|| [sx / n, sy / n]) @@ -307,12 +354,14 @@ pub fn with_fixes(sc: &Scenario, target: f64, max: usize, run: &dyn Fn(&Scenario /// A fixed-width table of the runs, grouped by scenario. pub fn table(runs: &[Run]) -> String { - let mut out = - format!("{:<10} {:<12} {:>8} {:>8} {:>8} {:>10} {:>9} {:>6}\n", "clip", "tracker", "meanIoU", "minIoU", "flicker", "after ret.", "s/frame", "fixes"); + let mut out = format!( + "{:<10} {:<12} {:>8} {:>8} {:>8} {:>10} {:>9} {:>6} {:>7}\n", + "clip", "tracker", "meanIoU", "minIoU", "flicker", "after ret.", "s/frame", "fixes", "peakMB" + ); for r in runs { let back = r.after_return.map_or_else(|| "-".to_string(), |v| format!("{v:.3}")); out.push_str(&format!( - "{:<10} {:<12} {:>8.3} {:>8.3} {:>8.3} {:>10} {:>9.3} {:>6}\n", + "{:<10} {:<12} {:>8.3} {:>8.3} {:>8.3} {:>10} {:>9.3} {:>6} {:>7}\n", r.scenario, r.tracker, r.mean_iou, @@ -320,7 +369,8 @@ pub fn table(runs: &[Run]) -> String { r.flicker, back, r.secs_per_frame, - r.fixes.as_deref().unwrap_or("-") + r.fixes.as_deref().unwrap_or("-"), + r.peak_mb.map_or_else(|| "-".to_string(), |v| format!("{v:.0}")) )); } out @@ -373,6 +423,23 @@ mod tests { assert!(table(&[easy, hard]).contains("fixes")); } + #[test] + fn a_users_clip_becomes_a_scenario_and_bad_ones_are_refused() { + let all = scenarios(24); + let src = all.iter().find(|s| s.name == "move").unwrap(); + let sc = scenario_from("mine", src.frames.clone(), src.truth.clone()).unwrap(); + assert_eq!((sc.w, sc.h, sc.base), (W as usize, H as usize, 0)); + assert!(sc.strokes.iter().any(|s| s.kind == StrokeKind::Fg) && sc.strokes.iter().any(|s| s.kind == StrokeKind::Bg)); + let run = run_flow(&sc, "classical", None); + assert!(run.mean_iou > 0.9, "{}", run.mean_iou); + assert!(scenario_from("x", vec![], vec![]).is_err()); + assert!(scenario_from("x", src.frames.clone(), src.truth[..3].to_vec()).is_err()); + assert!(scenario_from("x", src.frames.clone(), vec![vec![0; sc.w * sc.h]; src.frames.len()]).is_err(), "no object anywhere"); + let mut short = src.truth.clone(); + short[0].pop(); + assert!(scenario_from("x", src.frames.clone(), short).is_err()); + } + #[test] fn tiny_clip_requests_are_clamped() { assert!(scenarios(0).iter().all(|s| s.frames.len() == 12)); From 066bb2d031edf8ca2ab223833b11c232a965c305 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 16 Oct 2026 04:05:58 +0200 Subject: [PATCH 41/77] S2.44: the span bar's confidence colours, checked in a real render egui_kittest test: a scripted tracker (sure, doubtful, unsure, absent, sure again) tracks a layer and the Layer panel is rendered; the strip's pixels are green, amber, red, slate and green at the right frames. ROTO_SPAN_SNAPSHOT saves the panel image. Co-Authored-By: Claude Sonnet 5.5 --- Cargo.lock | 1 + crates/ui-egui/Cargo.toml | 1 + .../ui-egui/tests/ui_roto_span_confidence.rs | 108 ++++++++++++++++++ 3 files changed, 110 insertions(+) create mode 100644 crates/ui-egui/tests/ui_roto_span_confidence.rs diff --git a/Cargo.lock b/Cargo.lock index ce98aeadc..b4acc3a6c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1930,6 +1930,7 @@ dependencies = [ "effectcraft-path", "effectcraft-project", "effectcraft-render", + "effectcraft-segment", "effectcraft-text", "effectcraft-time", "eframe", diff --git a/crates/ui-egui/Cargo.toml b/crates/ui-egui/Cargo.toml index b1e281627..89532edb3 100644 --- a/crates/ui-egui/Cargo.toml +++ b/crates/ui-egui/Cargo.toml @@ -45,6 +45,7 @@ sys-locale = { workspace = true } egui = { workspace = true, features = ["default_fonts", "monochrome_emoji_fonts"] } egui_kittest = { workspace = true, features = ["eframe", "wgpu"] } pollster = "1" +effectcraft-segment = { workspace = true } effectcraft-host = { workspace = true } exr = { workspace = true } diff --git a/crates/ui-egui/tests/ui_roto_span_confidence.rs b/crates/ui-egui/tests/ui_roto_span_confidence.rs new file mode 100644 index 000000000..3b3148d33 --- /dev/null +++ b/crates/ui-egui/tests/ui_roto_span_confidence.rs @@ -0,0 +1,108 @@ +//! The Roto Brush span bar colours the frames a tracker followed by its confidence: green when it is sure, +//! amber when it is doubtful, red when it is not, slate where the object was absent (egui_kittest, UI logic +//! and a real render). A stand-in tracker supplies the confidences, so no weights are needed. +//! `ROTO_SPAN_SNAPSHOT=/abs/out.png` also saves the rendered panel. + +use std::sync::Arc; + +use effectcraft_engine::Session; +use effectcraft_segment::{ClipFrame, MaskModel, ModelInfo, ObjectTracker, Prompt, TrackRequest, TrackedFrame}; +use effectcraft_ui_egui::EffectcraftApp; +use effectcraft_ui_egui::dock::PanelKind; +use egui_kittest::Harness; +use serde_json::json; + +/// Follows a fixed disc; its confidence per frame is scripted. +struct Scripted; + +impl MaskModel for Scripted { + fn info(&self) -> &'static ModelInfo { + &effectcraft_segment::SAM2_TINY + } + fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &Prompt) -> effectcraft_segment::Result> { + Err("no weights".into()) + } + fn tracker(&self) -> Option<&dyn ObjectTracker> { + Some(self) + } +} + +/// (confidence, occluded) of layer frame `f`: sure, doubtful, unsure, absent, sure again. +fn script(f: i64) -> (f32, bool) { + match f { + 0..=5 => (0.99, false), + 6..=9 => (0.6, false), + 10..=12 => (0.2, false), + 13..=15 => (0.0, true), + _ => (0.99, false), + } +} + +impl ObjectTracker for Scripted { + fn track( + &self, + req: &TrackRequest, + frame: &mut dyn FnMut(i64) -> Option>, + _: &dyn Fn() -> bool, + out: &mut dyn FnMut(usize, i64, TrackedFrame), + ) -> effectcraft_segment::Result<()> { + for g in req.first..req.first + req.count as i64 { + let img = frame(g).ok_or("no frame")?; + let (confidence, occluded) = script(g); + let prob = (0..img.w * img.h).map(|i| if !occluded && (i % img.w).abs_diff(img.w / 2) < img.w / 4 && (i / img.w).abs_diff(img.h / 2) < img.h / 4 { 1.0 } else { 0.0 }).collect(); + out(0, g, TrackedFrame { prob, w: img.w, h: img.h, confidence, occluded }); + } + Ok(()) + } +} + +fn close(got: [u8; 4], want: [u8; 3]) -> bool { + got.iter().zip(want).all(|(a, b)| a.abs_diff(b) <= 24) +} + +#[test] +fn the_span_bar_shows_the_trackers_confidence() { + let mut s = Session::default(); + s.execute("comp.new", json!({"name": "View", "width": 640, "height": 360, "duration": 4})).unwrap(); + s.execute("layer.newSolid", json!({"name": "Box", "color": "#e04020", "width": 80, "height": 80})).unwrap(); + let id = s.active_comp().unwrap().layers.iter().find(|l| l.name == "Box").unwrap().id.0; + s.execute("roto.stroke", json!({"layer": id, "points": [[20.0, 40.0], [60.0, 40.0]], "radius": 6})).unwrap(); + s.execute("roto.stroke", json!({"layer": id, "kind": "bg", "points": [[5.0, 5.0], [75.0, 5.0]], "radius": 4})).unwrap(); + effectcraft_engine::effects::roto::set_model(Some(Arc::new(Scripted))); + // the model stays loaded while the panel draws (its id is part of the cache keys) + s.execute("roto.propagate", json!({"layer": id, "direction": "both", "wait": true})).unwrap(); + + let mut h = Harness::builder().with_size(egui::vec2(1600.0, 1000.0)).build_eframe(move |_| EffectcraftApp::new(s)); + h.run_steps(3); + h.state_mut().session.execute("layer.select", json!({"layers": [id]})).unwrap(); + let ctx = h.ctx.clone(); + effectcraft_ui_egui::menus::invoke(h.state_mut(), &ctx, "layer.openLayer", json!({"layer": id})).unwrap(); + h.run_steps(4); + assert_eq!(h.state().ui.focused, PanelKind::Layer); + + let track = h.state().auto.find("layerPanel.rotoSpan").expect("the span bar is on screen").rect; + let (comp, layer) = { + let comp = h.state().session.active_comp().unwrap().clone(); + let layer = comp.layer(effectcraft_project::LayerId(id)).unwrap().clone(); + (comp, layer) + }; + let lim = effectcraft_engine::roto::frame_limits(&comp, &layer); + let n = (lim[1] - lim[0] + 1) as f32; + let image = h.render().expect("the panel renders"); + if let Some(path) = std::env::var_os("ROTO_SPAN_SNAPSHOT") { + image.save(path).unwrap(); + } + let scale = image.width() as f32 / 1600.0; + let sample = |f: i64| -> [u8; 4] { + let x = (track[0] + ((f - lim[0]) as f32 + 0.5) / n * track[2]) * scale; + let y = (track[1] + track[3] - 2.0) * scale; + image.get_pixel(x as u32, y as u32).0 + }; + // the colours the bar uses (see `trust_color`) + let (green, amber, red, slate) = ([70, 200, 90], [230, 180, 50], [220, 70, 60], [110, 118, 150]); + for (f, want, what) in [(2, green, "sure"), (7, amber, "doubtful"), (11, red, "unsure"), (14, slate, "absent"), (18, green, "sure again")] { + let got = sample(f); + assert!(close(got, want), "frame {f} ({what}): drew {got:?}, wanted about {want:?}"); + } + effectcraft_engine::effects::roto::set_model(None); +} From 2acd0ca1a8dba0ed1a0a391c37431c8fc9aa9480 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Sat, 17 Oct 2026 13:36:23 +0200 Subject: [PATCH 42/77] S2.64: SAM 2.1 Small, Base+ and Large, each checked against PyTorch Hiera configs and size detection from the weights, registry entries pinned to commits with verified SHA-256, the loader picks the size by checksum. The GPU encoder submits once per block so Large fits in memory, and uploads no longer panic on a lost device. Reference tests run per size on the CPU and GPU. Co-Authored-By: Claude Sonnet 5.5 --- crates/engine/src/analysis/models.rs | 2 +- crates/engine/src/tests/models.rs | 4 +- crates/sam2/src/gpu/encoder.rs | 2 + crates/sam2/src/gpu/mod.rs | 12 +++- crates/sam2/src/hiera.rs | 39 +++++++++++++ crates/sam2/src/lib.rs | 13 +++++ crates/sam2/src/mask.rs | 21 +++++-- crates/sam2/src/model.rs | 2 +- crates/sam2/src/neck.rs | 8 ++- crates/sam2/src/weights.rs | 9 +++ crates/sam2/tests/reference.rs | 55 ++++++++++++++----- crates/segment/src/lib.rs | 55 ++++++++++++++++++- .../ui-egui/tests/ui_roto_span_confidence.rs | 4 +- docs/roto-brush-sam2.md | 18 ++++++ 14 files changed, 212 insertions(+), 32 deletions(-) diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index a584aebbb..d6d9ef63b 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -336,7 +336,7 @@ 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 { + if effectcraft_sam2::mask::SIZES.iter().any(|m| m.id == 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 fab025a3a..ee05ea6b8 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -36,8 +36,8 @@ fn models_are_listed_chosen_and_verified() { for (prefix, model, task) in [("roto", "mobilesam", Task::Mask), ("face", "mediapipe-face", Task::Face)] { let list = s.execute_checked(&format!("{prefix}.models"), json!({})).unwrap(); let ids: Vec<&str> = list["models"].as_array().unwrap().iter().map(|m| m["id"].as_str().unwrap()).collect(); - // Roto Brush also offers SAM 2.1 Tiny (a video tracker; see crates/sam2). - let expect: Vec<&str> = if task == Task::Mask { vec!["classical", model, "sam2.1-hiera-tiny"] } else { vec!["classical", model] }; + // Roto Brush also offers SAM 2.1 in four sizes (video trackers; see crates/sam2). + let expect: Vec<&str> = if task == Task::Mask { vec!["classical", model, "sam2.1-hiera-tiny", "sam2.1-hiera-small", "sam2.1-hiera-base-plus", "sam2.1-hiera-large"] } else { vec!["classical", model] }; assert_eq!(ids, expect); let m = &list["models"][1]; assert_eq!((m["licence"].as_str(), m["installed"].as_bool()), (Some("Apache-2.0"), Some(false))); diff --git a/crates/sam2/src/gpu/encoder.rs b/crates/sam2/src/gpu/encoder.rs index 630f4f09b..1d390a9d1 100644 --- a/crates/sam2/src/gpu/encoder.rs +++ b/crates/sam2/src/gpu/encoder.rs @@ -244,6 +244,8 @@ impl GpuEncoder { for (i, b) in self.blocks.iter().enumerate() { let (nx, nh, nw) = self.block(b, &x, h, w)?; (x, h, w) = (nx, nh, nw); + // submit per block so the GPU can free the previous block's intermediates (Large would otherwise run out of memory) + g.flush()?; if self.stage_ends.contains(&i) { stages.push((x.clone(), h, w, b.plan.dim_out)); } diff --git a/crates/sam2/src/gpu/mod.rs b/crates/sam2/src/gpu/mod.rs index 9ca180932..5b091062c 100644 --- a/crates/sam2/src/gpu/mod.rs +++ b/crates/sam2/src/gpu/mod.rs @@ -147,11 +147,17 @@ impl Gpu { } pub fn upload(&self, data: &[f32]) -> Buf { - let buffer = self.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + // `write_buffer` instead of a mapped-at-creation buffer: on a lost or out-of-memory device the latter panics inside wgpu, + // while this records the error for `take_error`. + let buffer = self.device.create_buffer(&wgpu::BufferDescriptor { label: None, - contents: bytemuck::cast_slice(data), - usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + size: (data.len().max(1) as u64) * 4, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, }); + if !data.is_empty() { + self.queue.write_buffer(&buffer, 0, bytemuck::cast_slice(data)); + } Buf { buffer, len: data.len() } } diff --git a/crates/sam2/src/hiera.rs b/crates/sam2/src/hiera.rs index 246448e86..82862c829 100644 --- a/crates/sam2/src/hiera.rs +++ b/crates/sam2/src/hiera.rs @@ -42,6 +42,45 @@ impl HieraConfig { pos_background: 7, } } + + /// `facebook/sam2.1-hiera-small`: the same widths as Tiny with deeper third and fourth stages. + pub fn small() -> HieraConfig { + HieraConfig { blocks: vec![1, 2, 11, 2], global_blocks: vec![7, 10, 13], ..HieraConfig::tiny() } + } + + /// `facebook/sam2.1-hiera-base-plus`. + pub fn base_plus() -> HieraConfig { + HieraConfig { + embed_dims: vec![112, 224, 448, 896], + blocks: vec![2, 3, 16, 3], + heads: vec![2, 4, 8, 16], + global_blocks: vec![12, 16, 20], + pos_background: 14, + ..HieraConfig::tiny() + } + } + + /// `facebook/sam2.1-hiera-large`. + pub fn large() -> HieraConfig { + HieraConfig { + embed_dims: vec![144, 288, 576, 1152], + blocks: vec![2, 6, 36, 4], + heads: vec![2, 4, 8, 16], + windows: vec![8, 4, 16, 8], + global_blocks: vec![23, 33, 43], + ..HieraConfig::tiny() + } + } + + /// The configuration of a released size. + pub fn for_size(size: crate::Size) -> HieraConfig { + match size { + crate::Size::Tiny => HieraConfig::tiny(), + crate::Size::Small => HieraConfig::small(), + crate::Size::BasePlus => HieraConfig::base_plus(), + crate::Size::Large => HieraConfig::large(), + } + } } fn at(v: &[T], i: usize, what: &str) -> Result { diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 59eceddb3..09cbb1466 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -82,6 +82,19 @@ impl Size { Size::Large => "facebook/sam2.1-hiera-large", } } + + /// The size a checkpoint is, from its tensors: the width of the patch embedding (96 for Tiny and Small, 112 for + /// Base+, 144 for Large) and the number of blocks (12 for Tiny, 16 for Small). + pub fn detect(w: &weights::Weights) -> Result { + let width = w.shape("vision_encoder.backbone.patch_embed.projection.weight").and_then(|s| s.first().copied()); + Ok(match width { + Some(112) => Size::BasePlus, + Some(144) => Size::Large, + Some(96) if w.contains("vision_encoder.backbone.blocks.15.layer_norm1.weight") => Size::Small, + Some(96) => Size::Tiny, + _ => return Err(Error::Weights("this checkpoint is not a SAM 2.1 Hiera size (patch embedding width)".into())), + }) + } } #[cfg(test)] diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index 4578af0f9..251bb05ca 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -53,6 +53,7 @@ fn fingerprint(rgb: &[[f32; 3]], w: usize, h: usize) -> u64 { /// SAM 2 for single frames. pub struct Sam2Mask { model: Arc, + info: &'static ModelInfo, last: Mutex)>>, used: Mutex, pub(crate) peak: std::sync::atomic::AtomicU64, @@ -60,7 +61,12 @@ pub struct Sam2Mask { impl Sam2Mask { pub fn new(model: Arc) -> Sam2Mask { - Sam2Mask { model, last: Mutex::new(None), used: Mutex::new(std::time::Instant::now()), peak: std::sync::atomic::AtomicU64::new(0) } + Sam2Mask::with_info(model, &SAM2_TINY) + } + + /// A model of a given registry entry (its size). + pub fn with_info(model: Arc, info: &'static ModelInfo) -> Sam2Mask { + Sam2Mask { model, info, last: Mutex::new(None), used: Mutex::new(std::time::Instant::now()), peak: std::sync::atomic::AtomicU64::new(0) } } /// Note that the model is being used (for [`MaskModel::idle_secs`]). @@ -106,7 +112,7 @@ impl Sam2Mask { impl MaskModel for Sam2Mask { fn info(&self) -> &'static ModelInfo { - &SAM2_TINY + self.info } fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &effectcraft_segment::Prompt) -> effectcraft_segment::Result> { @@ -134,14 +140,17 @@ impl MaskModel for Sam2Mask { /// 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 Some(info) = SIZES.iter().find(|m| m.sha256 == got) else { + return Err(format!("SAM 2.1: the file does not match any published size (SHA-256 {got})")); + }; let w = crate::weights::Weights::from_bytes(bytes).map_err(|e| e.to_string())?; let model = VideoModel::load_auto(&w).map_err(|e| e.to_string())?; - Ok(Arc::new(Sam2Mask::new(Arc::new(model)))) + Ok(Arc::new(Sam2Mask::with_info(Arc::new(model), info))) } +/// The registry entries this crate loads. +pub const SIZES: [&ModelInfo; 4] = [&SAM2_TINY, &effectcraft_segment::SAM2_SMALL, &effectcraft_segment::SAM2_BASE_PLUS, &effectcraft_segment::SAM2_LARGE]; + #[cfg(test)] mod tests { use super::*; diff --git a/crates/sam2/src/model.rs b/crates/sam2/src/model.rs index c30cbb5b9..36e30e310 100644 --- a/crates/sam2/src/model.rs +++ b/crates/sam2/src/model.rs @@ -32,7 +32,7 @@ impl ImageModel { pub fn load(w: &Weights) -> Result { let prompt = PromptEncoder::load(w)?; Ok(ImageModel { - vision: VisionEncoder::load_tiny(w)?, + vision: VisionEncoder::load(w)?, decoder: MaskDecoder::load(w)?, no_memory: w.get("no_memory_embedding", &[1, 1, 256])?.reshape((1, 256, 1, 1))?, image_pos: prompt.image_positions()?, diff --git a/crates/sam2/src/neck.rs b/crates/sam2/src/neck.rs index 01ab5523b..c897e9818 100644 --- a/crates/sam2/src/neck.rs +++ b/crates/sam2/src/neck.rs @@ -142,9 +142,11 @@ pub struct ImageFeatures { } impl VisionEncoder { - pub fn load_tiny(w: &Weights) -> Result { - let cfg = HieraConfig::tiny(); - Ok(VisionEncoder { hiera: std::sync::Arc::new(Hiera::load(w, &cfg)?), neck: Neck::load(w, &[768, 384, 192, 96], 256, &[2, 3])?, levels: 3 }) + /// The encoder of whichever size the checkpoint is. + pub fn load(w: &Weights) -> Result { + let cfg = HieraConfig::for_size(crate::Size::detect(w)?); + let channels: Vec = cfg.embed_dims.iter().rev().copied().collect(); + Ok(VisionEncoder { hiera: std::sync::Arc::new(Hiera::load(w, &cfg)?), neck: Neck::load(w, &channels, 256, &[2, 3])?, levels: 3 }) } pub fn forward(&self, pixels: &Tensor) -> Result { diff --git a/crates/sam2/src/weights.rs b/crates/sam2/src/weights.rs index cfbd6a506..ca246b143 100644 --- a/crates/sam2/src/weights.rs +++ b/crates/sam2/src/weights.rs @@ -46,6 +46,15 @@ impl Weights { Ok(Weights { map: out }) } + pub fn contains(&self, name: &str) -> bool { + self.map.contains_key(name) + } + + /// The shape of a tensor, if the checkpoint has it. + pub fn shape(&self, name: &str) -> Option> { + self.map.get(name).map(|t| t.dims().to_vec()) + } + pub fn len(&self) -> usize { self.map.len() } diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index d121df840..01648e329 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -33,6 +33,7 @@ fn load() -> Option<(Weights, HashMap)> { } let w = Weights::open(&wf).unwrap(); let r = candle_core::safetensors::load(&rf, &Device::Cpu).unwrap(); + note_size(&w); Some((w, r)) } @@ -41,6 +42,30 @@ fn max_abs(a: &Tensor, b: &Tensor) -> f32 { (a - b).unwrap().abs().unwrap().flatten_all().unwrap().max(0).unwrap().to_scalar::().unwrap() } +/// The largest magnitude in a tensor. +fn peak(t: &Tensor) -> f32 { + t.abs().unwrap().flatten_all().unwrap().max(0).unwrap().to_scalar::().unwrap() +} + +/// Set when the checkpoint is Base+ or Large (24 and 48 blocks): f32 rounding error is relative and accumulates with +/// depth, and these backbones carry much larger activations (Large peaks at 140). Tiny and Small keep the strict bounds. +static DEEP: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false); + +fn note_size(w: &Weights) { + let size = effectcraft_sam2::Size::detect(w).unwrap(); + DEEP.store(matches!(size, effectcraft_sam2::Size::BasePlus | effectcraft_sam2::Size::Large), std::sync::atomic::Ordering::Relaxed); +} + +/// `base`, or for the deep backbones the larger of `base` and 5e-4 of the activations' peak. +fn tol(base: f32, want: &Tensor) -> f32 { + if DEEP.load(std::sync::atomic::Ordering::Relaxed) { base.max(5e-4 * peak(want)) } else { base } +} + +/// `base`, ×4 for the deep backbones. +fn loose(base: f32) -> f32 { + if DEEP.load(std::sync::atomic::Ordering::Relaxed) { base * 4.0 } else { base } +} + /// f32 activations differ from PyTorch by accumulation order only; the Hiera blocks stay well /// inside this bound (measured max 2.6e-5, 2026-10-15). const HIERA_TOL: f32 = 1e-4; @@ -48,25 +73,26 @@ const HIERA_TOL: f32 = 1e-4; #[test] fn hiera_matches_pytorch_block_by_block() { let Some((w, r)) = load() else { return }; - let hiera = Hiera::load(&w, &HieraConfig::tiny()).unwrap(); + let hiera = Hiera::load(&w, &HieraConfig::for_size(effectcraft_sam2::Size::detect(&w).unwrap())).unwrap(); let pixels = r.get("input.pixel_values").unwrap(); let (emb, blocks) = hiera.forward_all(pixels).unwrap(); let e = max_abs(&emb, r.get("vision_encoder.backbone.patch_embed.0").unwrap()); eprintln!("patch_embed max abs error {e:e}"); assert!(e < 1e-4, "patch embed error {e}"); - assert_eq!(blocks.len(), 12); + assert_eq!(blocks.len(), HieraConfig::for_size(effectcraft_sam2::Size::detect(&w).unwrap()).blocks.iter().sum::()); for (i, b) in blocks.iter().enumerate() { let want = r.get(&format!("vision_encoder.backbone.blocks.{i}.0")).unwrap(); let e = max_abs(b, want); eprintln!("block {i}: max abs error {e:e}"); - assert!(e < HIERA_TOL, "block {i} error {e} (tolerance {HIERA_TOL})"); + let t = tol(HIERA_TOL, want); + assert!(e < t, "block {i} error {e} (tolerance {t}, peak {})", peak(want)); } } #[test] fn hiera_is_deterministic_and_rejects_bad_sizes() { let Some((w, r)) = load() else { return }; - let hiera = Hiera::load(&w, &HieraConfig::tiny()).unwrap(); + let hiera = Hiera::load(&w, &HieraConfig::for_size(effectcraft_sam2::Size::detect(&w).unwrap())).unwrap(); let pixels = r.get("input.pixel_values").unwrap(); let a = hiera.forward(pixels).unwrap(); let b = hiera.forward(pixels).unwrap(); @@ -85,7 +111,7 @@ const NECK_TOL: f32 = 1e-4; #[test] fn neck_and_vision_encoder_match_pytorch() { let Some((w, r)) = load() else { return }; - let enc = effectcraft_sam2::neck::VisionEncoder::load_tiny(&w).unwrap(); + let enc = effectcraft_sam2::neck::VisionEncoder::load(&w).unwrap(); let pixels = r.get("input.pixel_values").unwrap(); let stages = enc.hiera.forward(pixels).unwrap(); let out = enc.neck.forward(&stages).unwrap(); @@ -93,7 +119,8 @@ fn neck_and_vision_encoder_match_pytorch() { for (i, f) in out.features.iter().enumerate() { let e = max_abs(f, r.get(&format!("vision_encoder.neck.{i}")).unwrap()); eprintln!("neck feature {i}: max abs error {e:e}"); - assert!(e < NECK_TOL, "neck feature {i} error {e}"); + let want = r.get(&format!("vision_encoder.neck.{i}")).unwrap(); + assert!(e < tol(NECK_TOL, want), "neck feature {i} error {e} (peak {})", peak(want)); } for (i, p) in out.positions.iter().enumerate() { let e = max_abs(p, r.get(&format!("vision_encoder.neck.{}", i + 4)).unwrap()); @@ -105,7 +132,7 @@ fn neck_and_vision_encoder_match_pytorch() { let dims: Vec = feats.features.iter().map(|f| f.dim(2).unwrap()).collect(); assert_eq!(dims, vec![256, 128, 64]); let e = max_abs(feats.features.first().unwrap(), r.get("vision_encoder.neck.3").unwrap()); - assert!(e < NECK_TOL); + assert!(e < tol(NECK_TOL, r.get("vision_encoder.neck.3").unwrap())); } /// Click-prompted single image: prompt embeddings and decoder outputs against `Sam2Model`. @@ -133,7 +160,7 @@ fn prompt_encoder_and_decoder_match_pytorch() { let want_masks = r.get("mask_decoder.0").unwrap().squeeze(0).unwrap(); let e = max_abs(&out.masks, &want_masks); eprintln!("masks: max abs error {e:e} (logit range {:?})", want_masks.flatten_all().unwrap().max(0).unwrap().to_scalar::()); - assert!(e < DECODER_TOL, "mask logits error {e}"); + assert!(e < loose(DECODER_TOL), "mask logits error {e}"); let e = max_abs(&out.iou, &r.get("mask_decoder.1").unwrap().squeeze(0).unwrap()); eprintln!("iou: max abs error {e:e}"); assert!(e < 1e-4, "iou error {e}"); @@ -141,7 +168,7 @@ fn prompt_encoder_and_decoder_match_pytorch() { eprintln!("object score: max abs error {e:e}"); assert!(e < 1e-4, "object score error {e}"); let e = max_abs(&out.sam_tokens, &r.get("mask_decoder.2").unwrap().squeeze(0).unwrap()); - assert!(e < 1e-4, "sam token error {e}"); + assert!(e < loose(1e-4), "sam token error {e}"); } fn video_fixture_file() -> PathBuf { @@ -155,7 +182,9 @@ fn load_video() -> Option<(Weights, HashMap)> { eprintln!("skipped: SAM 2 weights or video reference tensors not found ({} / {})", wf.display(), rf.display()); return None; } - Some((Weights::open(&wf).unwrap(), candle_core::safetensors::load(&rf, &Device::Cpu).unwrap())) + let w = Weights::open(&wf).unwrap(); + note_size(&w); + Some((w, candle_core::safetensors::load(&rf, &Device::Cpu).unwrap())) } /// Measured max 3.5e-6 (memory features), 2.6e-5 (memory attention), 2026-10-15. @@ -642,13 +671,13 @@ fn gpu_encoder_matches_the_cpu_path_and_pytorch() { let low_cpu = cpu.features.features.get(2).unwrap(); let e = max_abs(&got.low, low_cpu); eprintln!("low level vs CPU: max abs error {e:e}"); - assert!(e < GPU_TOL, "low level error {e}"); + assert!(e < tol(GPU_TOL, low_cpu), "low level error {e} (peak {})", peak(low_cpu)); let e = max_abs(&got.low, r.get("vision_encoder.neck.1").unwrap()); eprintln!("low level vs PyTorch: max abs error {e:e}"); - assert!(e < GPU_TOL, "low level vs PyTorch error {e}"); + assert!(e < tol(GPU_TOL, low_cpu), "low level vs PyTorch error {e}"); let (e0, e1) = (max_abs(&got.high.0, &cpu.high.0), max_abs(&got.high.1, &cpu.high.1)); eprintln!("skip features vs CPU: {e0:e} / {e1:e}"); - assert!(e0 < GPU_TOL && e1 < GPU_TOL, "skip feature errors {e0} / {e1}"); + assert!(e0 < tol(GPU_TOL, &cpu.high.0) && e1 < tol(GPU_TOL, &cpu.high.1), "skip feature errors {e0} / {e1}"); } /// The GPU memory attention against PyTorch (the fixture's frame: one memory frame and four pointer tokens). diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index 3ebf964f9..fb16f01ce 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -185,6 +185,54 @@ pub const SAM2_TINY: ModelInfo = ModelInfo { sha256: "48c14467e5cf9e51870511feb72c89688e82dd74523142c0538b663e193ac2a7", }; +/// SAM 2.1 Small: the Small checkpoint, pinned to a commit like [`SAM2_TINY`]. +pub const SAM2_SMALL: ModelInfo = ModelInfo { + id: "sam2.1-hiera-small", + task: Task::Mask, + name: "SAM 2.1 Small", + description: "Meta's Segment Anything Model 2 (Hiera-Small, 46M parameters). The same streaming-memory tracker as SAM 2.1 Tiny with a larger backbone: a little more accurate on thin parts and busy backgrounds, a little slower.", + authors: "Meta AI (Nikhila Ravi et al., SAM 2), Hugging Face transformers checkpoint", + licence: "Apache-2.0", + licence_url: "https://github.com/facebookresearch/sam2/blob/main/LICENSE", + homepage: "https://github.com/facebookresearch/sam2", + url: "https://huggingface.co/facebook/sam2.1-hiera-small/resolve/ee5bba1d82bb8749febdf90f45e84b687142ba03/model.safetensors", + file_name: "sam2.1-hiera-small.safetensors", + size: 184_305_280, + sha256: "0a4067b11ce1e23d5229203f11c718a823060d15a4b23fa2372a7d4b77cbbc60", +}; + +/// SAM 2.1 Base+: the Base+ checkpoint, pinned to a commit like [`SAM2_TINY`]. +pub const SAM2_BASE_PLUS: ModelInfo = ModelInfo { + id: "sam2.1-hiera-base-plus", + task: Task::Mask, + name: "SAM 2.1 Base+", + description: "Meta's Segment Anything Model 2 (Hiera-Base+, 81M parameters). The same streaming-memory tracker as SAM 2.1 Tiny with a larger backbone: better on difficult footage; wants a GPU for comfortable speed.", + authors: "Meta AI (Nikhila Ravi et al., SAM 2), Hugging Face transformers checkpoint", + licence: "Apache-2.0", + licence_url: "https://github.com/facebookresearch/sam2/blob/main/LICENSE", + homepage: "https://github.com/facebookresearch/sam2", + url: "https://huggingface.co/facebook/sam2.1-hiera-base-plus/resolve/b7320756a13354e7530a63935656d35b2f91a290/model.safetensors", + file_name: "sam2.1-hiera-base-plus.safetensors", + size: 323_476_296, + sha256: "2012733a0de5d03efd1bba550a2847c4551be9ef2e0d497c83074df66189f780", +}; + +/// SAM 2.1 Large: the Large checkpoint, pinned to a commit like [`SAM2_TINY`]. +pub const SAM2_LARGE: ModelInfo = ModelInfo { + id: "sam2.1-hiera-large", + task: Task::Mask, + name: "SAM 2.1 Large", + description: "Meta's Segment Anything Model 2 (Hiera-Large, 224M parameters). The same streaming-memory tracker as SAM 2.1 Tiny with a larger backbone: the most accurate; about 900 MB and meant for a GPU.", + authors: "Meta AI (Nikhila Ravi et al., SAM 2), Hugging Face transformers checkpoint", + licence: "Apache-2.0", + licence_url: "https://github.com/facebookresearch/sam2/blob/main/LICENSE", + homepage: "https://github.com/facebookresearch/sam2", + url: "https://huggingface.co/facebook/sam2.1-hiera-large/resolve/665f8e2ad61cf5f53d65644ff27c8ee525124610/model.safetensors", + file_name: "sam2.1-hiera-large.safetensors", + size: 897_897_416, + sha256: "dc407dce21301fd94abb395c5099b4f2c455fdc8a8f261ac3d0ea6d4cd197230", +}; + pub const FACE_LANDMARKER: ModelInfo = ModelInfo { id: "mediapipe-face", task: Task::Face, @@ -201,7 +249,7 @@ pub const FACE_LANDMARKER: ModelInfo = ModelInfo { }; /// The trained models, besides the built-in classical engines. -pub const MODELS: &[ModelInfo] = &[MOBILE_SAM, SAM2_TINY, FACE_LANDMARKER]; +pub const MODELS: &[ModelInfo] = &[MOBILE_SAM, SAM2_TINY, SAM2_SMALL, SAM2_BASE_PLUS, SAM2_LARGE, FACE_LANDMARKER]; /// The id of the built-in classical engines (graph cut segmenter, shape-model face tracker): no /// weights, always available. @@ -353,7 +401,10 @@ mod tests { assert!(m.url.starts_with("https://") && m.licence_url.starts_with("https://") && m.homepage.starts_with("https://")); assert!(!m.authors.is_empty() && notice(m).contains(m.authors)); } - assert_eq!(models(Task::Mask).map(|m| m.id).collect::>(), ["mobilesam", "sam2.1-hiera-tiny"]); + assert_eq!( + models(Task::Mask).map(|m| m.id).collect::>(), + ["mobilesam", "sam2.1-hiera-tiny", "sam2.1-hiera-small", "sam2.1-hiera-base-plus", "sam2.1-hiera-large"] + ); assert!(SAM2_TINY.url.contains("/resolve/") && SAM2_TINY.file_name.ends_with(".safetensors"), "pinned to a commit, not a branch"); assert_eq!(models(Task::Face).map(|m| m.id).collect::>(), ["mediapipe-face"]); assert!(load("mediapipe-face", b"nope").err().is_some_and(|e| e.contains("does not match"))); diff --git a/crates/ui-egui/tests/ui_roto_span_confidence.rs b/crates/ui-egui/tests/ui_roto_span_confidence.rs index 3b3148d33..de561014a 100644 --- a/crates/ui-egui/tests/ui_roto_span_confidence.rs +++ b/crates/ui-egui/tests/ui_roto_span_confidence.rs @@ -49,7 +49,9 @@ impl ObjectTracker for Scripted { for g in req.first..req.first + req.count as i64 { let img = frame(g).ok_or("no frame")?; let (confidence, occluded) = script(g); - let prob = (0..img.w * img.h).map(|i| if !occluded && (i % img.w).abs_diff(img.w / 2) < img.w / 4 && (i / img.w).abs_diff(img.h / 2) < img.h / 4 { 1.0 } else { 0.0 }).collect(); + let prob = (0..img.w * img.h) + .map(|i| if !occluded && (i % img.w).abs_diff(img.w / 2) < img.w / 4 && (i / img.w).abs_diff(img.h / 2) < img.h / 4 { 1.0 } else { 0.0 }) + .collect(); out(0, g, TrackedFrame { prob, w: img.w, h: img.h, confidence, occluded }); } Ok(()) diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index a08c8a5e7..56044f3fb 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -29,6 +29,24 @@ strokes and span they came from are unchanged. Changing a stroke or the span mak Choosing another model, or the classic engine, leaves them in place and `roto.status` says so; a classic Propagate replaces them with classic mattes. Object Matte saves nothing. +## Sizes +Four checkpoints are in the registry (Settings ▸ Roto Brush); the size is detected from the file, and each is +pinned to a Hugging Face commit with its SHA-256. All run the same tracker; a bigger backbone costs time and memory. + +| Model | Download | GPU encode (RX 7900 GRE) | Agreement with PyTorch | +|---|---|---|---| +| SAM 2.1 Tiny (`sam2.1-hiera-tiny`) | 155.9 MB | 0.24 s | encoder 6.8e-6 | +| SAM 2.1 Small (`sam2.1-hiera-small`) | 184.3 MB | 0.20 s | encoder 1.0e-5 | +| SAM 2.1 Base+ (`sam2.1-hiera-base-plus`) | 323.5 MB | 0.37 s | encoder 2.3e-5 | +| SAM 2.1 Large (`sam2.1-hiera-large`) | 897.9 MB | 0.72 s | encoder 5.2e-4 | + +Every size was checked stage by stage against PyTorch (patch embedding, each Hiera block, neck, prompt encoder, +decoder, memory encoder, memory attention) and on a whole clip forward and backward, on the CPU and the GPU: the video +logit error is at most 0.005 with the same masks. Large carries much larger activations (peaks near 300 in the third +stage), so its f32 rounding error is larger in absolute terms; its tests use a bound relative to the peak. Large needs +the GPU to run per block (the encode submits once per block so intermediates are freed); a GPU error of any kind falls +back to the CPU. + ## Corrections re-track only what they change A correction stroke changes the keys of the frames from it on (in both directions from the base), so Propagate tracks only those again: the window starts at the nearest prompted frame at or before the first frame it needs (the corrected From ac8765bf67b42dd6b4bcbc75b3fe09cf02432848 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Sat, 17 Oct 2026 14:31:25 +0200 Subject: [PATCH 43/77] S2.65: help and bench follow all four SAM 2 sizes; tests for the size configs and detection Co-Authored-By: Claude Sonnet 5.5 --- apps/effectcraft-cli/src/main.rs | 9 +++++---- crates/engine/src/analysis/models.rs | 2 +- crates/engine/src/tests/models.rs | 6 +++++- crates/sam2/src/lib.rs | 30 ++++++++++++++++++++++++++++ docs/roto-brush-sam2.md | 2 +- 5 files changed, 42 insertions(+), 7 deletions(-) diff --git a/apps/effectcraft-cli/src/main.rs b/apps/effectcraft-cli/src/main.rs index e6a3352dc..1c528c8fb 100644 --- a/apps/effectcraft-cli/src/main.rs +++ b/apps/effectcraft-cli/src/main.rs @@ -28,7 +28,7 @@ //! effectcraft-cli bench --ops [--small] [--layers N] [--comps N] [--footage N] everyday-operation timings //! on a large generated project (open, save, auto-save, undo/redo, timeline, Project panel) //! effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt] [--fixes] [--clip FRAMES --truth MASKS] [--json] Roto Brush trackers on synthetic clips with known mattes -//! (classical engine, MobileSAM and SAM 2.1 Tiny when installed; IoU, flicker, re-appearance, seconds per frame) +//! (classical engine, MobileSAM and SAM 2.1 (any installed size) when installed; IoU, flicker, re-appearance, seconds per frame) //! effectcraft-cli script FILE.jsx [F.ecproj] | --eval CODE run an After Effects-style script //! effectcraft-cli mcp [--autosave | --bridge PORT] MCP server on stdio //! @@ -1030,11 +1030,11 @@ fn bench_roto(args: &Args) -> Result<(), Failure> { note!("bench --roto: MobileSAM is not installed; comparing the classical engine only"); } // `--sam2 FILE` uses the weights at FILE; otherwise the installed ones. - let sam2_file = args.opt("--sam2").map(std::path::PathBuf::from).or_else(|| session.model_path(&effectcraft_segment::SAM2_TINY)); + let sam2_file = args.opt("--sam2").map(std::path::PathBuf::from).or_else(|| effectcraft_sam2::mask::SIZES.iter().find_map(|m| session.model_path(m))); let sam2: Option> = sam2_file.and_then(|p| std::fs::read(p).ok()).and_then(|bytes| effectcraft_sam2::mask::load(&bytes).ok()).map(|m| m as std::sync::Arc); if sam2.is_none() { - note!("bench --roto: SAM 2.1 Tiny is not installed; no SAM 2 column"); + note!("bench --roto: no SAM 2.1 size is installed; no SAM 2 column"); } // `--fixes`: also count the correction strokes each needs before every frame reaches IoU 0.9 (up to 6). let fixes = args.flag("--fixes"); @@ -1053,7 +1053,8 @@ fn bench_roto(args: &Args) -> Result<(), Failure> { runs.extend(measure(&sc, &|s| Some(bench::run_flow(s, "mobilesam", Some(m.as_ref()))))); } if let Some(m) = &sam2 { - runs.extend(measure(&sc, &|s| bench::run_tracker(s, "sam2-tiny", m.as_ref()))); + let label = m.info().id.replacen("sam2.1-hiera-", "sam2-", 1); + runs.extend(measure(&sc, &|s| bench::run_tracker(s, &label, m.as_ref()))); } } if args.flag("--json") { diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index d6d9ef63b..a00dd660e 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -544,7 +544,7 @@ pub fn specs() -> Vec { "Use Roto Brush Model", [], None, - "{id: classical|mobilesam|sam2.1-hiera-tiny, wait?} (Roto Brush 2.0 / 3.0 use it; wait: loaded before returning)", + "{id: classical|mobilesam|sam2.1-hiera-tiny|sam2.1-hiera-small|sam2.1-hiera-base-plus|sam2.1-hiera-large, wait?} (Roto Brush 2.0 / 3.0 use it; wait: loaded before returning)", always, |s, p| select(s, p, Task::Mask) ), diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index ee05ea6b8..3dee87e74 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -37,7 +37,11 @@ fn models_are_listed_chosen_and_verified() { let list = s.execute_checked(&format!("{prefix}.models"), json!({})).unwrap(); let ids: Vec<&str> = list["models"].as_array().unwrap().iter().map(|m| m["id"].as_str().unwrap()).collect(); // Roto Brush also offers SAM 2.1 in four sizes (video trackers; see crates/sam2). - let expect: Vec<&str> = if task == Task::Mask { vec!["classical", model, "sam2.1-hiera-tiny", "sam2.1-hiera-small", "sam2.1-hiera-base-plus", "sam2.1-hiera-large"] } else { vec!["classical", model] }; + let expect: Vec<&str> = if task == Task::Mask { + vec!["classical", model, "sam2.1-hiera-tiny", "sam2.1-hiera-small", "sam2.1-hiera-base-plus", "sam2.1-hiera-large"] + } else { + vec!["classical", model] + }; assert_eq!(ids, expect); let m = &list["models"][1]; assert_eq!((m["licence"].as_str(), m["installed"].as_bool()), (Some("Apache-2.0"), Some(false))); diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 09cbb1466..f2aba0390 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -107,6 +107,36 @@ mod tests { assert!(Size::Large.repo().ends_with("large")); } + #[test] + fn configs_match_the_released_sizes() { + let blocks = |s| hiera::HieraConfig::for_size(s).blocks.iter().sum::(); + assert_eq!([Size::Tiny, Size::Small, Size::BasePlus, Size::Large].map(blocks), [12, 16, 24, 48]); + for s in [Size::Tiny, Size::Small, Size::BasePlus, Size::Large] { + let c = hiera::HieraConfig::for_size(s); + assert_eq!((c.embed_dims.len(), c.heads.len(), c.windows.len()), (4, 4, 4)); + assert!(c.global_blocks.iter().all(|g| *g < blocks(s))); + } + } + + #[test] + fn size_comes_from_the_patch_embedding_and_depth() { + use candle_core::{DType, Device, Tensor}; + let make = |width: usize, deep: bool| { + let mut m = std::collections::HashMap::new(); + m.insert("vision_encoder.backbone.patch_embed.projection.weight".to_string(), Tensor::zeros((width, 3, 7, 7), DType::F32, &Device::Cpu).unwrap()); + if deep { + m.insert("vision_encoder.backbone.blocks.15.layer_norm1.weight".to_string(), Tensor::zeros(4, DType::F32, &Device::Cpu).unwrap()); + } + weights::Weights::from_map(m).unwrap() + }; + assert_eq!(Size::detect(&make(96, false)).unwrap(), Size::Tiny); + assert_eq!(Size::detect(&make(96, true)).unwrap(), Size::Small); + assert_eq!(Size::detect(&make(112, false)).unwrap(), Size::BasePlus); + assert_eq!(Size::detect(&make(144, false)).unwrap(), Size::Large); + assert!(Size::detect(&make(64, false)).is_err()); + assert!(Size::detect(&weights::Weights::from_map(Default::default()).unwrap()).is_err()); + } + #[test] fn errors_say_what_failed() { assert!(Error::Input("empty frame".into()).to_string().contains("empty frame")); diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index 56044f3fb..b25506b2b 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -106,7 +106,7 @@ encoder's 64 × 64 level matches PyTorch to 7e-6, the memory attention to 1.8e-5 times the pieces on both devices (`EFFECTCRAFT_SAM2_PROFILE=1` prints a frame's breakdown). ## Benchmark -`effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt]` runs synthetic clips with +`effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors (any size; default: the first installed)] [--mobilesam mobile_sam.pt]` runs synthetic clips with known mattes through all three. 24 frames, debug build, SAM 2 on the GPU; mean IoU (mean IoU after the object returns), seconds per frame: From e30db424936ca0f9d2237f62f4f2dd7b7a8fa61d Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Sat, 17 Oct 2026 14:37:50 +0200 Subject: [PATCH 44/77] S2.66: Settings suggests the SAM 2 size that suits this computer Co-Authored-By: Claude Sonnet 5.5 --- crates/engine/src/analysis/models.rs | 6 ++++++ crates/sam2/src/mask.rs | 21 +++++++++++++++++++++ crates/ui-egui/src/panels/settings.rs | 3 +++ 3 files changed, 30 insertions(+) diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index a00dd660e..d166702ac 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -420,7 +420,13 @@ fn list(s: &mut Session, task: Task) -> Result { "active": st.active.is_none(), })]; for m in seg::models(task) { + #[cfg(not(target_arch = "wasm32"))] + let (recommended, hint) = (effectcraft_sam2::mask::recommended().id == m.id, effectcraft_sam2::mask::hint(m.id)); + #[cfg(target_arch = "wasm32")] + let (recommended, hint) = (false, None::); v.push(json!({ + "recommended": recommended, + "hint": hint, "id": m.id, "name": m.name, "description": m.description, diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index 251bb05ca..d840d4052 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -148,6 +148,27 @@ pub fn load(bytes: &[u8]) -> std::result::Result, String> { Ok(Arc::new(Sam2Mask::with_info(Arc::new(model), info))) } +/// The size to suggest on this computer: Small when a GPU is usable (no slower to run than Tiny there, and more accurate), +/// Tiny otherwise (the bigger backbones take several seconds a frame on a CPU). +pub fn recommended() -> &'static ModelInfo { + if crate::gpu::Gpu::shared().is_ok() { &effectcraft_segment::SAM2_SMALL } else { &SAM2_TINY } +} + +/// A one-line note for the Settings list: what suits this computer, and what the bigger sizes cost without a GPU. +pub fn hint(id: &str) -> Option { + if !SIZES.iter().any(|m| m.id == id) { + return None; + } + let gpu = crate::gpu::Gpu::shared().is_ok(); + Some(match (id == recommended().id, gpu, id) { + (true, true, _) => "Recommended: a GPU was found, and this size is no slower than Tiny on it.".to_string(), + (true, false, _) => "Recommended: no GPU was found, and this is the quickest on a CPU.".to_string(), + (false, false, "sam2.1-hiera-base-plus" | "sam2.1-hiera-large") => "Slow without a GPU (several seconds a frame); none was found.".to_string(), + (false, true, "sam2.1-hiera-large") => "For difficult footage; about 0.7 s of GPU time to look at each frame.".to_string(), + _ => return None, + }) +} + /// The registry entries this crate loads. pub const SIZES: [&ModelInfo; 4] = [&SAM2_TINY, &effectcraft_segment::SAM2_SMALL, &effectcraft_segment::SAM2_BASE_PLUS, &effectcraft_segment::SAM2_LARGE]; diff --git a/crates/ui-egui/src/panels/settings.rs b/crates/ui-egui/src/panels/settings.rs index 816b0ca01..8b4133fbc 100644 --- a/crates/ui-egui/src/panels/settings.rs +++ b/crates/ui-egui/src/panels/settings.rs @@ -489,6 +489,9 @@ fn models_ui(app: &mut EffectcraftApp, ui: &mut egui::Ui, t: &Tokens, acts: &mut ui.label(RichText::new(state).color(if active { t.accent } else { t.text_dim })); }); ui.label(RichText::new(m["description"].as_str().unwrap_or_default()).color(t.text_dim)); + if let Some(hint) = m["hint"].as_str() { + ui.label(RichText::new(hint).color(if m["recommended"].as_bool().unwrap_or(false) { t.accent } else { t.text_dim })); + } if let Some(by) = m["authors"].as_str() { ui.label(RichText::new(format!("By {by}")).color(t.text_dim)); } From 2dfe1b8dbec73c49e103351688cf439c0ee6055b Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Sat, 17 Oct 2026 14:37:50 +0200 Subject: [PATCH 45/77] S2.67: optional half-precision GPU weights (EFFECTCRAFT_SAM2_F16=1) Halves the weight memory (Large: about 900 MB to 450 MB). Off by default: the video logit error grows from about 0.003 to 0.06-0.12 (masks unchanged) for a speed gain of under 10 percent. Co-Authored-By: Claude Sonnet 5.5 --- Cargo.lock | 1 + crates/sam2/Cargo.toml | 1 + crates/sam2/src/gpu/kernels.wgsl | 67 ++++++++++++++++++++++++++++++++ crates/sam2/src/gpu/mod.rs | 48 ++++++++++++++++++++--- crates/sam2/src/gpu/tests.rs | 33 ++++++++++++++++ 5 files changed, 144 insertions(+), 6 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b4acc3a6c..b104f63cf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1836,6 +1836,7 @@ dependencies = [ "candle-core", "candle-nn", "effectcraft-segment", + "half", "log", "pollster 0.4.0", "rayon", diff --git a/crates/sam2/Cargo.toml b/crates/sam2/Cargo.toml index ff8e64c04..e45cc8def 100644 --- a/crates/sam2/Cargo.toml +++ b/crates/sam2/Cargo.toml @@ -22,3 +22,4 @@ candle-nn = "=0.9.2" wgpu = { workspace = true } pollster = { workspace = true } bytemuck = { workspace = true } +half = "2.7" diff --git a/crates/sam2/src/gpu/kernels.wgsl b/crates/sam2/src/gpu/kernels.wgsl index cce4c2a05..faea20e76 100644 --- a/crates/sam2/src/gpu/kernels.wgsl +++ b/crates/sam2/src/gpu/kernels.wgsl @@ -96,3 +96,70 @@ fn matmul(@builtin(workgroup_id) wg: vec3, @builtin(local_invocation_id) li } } } + +// The same, with B (a weight matrix) stored as half precision, two values per word. +@group(0) @binding(1) var w16: array; + +@compute @workgroup_size(16, 16) +fn matmul_w16(@builtin(workgroup_id) wg: vec3, @builtin(local_invocation_id) li: vec3) { + let m = u(0u); let n = u(1u); let k = u(2u); + let a_base = u(4u) + wg.z * u(5u); + let b_base = u(8u) + wg.z * u(9u); + let a_rs = u(6u); let a_cs = u(7u); let b_rs = u(10u); let b_cs = u(11u); + let row0 = wg.y * 64u; + let col0 = wg.x * 64u; + let tid = li.y * 16u + li.x; + let a_k_contig = a_cs == 1u; + let b_n_contig = b_cs == 1u; + var acc: array; + for (var i = 0u; i < 16u; i++) { acc[i] = 0.0; } + for (var kt = 0u; kt < k; kt += 16u) { + for (var i = 0u; i < 4u; i++) { + let idx = tid + i * 256u; + var r: u32; var c: u32; + if (a_k_contig) { r = idx / 16u; c = idx % 16u; } else { r = idx % 64u; c = idx / 64u; } + let gr = row0 + r; let gc = kt + c; + var va = 0.0; + if (gr < m && gc < k) { va = in0[a_base + gr * a_rs + gc * a_cs]; } + sa[r * 16u + c] = va; + var r2: u32; var c2: u32; + if (b_n_contig) { r2 = idx / 64u; c2 = idx % 64u; } else { r2 = idx % 16u; c2 = idx / 16u; } + let gr2 = kt + r2; let gc2 = col0 + c2; + var vb = 0.0; + if (gr2 < k && gc2 < n) { let bi = b_base + gr2 * b_rs + gc2 * b_cs; vb = unpack2x16float(w16[bi >> 1u])[bi & 1u]; } + sb[r2 * 64u + c2] = vb; + } + workgroupBarrier(); + for (var kk = 0u; kk < 16u; kk++) { + var ar: array; + var br: array; + for (var i = 0u; i < 4u; i++) { + ar[i] = sa[(li.y * 4u + i) * 16u + kk]; + br[i] = sb[kk * 64u + li.x * 4u + i]; + } + for (var i = 0u; i < 4u; i++) { + for (var j = 0u; j < 4u; j++) { acc[i * 4u + j] += ar[i] * br[j]; } + } + } + workgroupBarrier(); + } + let flags = u(15u); + let alpha = f(16u); + let c_base = u(12u) + wg.z * u(13u); + let c_rs = u(14u); + let act = (flags >> 1u) & 3u; + for (var i = 0u; i < 4u; i++) { + for (var j = 0u; j < 4u; j++) { + let gr = row0 + li.y * 4u + i; + let gc = col0 + li.x * 4u + j; + if (gr < m && gc < n) { + var v = acc[i * 4u + j] * alpha; + if ((flags & 1u) != 0u) { v += in2[gc]; } + if (act == 1u) { v = gelu_(v); } else if (act == 2u) { v = max(v, 0.0); } + let idx = c_base + gr * c_rs + gc; + if ((flags & 8u) != 0u) { v += in3[idx]; } + out0[idx] = v; + } + } + } +} diff --git a/crates/sam2/src/gpu/mod.rs b/crates/sam2/src/gpu/mod.rs index 5b091062c..152c67c50 100644 --- a/crates/sam2/src/gpu/mod.rs +++ b/crates/sam2/src/gpu/mod.rs @@ -23,6 +23,8 @@ pub struct Buf { pub(crate) buffer: wgpu::Buffer, /// Elements (f32). pub len: usize, + /// Two half-precision values per 32-bit word (a weight matrix only; `len` still counts values). + pub(crate) f16: bool, } /// One strided matrix operand: `(offset, batch stride, row stride, column stride)` in elements. @@ -89,6 +91,11 @@ fn u32_of(v: usize, what: &str) -> Result { u32::try_from(v).map_err(|_| Error::Compute(format!("{what} {v} does not fit a GPU index"))) } +/// Whether weight matrices are kept on the GPU in half precision (`EFFECTCRAFT_SAM2_F16=1`). +pub fn half_weights() -> bool { + std::env::var("EFFECTCRAFT_SAM2_F16").is_ok_and(|v| matches!(v.trim(), "1" | "on" | "true")) +} + impl Gpu { /// The process's one GPU device, created on first use. Several devices at once upset some Vulkan loaders /// (a crash when two threads create or drop them together), and one is all SAM 2 needs: models come and go, @@ -158,7 +165,36 @@ impl Gpu { if !data.is_empty() { self.queue.write_buffer(&buffer, 0, bytemuck::cast_slice(data)); } - Buf { buffer, len: data.len() } + Buf { buffer, len: data.len(), f16: false } + } + + /// A weight matrix: stored as half precision (half the memory) when `EFFECTCRAFT_SAM2_F16` asks for it and every + /// value fits, as f32 otherwise. + pub fn upload_weights(&self, data: &[f32]) -> Buf { + if !half_weights() || data.is_empty() || data.iter().any(|v| v.abs() > 60_000.0) { + return self.upload(data); + } + self.upload_f16(data) + } + + /// `data` as half precision, two values per word. + pub(crate) fn upload_f16(&self, data: &[f32]) -> Buf { + let words: Vec = data + .chunks(2) + .map(|p| { + let lo = half::f16::from_f32(p.first().copied().unwrap_or(0.0)).to_bits(); + let hi = half::f16::from_f32(p.get(1).copied().unwrap_or(0.0)).to_bits(); + u32::from(lo) | (u32::from(hi) << 16) + }) + .collect(); + let buffer = self.device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: (words.len() as u64) * 4, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + self.queue.write_buffer(&buffer, 0, bytemuck::cast_slice(&words)); + Buf { buffer, len: data.len(), f16: true } } /// A zeroed buffer of `len` elements. @@ -173,7 +209,7 @@ impl Gpu { usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST, mapped_at_creation: false, }); - Ok(Buf { buffer, len }) + Ok(Buf { buffer, len, f16: false }) } fn pipeline(&self, entry: &'static str) -> Arc { @@ -315,7 +351,7 @@ impl Gpu { } }; self.run( - "matmul", + if b.f16 { "matmul_w16" } else { "matmul" }, &[(0, a), (1, b), (2, bias), (3, c), (4, res)], ¶ms, (u32_of(op.n.div_ceil(64), "groups")?, u32_of(op.m.div_ceil(64), "groups")?, u32_of(op.batch, "batch")?), @@ -351,7 +387,7 @@ impl GLinear { /// From `{name}.weight` (`[out, in]`) and `{name}.bias`. pub fn load(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize) -> Result { Ok(GLinear { - w: g.upload(&flat(w.get(&format!("{name}.weight"), &[out, inp])?)?), + w: g.upload_weights(&flat(w.get(&format!("{name}.weight"), &[out, inp])?)?), b: g.upload(&flat(w.get(&format!("{name}.bias"), &[out])?)?), inp, out, @@ -373,13 +409,13 @@ impl GLinear { *b *= s; } } - Ok(GLinear { w: g.upload(&wt), b: g.upload(&bias), inp, out }) + Ok(GLinear { w: g.upload_weights(&wt), b: g.upload(&bias), inp, out }) } /// A 1 x 1 convolution (`[out, in, 1, 1]`) as a linear layer over channels-last tokens. pub fn load_conv1x1(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize) -> Result { Ok(GLinear { - w: g.upload(&flat(w.get(&format!("{name}.weight"), &[out, inp, 1, 1])?)?), + w: g.upload_weights(&flat(w.get(&format!("{name}.weight"), &[out, inp, 1, 1])?)?), b: g.upload(&flat(w.get(&format!("{name}.bias"), &[out])?)?), inp, out, diff --git a/crates/sam2/src/gpu/tests.rs b/crates/sam2/src/gpu/tests.rs index ed5b3528a..544df73e0 100644 --- a/crates/sam2/src/gpu/tests.rs +++ b/crates/sam2/src/gpu/tests.rs @@ -344,3 +344,36 @@ fn convolutions_and_the_layernorm_gelu_epilogue_match_candle() { .collect(); assert_close(&run_and_read(&g, &gy), &want, 1e-3, "layernorm + gelu"); } + +#[test] +fn half_precision_weights_give_the_same_product_to_half_precision() { + let Some(g) = gpu() else { return }; + // odd sizes: the last word holds one value; the weight is a linear layer's [n, k] + let (m, n, k) = (70, 91, 41); + let (av, wv) = (noise(m * k, 3), noise(n * k, 4)); + let (a, w16, w32) = (g.upload(&av), g.upload_f16(&wv), g.upload(&wv)); + assert!(w16.f16 && w16.len == n * k); + let (c16, c32) = (g.zeros(m * n).unwrap(), g.zeros(m * n).unwrap()); + let sb = Strides { offset: 0, batch: 0, row: 1, col: k }; + for (w, c) in [(&w16, &c16), (&w32, &c32)] { + g.matmul(&Matmul { + m, + n, + k, + batch: 1, + a: (&a, Strides::row_major(k)), + b: (w, sb), + bias: None, + c: (c, Strides::row_major(n)), + residual: None, + alpha: 1.0, + act: Act::None, + }) + .unwrap(); + } + g.flush().unwrap(); + let (got, want) = (g.download(&c16).unwrap(), g.download(&c32).unwrap()); + // f16 keeps 11 bits: each weight is off by at most 2^-11 relative + assert_close(&got, &want, 2e-3 * (k as f32).sqrt(), "f16 vs f32 weights"); + assert!(got != want, "the weights really were rounded to half precision"); +} From fd1d6e623f978e852faa5acda81862729624d678 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Sat, 17 Oct 2026 14:43:13 +0200 Subject: [PATCH 46/77] S2.68: document the recommended size, half-precision weights and the decoder measurements Co-Authored-By: Claude Sonnet 5.5 --- docs/roto-brush-sam2.md | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index b25506b2b..8b1931494 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -40,6 +40,20 @@ pinned to a Hugging Face commit with its SHA-256. All run the same tracker; a bi | SAM 2.1 Base+ (`sam2.1-hiera-base-plus`) | 323.5 MB | 0.37 s | encoder 2.3e-5 | | SAM 2.1 Large (`sam2.1-hiera-large`) | 897.9 MB | 0.72 s | encoder 5.2e-4 | +Settings marks the size that suits the computer: **Small** when a usable GPU is found (it encodes no slower than +Tiny there and is more accurate), **Tiny** otherwise; Base+ and Large say so when they would be slow on a CPU. The list +carries `recommended` and `hint` per model (`roto.models`). + +### Half-precision weights (optional) +`EFFECTCRAFT_SAM2_F16=1` keeps the weight matrices on the GPU as half precision, halving their memory (Large: about +900 MB to 450 MB). It is off by default: the video logit error against PyTorch grows from about 0.003 to 0.06 (Tiny) +and 0.12 (Large) with the same masks, and the speed gain is under 10% (Large encode 0.72 s to 0.66 s). + +### What stays on the CPU +The mask decoder takes about 110 ms of a tracked frame (attention layers 42 ms, the two upscaling convolutions 50 ms). +Rewriting the upscaling convolutions as a matrix product and a pixel shuffle did not make them faster, and a GPU port +would save under 0.1 s of a 0.7 s frame, so it stays on the CPU. + Every size was checked stage by stage against PyTorch (patch embedding, each Hiera block, neck, prompt encoder, decoder, memory encoder, memory attention) and on a whole clip forward and backward, on the CPU and the GPU: the video logit error is at most 0.005 with the same masks. Large carries much larger activations (peaks near 300 in the third From 08bf6e43fe5ffd86d666616211ce708ee650cede Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Sat, 17 Oct 2026 15:57:59 +0200 Subject: [PATCH 47/77] S2.69: several objects in one Roto Brush instance, tracked in one SAM 2 pass Strokes carry an object (0 for older projects, nothing written for 0). The tracker follows every object in one pass and the mattes are united; the confidence of a frame is the least sure object's. Object picker in the Layer panel's bar; roto.options/stroke/clearStrokes/status take or report objects. Co-Authored-By: Claude Sonnet 5.5 --- crates/effects/src/migrate.rs | 8 +- crates/effects/src/object_matte.rs | 8 +- crates/effects/src/roto.rs | 192 +++++++++++++++--- crates/engine/src/analysis/roto.rs | 7 + crates/engine/src/cmd/roto_cmds.rs | 27 ++- crates/engine/src/tests/roto.rs | 25 +++ crates/track/src/roto.rs | 41 ++++ crates/track/src/roto/bench.rs | 2 +- crates/track/src/roto/tests.rs | 74 ++++++- crates/ui-egui/src/panels/roto_tool.rs | 27 ++- .../ui-egui/tests/ui_roto_span_confidence.rs | 9 + docs/roto-brush-sam2.md | 17 ++ 12 files changed, 382 insertions(+), 55 deletions(-) diff --git a/crates/effects/src/migrate.rs b/crates/effects/src/migrate.rs index 35d0960e0..ecea0fe19 100644 --- a/crates/effects/src/migrate.rs +++ b/crates/effects/src/migrate.rs @@ -765,7 +765,13 @@ not a stroke", let data = effectcraft_track::roto::RotoData { base: Some(0), span: [0, 10], - strokes: vec![effectcraft_track::roto::Stroke { kind: effectcraft_track::roto::StrokeKind::Fg, frame: 0, radius: 5.0, points: vec![[2.0, 2.0]] }], + strokes: vec![effectcraft_track::roto::Stroke { + kind: effectcraft_track::roto::StrokeKind::Fg, + frame: 0, + radius: 5.0, + points: vec![[2.0, 2.0]], + object: 0, + }], }; roto.get_mut(crate::roto::STROKES).unwrap().value = Value::Str(data.to_json()); to_effect_space(spec, &mut roto, SIZE, D, &own); diff --git a/crates/effects/src/object_matte.rs b/crates/effects/src/object_matte.rs index 4b5e1fd42..3a75bec1b 100644 --- a/crates/effects/src/object_matte.rs +++ b/crates/effects/src/object_matte.rs @@ -144,7 +144,7 @@ pub fn prompt_strokes(params: &Params, layer_size: [f64; 2], frame: i64, with_mo if !inside(pt) { return None; } - let mut out = vec![Stroke { kind: StrokeKind::Fg, frame, radius: POINT_RADIUS, points: vec![pt] }]; + let mut out = vec![Stroke { kind: StrokeKind::Fg, frame, radius: POINT_RADIUS, points: vec![pt], object: 0 }]; let rect = |x0: f64, y0: f64, x1: f64, y1: f64| vec![[x0, y0], [x1, y0], [x1, y1], [x0, y1], [x0, y0]]; let boxed = b(params, "useBox").then(|| Some((v2(params, "boxTopLeft")?, v2(params, "boxBottomRight")?))).flatten(); match boxed { @@ -158,9 +158,9 @@ pub fn prompt_strokes(params: &Params, layer_size: [f64; 2], frame: i64, with_mo } out.push(if with_model { // The model takes the box itself as its box prompt. - Stroke { kind: StrokeKind::Box, frame, radius: 0.0, points: vec![[x0, y0], [x1, y1]] } + Stroke { kind: StrokeKind::Box, frame, radius: 0.0, points: vec![[x0, y0], [x1, y1]], object: 0 } } else { - Stroke { kind: StrokeKind::Bg, frame, radius: EDGE_RADIUS, points: rect(x0, y0, x1, y1) } + Stroke { kind: StrokeKind::Bg, frame, radius: EDGE_RADIUS, points: rect(x0, y0, x1, y1), object: 0 } }); } None if !with_model => { @@ -168,7 +168,7 @@ pub fn prompt_strokes(params: &Params, layer_size: [f64; 2], frame: i64, with_mo if w <= 4.0 * i || h <= 4.0 * i { return None; } - out.push(Stroke { kind: StrokeKind::Bg, frame, radius: EDGE_RADIUS, points: rect(i, i, w - i, h - i) }); + out.push(Stroke { kind: StrokeKind::Bg, frame, radius: EDGE_RADIUS, points: rect(i, i, w - i, h - i), object: 0 }); } None => {} } diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index c31c4836b..ffa41f01e 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -564,32 +564,72 @@ impl Chain { let (Some(&need_lo), Some(&need_hi)) = (needed.first(), needed.last()) else { return Ok(Some(TrackSummary::default())); }; - // A prompt is a few clicks from the strokes and a box around what the classic engine makes of - // them on that frame (clicks alone pick a part of the object, not all of it). - let mut all_prompts: Vec<(i64, effectcraft_segment::Prompt)> = vec![]; - for f in self.data.stroke_frames().into_iter().filter(|f| self.data.in_span(*f)) { - let strokes = self.data.strokes_at(f); - let p = match frame(f) { - Some(img) => rb::tracker_prompt_boxed(&strokes, self.opts.scale, &rb::segment(&img, &strokes, None, &self.opts, self.refine_radius)), - None => rb::tracker_prompt(&strokes, self.opts.scale), - }; - if p.points.iter().any(|(_, fg)| *fg) { - all_prompts.push((f, p)); + // Each object has its own prompts: a few clicks from its strokes and a box around what the classic engine + // makes of them on that frame (clicks alone pick a part of the object, not all of it). + let objects = self.data.objects(); + let mut by_object: Vec> = vec![]; + for &obj in &objects { + let mut prompts = vec![]; + for f in self.data.stroke_frames().into_iter().filter(|f| self.data.in_span(*f)) { + let strokes = self.data.strokes_of(f, obj); + let p = match frame(f) { + Some(img) => rb::tracker_prompt_boxed(&strokes, self.opts.scale, &rb::segment(&img, &strokes, None, &self.opts, self.refine_radius)), + None => rb::tracker_prompt(&strokes, self.opts.scale), + }; + if p.points.iter().any(|(_, fg)| *fg) { + prompts.push((f, p)); + } } + by_object.push(prompts); } - if all_prompts.is_empty() { + // objects without a foreground prompt cannot be followed + let followed: Vec<(u32, Vec<(i64, effectcraft_segment::Prompt)>)> = objects.into_iter().zip(by_object).filter(|(_, p)| !p.is_empty()).collect(); + if followed.is_empty() { return Err("no foreground strokes to follow".into()); } - // The window to track: from the nearest prompted frame at or before the first frame needed - // (a correction restarts the object there) to the last frame needed, reaching the first - // prompted frame when everything needed lies before it. - let first_prompt = all_prompts.first().map_or(need_lo, |p| p.0); - let lo = all_prompts.iter().map(|p| p.0).filter(|f| *f <= need_lo).max().unwrap_or(need_lo); - let hi = if need_hi < first_prompt { first_prompt } else { need_hi }; - let prompts: Vec<(i64, effectcraft_segment::Prompt)> = all_prompts.into_iter().filter(|p| p.0 >= lo && p.0 <= hi).collect(); - let req = TrackRequest::single(lo, (hi - lo + 1) as usize, prompts); + // The window to track: every object restarts at the nearest of its prompted frames at or before the first + // frame needed (a correction restarts the object there), so the window starts at the earliest restart; it + // ends at the last frame needed, reaching the farthest first prompt when an object's needed frames lie + // before its prompt. + let lo = followed.iter().map(|(_, p)| p.iter().map(|q| q.0).filter(|f| *f <= need_lo).max().unwrap_or(need_lo)).min().unwrap_or(need_lo); + let first_prompt = followed.iter().filter_map(|(_, p)| p.first().map(|q| q.0)).max().unwrap_or(need_hi); + let hi = need_hi.max(first_prompt); + let objects: Vec = followed.iter().map(|(o, _)| *o).collect(); + let prompts: Vec> = + followed.into_iter().map(|(_, p)| p.into_iter().filter(|q| q.0 >= lo && q.0 <= hi).collect()).collect(); + let req = TrackRequest { first: lo, count: (hi - lo + 1) as usize, objects: prompts }; let frame = std::cell::RefCell::new(frame); let mut summary = TrackSummary::default(); + // A frame is finished when every object has reported it (the objects' mattes are united). + let mut pending: BTreeMap>> = BTreeMap::new(); + let finish = |g: i64, parts: Vec>, summary: &mut TrackSummary, progress: &mut dyn FnMut(usize)| { + let (Some(key), Some(img)) = (self.keys.get(&g).copied(), (frame.borrow_mut())(g)) else { return }; + let mut segs = vec![]; + let (mut confidence, mut present) = (f32::INFINITY, false); + let mut worst = f32::INFINITY; + for (obj, t) in objects.iter().zip(parts) { + let Some(t) = t else { continue }; + let strokes = self.data.strokes_of(g, *obj); + segs.push(rb::segment_tracked(&img, &strokes, t.prob, t.occluded, &self.opts, self.refine_radius)); + worst = worst.min(t.confidence); + if !t.occluded { + present = true; + confidence = confidence.min(t.confidence); + } + } + if segs.is_empty() { + return; + } + // a frame is occluded when every object is absent; otherwise the doubtful object sets the confidence + let (occluded, confidence) = (!present, if present { confidence } else { worst }); + store(tracked_key(key), self.opts.scale, Arc::new(effectcraft_track::roto::FrameSeg::union(segs))); + trust_cache().lock().unwrap_or_else(|e| e.into_inner()).insert(tracked_key(key), FrameTrust { confidence, occluded }); + summary.frames += 1; + summary.occluded += usize::from(occluded); + summary.min_confidence = if summary.frames == 1 { confidence } else { summary.min_confidence.min(confidence) }; + progress(summary.frames); + }; + let n_objects = objects.len(); tracker.track( &req, &mut |g| { @@ -597,18 +637,21 @@ impl Chain { Some(Arc::new(ClipFrame { rgb: rb::rgb_of(&img), w: img.width as usize, h: img.height as usize })) }, cancel, - &mut |_, g, t| { - let (Some(key), Some(img)) = (self.keys.get(&g).copied(), (frame.borrow_mut())(g)) else { return }; - let strokes = self.data.strokes_at(g); - let seg = rb::segment_tracked(&img, &strokes, t.prob, t.occluded, &self.opts, self.refine_radius); - store(tracked_key(key), self.opts.scale, Arc::new(seg)); - trust_cache().lock().unwrap_or_else(|e| e.into_inner()).insert(tracked_key(key), FrameTrust { confidence: t.confidence, occluded: t.occluded }); - summary.frames += 1; - summary.occluded += usize::from(t.occluded); - summary.min_confidence = if summary.frames == 1 { t.confidence } else { summary.min_confidence.min(t.confidence) }; - progress(summary.frames); + &mut |o, g, t| { + let parts = pending.entry(g).or_insert_with(|| vec![None; n_objects]); + if let Some(slot) = parts.get_mut(o) { + *slot = Some(t); + } + if parts.iter().all(Option::is_some) { + let parts = pending.remove(&g).unwrap_or_default(); + finish(g, parts, &mut summary, progress); + } }, )?; + // frames some object never reported (it was not followed there): united from the rest + for (g, parts) in std::mem::take(&mut pending) { + finish(g, parts, &mut summary, progress); + } Ok(Some(summary)) } } @@ -1038,8 +1081,8 @@ mod tests { #[test] fn renders_the_segmented_disk_on_demand() { let mut d = RotoData::default(); - d.add(Stroke { kind: StrokeKind::Fg, frame: 0, radius: 3.0, points: vec![[35.0, 40.0], [45.0, 40.0]] }, 20, [0, 24]); - d.add(Stroke { kind: StrokeKind::Bg, frame: 0, radius: 3.0, points: vec![[5.0, 5.0], [95.0, 5.0], [95.0, 75.0]] }, 20, [0, 24]); + d.add(Stroke { kind: StrokeKind::Fg, frame: 0, radius: 3.0, points: vec![[35.0, 40.0], [45.0, 40.0]], object: 0 }, 20, [0, 24]); + d.add(Stroke { kind: StrokeKind::Bg, frame: 0, radius: 3.0, points: vec![[5.0, 5.0], [95.0, 5.0], [95.0, 75.0]], object: 0 }, 20, [0, 24]); let params = params_with(&d); let host = Seq; let run = |t: f64| { @@ -1083,6 +1126,91 @@ mod tests { assert!(o.data.iter().all(|p| p[3] == 0.0)); } + /// A stand-in tracker for two objects: object 0 is the moving disk, object 1 a fixed square. It records how many + /// objects (and prompts per object) it was asked for. + struct TwoObjects(std::sync::Mutex>>); + impl effectcraft_segment::MaskModel for TwoObjects { + fn info(&self) -> &'static effectcraft_segment::ModelInfo { + &effectcraft_segment::SAM2_TINY + } + fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &effectcraft_segment::Prompt) -> effectcraft_segment::Result> { + Err("no weights".into()) + } + fn tracker(&self) -> Option<&dyn effectcraft_segment::ObjectTracker> { + Some(self) + } + } + impl effectcraft_segment::ObjectTracker for TwoObjects { + fn track( + &self, + req: &effectcraft_segment::TrackRequest, + frame: &mut dyn FnMut(i64) -> Option>, + _: &dyn Fn() -> bool, + out: &mut dyn FnMut(usize, i64, effectcraft_segment::TrackedFrame), + ) -> effectcraft_segment::Result<()> { + self.0.lock().unwrap().push(req.objects.iter().map(Vec::len).collect()); + for g in req.first..req.first + req.count as i64 { + let img = frame(g).ok_or("no frame")?; + let c = [40.0 + 2.0 * g as f64, 40.0]; + // the objects report in reverse order on odd frames: the union must not depend on it + let order: Vec = if g % 2 == 1 { (0..req.objects.len()).rev().collect() } else { (0..req.objects.len()).collect() }; + for o in order { + let prob = (0..img.w * img.h) + .map(|i| { + let (x, y) = ((i % img.w) as f64 + 0.5, (i / img.w) as f64 + 0.5); + let inside = if o == 0 { (x - c[0]).hypot(y - c[1]) <= 15.0 } else { (8.0..22.0).contains(&x) && (8.0..22.0).contains(&y) }; + f32::from(inside) + }) + .collect(); + out(o, g, effectcraft_segment::TrackedFrame { prob, w: img.w, h: img.h, confidence: if o == 0 { 0.9 } else { 0.4 }, occluded: false }); + } + } + Ok(()) + } + } + + #[test] + fn objects_are_followed_in_one_pass_and_their_mattes_united() { + let mut d = RotoData::default(); + let stroke = |kind, points: Vec<[f64; 2]>, object| Stroke { kind, frame: 0, radius: 3.0, points, object }; + d.add(stroke(StrokeKind::Fg, vec![[35.0, 40.0], [45.0, 40.0]], 0), 8, [0, 24]); + d.add(stroke(StrokeKind::Fg, vec![[12.0, 15.0], [18.0, 15.0]], 1), 8, [0, 24]); + d.add(stroke(StrokeKind::Bg, vec![[85.0, 70.0], [95.0, 70.0]], 0), 8, [0, 24]); + assert_eq!(d.objects(), vec![0, 1]); + // the Refine Edge strokes belong to every object; the others only to their own + assert_eq!(d.strokes_of(0, 1).len(), 1); + assert_eq!(d.strokes_of(0, 0).len(), 2); + let params = params_with(&d); + let mut chain = Chain::new(¶ms, [100.0, 80.0], 25.0, 1.0); + let fake = Arc::new(TwoObjects(Default::default())); + chain.model = Some(fake.clone()); + // the clip: the moving disk and a fixed yellow square + let mut get = |f: i64| { + let mut img = frame(f as f64 / 25.0 + 0.001); + for y in 8..22 { + for x in 8..22 { + img.set(x, y, [0.95, 0.9, 0.1 + ((x * 3 + y * 5) % 7) as f32 / 70.0, 1.0]); + } + } + Some(Arc::new(img)) + }; + let mut seen = 0; + let summary = chain.track(&mut get, &|| false, &mut |n| seen = n).unwrap().unwrap(); + // one request for both objects; object 0 has its foreground and background prompt on the frame, object 1 its own + assert_eq!(fake.0.lock().unwrap().len(), 1); + assert_eq!(fake.0.lock().unwrap()[0], vec![1, 1]); + assert_eq!(summary.frames, 9); + assert_eq!(seen, 9); + // the least sure object sets the frame's confidence + assert!((summary.min_confidence - 0.4).abs() < 1e-6, "{}", summary.min_confidence); + assert_eq!(summary.occluded, 0); + let seg = cached(tracked_key(chain.keys[&3]), 1.0).expect("frame 3 is tracked"); + let at = |x: usize, y: usize| seg.matte[y * 100 + x]; + assert!(at(46, 40) != 0, "the moving disk (frame 3 centre x = 46)"); + assert!(at(15, 15) != 0, "the fixed square"); + assert_eq!(at(60, 70), 0, "nothing elsewhere"); + } + /// Instances saved before the propagation controls got their own twirl-down (and motion /// blur / decontamination their nested ones) load with their values in the new places. #[test] diff --git a/crates/engine/src/analysis/roto.rs b/crates/engine/src/analysis/roto.rs index 9cc950e07..626de3030 100644 --- a/crates/engine/src/analysis/roto.rs +++ b/crates/engine/src/analysis/roto.rs @@ -51,6 +51,9 @@ pub struct RotoOptions { pub boundary_color: [f64; 3], /// Start propagating in the background after each stroke. pub auto_propagate: bool, + /// The object new strokes belong to (0 is the first). With a tracker every object is followed on its own + /// and their mattes are united; the classic engine treats all foreground strokes as one object. + pub object: u32, } impl Default for RotoOptions { @@ -63,10 +66,14 @@ impl Default for RotoOptions { overlay_opacity: 50.0, boundary_color: [1.0, 0.25, 0.75], auto_propagate: true, + object: 0, } } } +/// The highest object id (SAM 2 follows up to 32 objects in one pass). +pub const MAX_OBJECT_ID: u32 = 31; + pub const VIEWS: [&str; 4] = ["alphaBoundary", "alpha", "alphaOverlay", "none"]; /// What a job does. diff --git a/crates/engine/src/cmd/roto_cmds.rs b/crates/engine/src/cmd/roto_cmds.rs index 5849fa2c8..5124d413d 100644 --- a/crates/engine/src/cmd/roto_cmds.rs +++ b/crates/engine/src/cmd/roto_cmds.rs @@ -27,7 +27,7 @@ pub fn specs() -> Vec { "Roto Brush Stroke", [], None, - "{layer?, kind?: fg|bg|refine|refineErase, points: [[x, y], …] (layer pixels), frame? (layer frame; default current), radius? (layer px; default the tool's diameter / 2), effect?}", + "{layer?, kind?: fg|bg|refine|refineErase, points: [[x, y], …] (layer pixels), frame? (layer frame; default current), radius? (layer px; default the tool's diameter / 2), effect?, object? (0–31; default the current object: with a tracker each object is followed on its own and the mattes are united)}", has_comp, stroke ), @@ -43,14 +43,22 @@ pub fn specs() -> Vec { cmd!("roto.span", "Set Segmentation Span", [], None, "{layer?, effect?, start?, end? (layer frames)}", has_comp, span), cmd!("roto.freeze", "Freeze", [], None, "{layer?, effect?, wait?}", has_comp, freeze), cmd!("roto.unfreeze", "Unfreeze", [], None, "{layer?, effect?}", has_comp, unfreeze), - cmd!("roto.clearStrokes", "Remove Roto Brush Strokes", [], None, "{layer?, effect?, frame? (only this layer frame), kind?}", has_comp, clear_strokes), + cmd!( + "roto.clearStrokes", + "Remove Roto Brush Strokes", + [], + None, + "{layer?, effect?, frame? (only this layer frame), kind?, object? (only this object's strokes)}", + has_comp, + clear_strokes + ), cmd!("roto.cancel", "Stop Roto Brush", [], None, "{}", is_running, |s, _| Ok(json!({"stopped": s.stop_roto()}))), cmd!( "roto.options", "Roto Brush Options", [], None, - "{diameter?, refineDiameter?, view?: alphaBoundary|alpha|alphaOverlay|none, overlayColor?: [r,g,b], overlayOpacity?, boundaryColor?: [r,g,b], autoPropagate?}", + "{diameter?, refineDiameter?, view?: alphaBoundary|alpha|alphaOverlay|none, overlayColor?: [r,g,b], overlayOpacity?, boundaryColor?: [r,g,b], autoPropagate?, object? (the object new strokes belong to, 0–31)}", always, options ), @@ -232,7 +240,8 @@ fn stroke(s: &mut Session, p: &Value) -> Result { if existing.is_none() && !kind.is_segmentation() { return Err(bad(cmd, "draw a foreground stroke with the Roto Brush tool first")); } - let st = Stroke { kind, frame, radius, points }; + let object = f_p(p, "object").map_or(s.state.roto.object, |v| v.clamp(0.0, rt::MAX_OBJECT_ID as f64).round() as u32); + let st = Stroke { kind, frame, radius, points, object }; let label = if kind.is_segmentation() { "Roto Brush Stroke" } else { "Refine Edge Stroke" }; let (uid, d) = s.edit(label, None, |proj, _| { let mut next = proj.next_id; @@ -263,7 +272,7 @@ fn stroke(s: &mut Session, p: &Value) -> Result { s.state.selected_layers = vec![lid]; } Ok(json!({ - "layer": lid.0, "effect": uid, "frame": frame, "kind": kind, "radius": radius, + "layer": lid.0, "effect": uid, "frame": frame, "kind": kind, "radius": radius, "object": object, "base": d.base, "span": d.span, "strokes": d.strokes.len(), })) } @@ -369,8 +378,9 @@ fn clear_strokes(s: &mut Session, p: &Value) -> Result { Some(k) => Some(StrokeKind::from_name(k).ok_or_else(|| bad(cmd, format!("unknown kind `{k}`")))?), None => None, }; + let object = f_p(p, "object").map(|v| v.round() as u32); let n0 = d.strokes.len(); - d.strokes.retain(|st| !(frame.is_none_or(|f| st.frame == f) && kind.is_none_or(|k| st.kind == k))); + d.strokes.retain(|st| !(frame.is_none_or(|f| st.frame == f) && kind.is_none_or(|k| st.kind == k) && object.is_none_or(|o| st.object == o))); if d.strokes.iter().all(|s| !s.kind.is_segmentation()) { d.base = None; } else if d.base.is_some_and(|b| !d.strokes.iter().any(|s| s.frame == b && s.kind.is_segmentation())) { @@ -412,6 +422,9 @@ fn options(s: &mut Session, p: &Value) -> Result { if let Some(b) = b_p(p, "autoPropagate") { o.auto_propagate = b; } + if let Some(v) = f_p(p, "object") { + o.object = v.clamp(0.0, rt::MAX_OBJECT_ID as f64).round() as u32; + } s.state.roto = o.clone(); Ok(serde_json::to_value(o).unwrap_or(Value::Null)) } @@ -430,6 +443,8 @@ fn status(s: &mut Session, p: &Value) -> Result { out["base"] = json!(d.base); out["span"] = json!(d.span); out["strokes"] = json!(d.strokes.len()); + out["objects"] = json!(d.objects().iter().map(|o| json!({"id": o, "strokes": d.strokes.iter().filter(|s| s.object == *o).count()})).collect::>()); + out["currentObject"] = json!(s.state.roto.object); out["strokeFrames"] = json!(d.stroke_frames()); out["frozen"] = json!(frozen_of(&g)); out["pending"] = json!(s.roto_pending.contains(&(cid, lid, uid))); diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index 66baee148..e4db67647 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -1007,3 +1007,28 @@ fn a_correction_re_tracks_only_the_frames_it_changes() { assert_eq!(trust.len(), 21); assert!(trust.values().all(|t| t.confidence > 0.9 && !t.occluded)); } + +#[test] +fn strokes_belong_to_the_current_object() { + let mut s = Session::default(); + s.execute("comp.new", json!({"width": 64, "height": 64, "frameRate": 25, "duration": 1})).unwrap(); + let layer = s.execute("layer.newSolid", json!({"width": 64, "height": 64})).unwrap()["layer"].as_u64().unwrap(); + s.execute("roto.stroke", json!({"layer": layer, "points": [[10, 10]], "radius": 4})).unwrap(); + // the current object is an option shared by the tool and the commands + s.execute("roto.options", json!({"object": 1})).unwrap(); + s.execute("roto.stroke", json!({"layer": layer, "points": [[50, 50]], "radius": 4})).unwrap(); + s.execute("roto.stroke", json!({"layer": layer, "kind": "bg", "points": [[30, 30]], "radius": 4})).unwrap(); + // an explicit object wins over the option, and is clamped to the supported range + s.execute("roto.stroke", json!({"layer": layer, "points": [[20, 50]], "radius": 4, "object": 99})).unwrap(); + let st = s.execute("roto.status", json!({"layer": layer})).unwrap(); + assert_eq!(st["currentObject"], 1); + assert_eq!(st["objects"], json!([{"id": 0, "strokes": 1}, {"id": 1, "strokes": 2}, {"id": 31, "strokes": 1}])); + // removing one object's strokes leaves the others + let r = s.execute("roto.clearStrokes", json!({"layer": layer, "object": 1})).unwrap(); + assert_eq!(r["removed"], 2); + let st = s.execute("roto.status", json!({"layer": layer})).unwrap(); + assert_eq!(st["objects"], json!([{"id": 0, "strokes": 1}, {"id": 31, "strokes": 1}])); + // the object survives a save and a load + let json = s.project.to_json(); + assert!(json.contains("\\\"object\\\":31"), "the project file carries the object"); +} diff --git a/crates/track/src/roto.rs b/crates/track/src/roto.rs index 000f7b55f..18c6eaa77 100644 --- a/crates/track/src/roto.rs +++ b/crates/track/src/roto.rs @@ -88,6 +88,15 @@ pub struct Stroke { pub frame: i64, pub radius: f64, pub points: Vec<[f64; 2]>, + /// Which object of the instance the stroke belongs to (0 for the first, and for every older project). + /// With a tracker each object is followed on its own and the mattes are combined; the classic engine + /// treats all foreground strokes as one object. + #[serde(default, skip_serializing_if = "is_zero")] + pub object: u32, +} + +fn is_zero(v: &u32) -> bool { + *v == 0 } /// A Roto Brush instance's user data: strokes, the base frame and the segmentation span @@ -130,6 +139,18 @@ impl RotoData { v.dedup(); v } + /// The objects that have a foreground stroke, in order. + pub fn objects(&self) -> Vec { + let mut v: Vec = self.strokes.iter().filter(|s| s.kind == StrokeKind::Fg).map(|s| s.object).collect(); + v.sort_unstable(); + v.dedup(); + v + } + /// The strokes of `object` on `frame`: its own foreground, background and box strokes, and the Refine Edge + /// strokes (which belong to the whole matte). + pub fn strokes_of(&self, frame: i64, object: u32) -> Vec<&Stroke> { + self.strokes.iter().filter(|s| s.frame == frame && (s.object == object || matches!(s.kind, StrokeKind::Refine | StrokeKind::RefineErase))).collect() + } pub fn in_span(&self, f: i64) -> bool { self.base.is_some() && f >= self.span[0] && f <= self.span[1] } @@ -224,6 +245,26 @@ pub struct FrameSeg { } impl FrameSeg { + /// The union of several objects' segmentations of one frame (the largest keeps its colour models). The + /// parts must share one size; any that does not is ignored. + pub fn union(parts: Vec) -> FrameSeg { + let Some(first) = parts.first() else { return FrameSeg::default() }; + let (w, h) = (first.w, first.h); + let mut parts: Vec = parts.into_iter().filter(|p| p.w == w && p.h == h && p.matte.len() == w * h).collect(); + parts.sort_by_key(|p| std::cmp::Reverse(p.area())); + let mut it = parts.into_iter(); + let Some(mut out) = it.next() else { return FrameSeg::default() }; + for p in it { + for (a, b) in out.matte.iter_mut().zip(&p.matte) { + *a = (*a).max(*b); + } + for (a, b) in out.refine.iter_mut().zip(&p.refine) { + *a = (*a).max(*b); + } + } + out + } + pub fn empty(w: usize, h: usize) -> FrameSeg { let Some(n) = w.checked_mul(h).filter(|n| *n <= rle::MAX_PLANE_BYTES) else { return FrameSeg::default(); diff --git a/crates/track/src/roto/bench.rs b/crates/track/src/roto/bench.rs index 2dafb9f3f..c086bfd57 100644 --- a/crates/track/src/roto/bench.rs +++ b/crates/track/src/roto/bench.rs @@ -122,7 +122,7 @@ fn line(kind: StrokeKind, frame: i64, a: [f64; 2], b: [f64; 2], radius: f64) -> let t = f64::from(k) / 10.0; [a[0] + (b[0] - a[0]) * t, a[1] + (b[1] - a[1]) * t] }; - Stroke { kind, frame, radius, points: (0..=10).map(pt).collect() } + Stroke { kind, frame, radius, points: (0..=10).map(pt).collect(), object: 0 } } /// The strokes an artist would paint on the base frame: one through the disk, two around it. diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index b7d8c7809..50ccaeaa4 100644 --- a/crates/track/src/roto/tests.rs +++ b/crates/track/src/roto/tests.rs @@ -65,7 +65,13 @@ pub(crate) fn frame_at(c: [f64; 2]) -> (Image, Vec) { } fn line(kind: StrokeKind, frame: i64, a: [f64; 2], b: [f64; 2], radius: f64) -> Stroke { - Stroke { kind, frame, radius, points: (0..=10).map(|k| [a[0] + (b[0] - a[0]) * k as f64 / 10.0, a[1] + (b[1] - a[1]) * k as f64 / 10.0]).collect() } + Stroke { + kind, + frame, + radius, + points: (0..=10).map(|k| [a[0] + (b[0] - a[0]) * k as f64 / 10.0, a[1] + (b[1] - a[1]) * k as f64 / 10.0]).collect(), + object: 0, + } } /// The strokes an artist would paint on frame `f`: one through the disk, one around it. @@ -231,6 +237,7 @@ fn ring(frame: i64, c: [f64; 2], r: f64, radius: f64) -> Stroke { frame, radius, points: (0..=48).map(|k| k as f64 / 48.0 * std::f64::consts::TAU).map(|a| [c[0] + r * a.cos(), c[1] + r * a.sin()]).collect(), + object: 0, } } @@ -295,7 +302,7 @@ fn later_refine_strokes_override_earlier_erases() { let (img, _) = soft_frame(c, 10.0, [0.9, 0.25, 0.1], [0.1, 0.3, 0.9]); let spot = [c[0] + R, c[1]]; let i = spot[1] as usize * W as usize + spot[0] as usize; - let erase = Stroke { kind: StrokeKind::RefineErase, frame: 0, radius: 6.0, points: vec![spot] }; + let erase = Stroke { kind: StrokeKind::RefineErase, frame: 0, radius: 6.0, points: vec![spot], object: 0 }; let refine = Stroke { kind: StrokeKind::Refine, ..erase.clone() }; let band = |st: &[Stroke]| segment(&img, &st.iter().collect::>(), None, &SegOpts::default(), 9.0).refine[i]; let mut st = vec![ @@ -389,7 +396,7 @@ fn roto_perf_1080p() { let mut seg2 = next.clone(); let c = center(1); let ring: Vec<[f64; 2]> = (0..=64).map(|k| k as f64 / 64.0 * std::f64::consts::TAU).map(|t| [(c[0] + R * t.cos()) * s, (c[1] + R * t.sin()) * s]).collect(); - let st = [Stroke { kind: StrokeKind::Refine, frame: 1, radius: 20.0, points: ring }]; + let st = [Stroke { kind: StrokeKind::Refine, frame: 1, radius: 20.0, points: ring, object: 0 }]; let r = rasterize(&st.iter().collect::>(), &[StrokeKind::Refine], 1.0, 1920, 1080); seg2.refine = r.iter().map(|v| *v as u8).collect(); let rgb = rgb_of(&b); @@ -500,8 +507,8 @@ fn mobilesam_segments_and_propagates_the_disk() { /// ignores it; a prior's box does not replace it. #[test] fn box_stroke_is_the_model_box_prompt() { - let fg = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 2.0, points: vec![[50.0, 40.0]] }; - let bx = Stroke { kind: StrokeKind::Box, frame: 0, radius: 0.0, points: vec![[80.0, 60.0], [20.0, 10.0]] }; + let fg = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 2.0, points: vec![[50.0, 40.0]], object: 0 }; + let bx = Stroke { kind: StrokeKind::Box, frame: 0, radius: 0.0, points: vec![[80.0, 60.0], [20.0, 10.0]], object: 0 }; let p = stroke_prompt(&[&fg, &bx], &[], 0.5, 0, 0); assert_eq!(p.bbox, Some([10.0, 5.0, 40.0, 30.0])); assert_eq!(p.points, vec![([25.0, 20.0], true)]); @@ -523,14 +530,14 @@ fn box_stroke_is_the_model_box_prompt() { /// out what this frame's background strokes removed. #[test] fn model_prompts_follow_the_latest_stroke() { - let green = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 4.0, points: vec![[50.0, 40.0]] }; + let green = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 4.0, points: vec![[50.0, 40.0]], object: 0 }; let red = Stroke { kind: StrokeKind::Bg, ..green.clone() }; let (w, h) = (100, 80); let prompt = |st: &[&Stroke]| stroke_prompt(st, &resolve(st, StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h), 1.0, w, h).points; assert_eq!(prompt(&[&green, &red]), vec![([50.0, 40.0], false)]); assert_eq!(prompt(&[&green, &red, &green]), vec![([50.0, 40.0], true); 2]); let soft: Vec = (0..w * h).map(|i| ((i % w) < 50) as u8 as f32).collect(); - let erase = Stroke { kind: StrokeKind::Bg, frame: 0, radius: 15.0, points: vec![[25.0, 40.0]] }; + let erase = Stroke { kind: StrokeKind::Bg, frame: 0, radius: 15.0, points: vec![[25.0, 40.0]], object: 0 }; let hard = resolve(&[&erase], StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h); let mut q = Prompt::default(); assert!(add_prior_prompt(&mut q, &soft, &hard, w, h)); @@ -561,7 +568,7 @@ impl effectcraft_segment::MaskModel for CountingModel { fn a_correction_keeps_the_model_prompted_with_what_is_left() { let (img, gt, rect) = bar_frame(); let (w, h) = (img.width as usize, img.height as usize); - let dot = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 15.0, points: vec![[80.0, 50.0]] }; + let dot = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 15.0, points: vec![[80.0, 50.0]], object: 0 }; let tap = Stroke { kind: StrokeKind::Bg, radius: 2.0, ..dot.clone() }; let refs = [&dot, &tap]; let model = CountingModel(rect, Default::default()); @@ -588,7 +595,7 @@ fn tracker_prompts_carry_a_handful_of_points() { let (fg, bg) = (p.points.iter().filter(|(_, f)| *f).count(), p.points.iter().filter(|(_, f)| !*f).count()); assert!((1..=TRACKER_FG_POINTS).contains(&fg) && (1..=TRACKER_BG_POINTS).contains(&bg), "{fg} fg, {bg} bg"); // a long, finely sampled stroke still gives at most the limit, and an empty one gives none - let long = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 2.0, points: (0..400).map(|i| [i as f64 * 0.4, 50.0]).collect() }; + let long = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 2.0, points: (0..400).map(|i| [i as f64 * 0.4, 50.0]).collect(), object: 0 }; assert!(tracker_prompt(&[&long], 1.0).points.len() <= TRACKER_FG_POINTS); assert!(tracker_prompt(&[], 1.0).points.is_empty()); } @@ -630,8 +637,55 @@ fn tracker_prompts_carry_the_box_of_the_classic_matte() { // no matte, no box; a box the user drew wins let none = tracker_prompt_boxed(&refs, 1.0, &FrameSeg::empty(W as usize, H as usize)); assert_eq!(none.bbox, None); - let bx = Stroke { kind: StrokeKind::Box, frame: 0, radius: 0.0, points: vec![[10.0, 10.0], [50.0, 60.0]] }; + let bx = Stroke { kind: StrokeKind::Box, frame: 0, radius: 0.0, points: vec![[10.0, 10.0], [50.0, 60.0]], object: 0 }; let mut with_box: Vec<&Stroke> = refs.clone(); with_box.push(&bx); assert_eq!(tracker_prompt_boxed(&with_box, 1.0, &classic).bbox, Some([10.0, 10.0, 50.0, 60.0])); } + +#[test] +fn strokes_carry_an_object_and_old_data_is_object_zero() { + // saved before objects existed: no field, and the field is not written for object 0 + let old = r#"{"base":0,"span":[0,4],"strokes":[{"kind":"fg","frame":0,"radius":3.0,"points":[[1.0,2.0]]}]}"#; + let d = RotoData::from_json(old); + assert_eq!(d.strokes.first().map(|s| s.object), Some(0)); + assert!(!d.to_json().contains("object"), "{}", d.to_json()); + // the keys of existing projects do not change + let same = RotoData::from_json(&d.to_json()); + assert_eq!(d.chain_keys(1), same.chain_keys(1)); + + let mut d = d; + d.add(Stroke { kind: StrokeKind::Fg, frame: 0, radius: 3.0, points: vec![[9.0, 9.0]], object: 2 }, 4, [0, 10]); + d.add(Stroke { kind: StrokeKind::Bg, frame: 0, radius: 3.0, points: vec![[0.0, 0.0]], object: 2 }, 4, [0, 10]); + d.add(Stroke { kind: StrokeKind::Refine, frame: 0, radius: 3.0, points: vec![[5.0, 5.0]], object: 0 }, 4, [0, 10]); + assert_eq!(d.objects(), vec![0, 2]); + assert!(d.to_json().contains("\"object\":2")); + assert_eq!(RotoData::from_json(&d.to_json()), d); + // an object's own strokes, plus the Refine Edge strokes that belong to the whole matte + assert_eq!(d.strokes_of(0, 2).len(), 3); + assert_eq!(d.strokes_of(0, 0).len(), 2); + assert!(d.strokes_of(0, 1).len() == 1); + // moving a stroke to another object changes the frame keys (the tracked frames are stale) + let before = d.chain_keys(1); + if let Some(s) = d.strokes.first_mut() { + s.object = 1; + } + assert_ne!(before, d.chain_keys(1)); +} + +#[test] +fn the_union_of_objects_is_their_largest_extent() { + let seg = |cells: &[usize]| { + let mut s = FrameSeg::empty(4, 1); + for c in cells { + s.matte[*c] = 255; + } + s + }; + let u = FrameSeg::union(vec![seg(&[0]), seg(&[2, 3]), seg(&[3])]); + assert_eq!(u.matte, vec![255, 0, 255, 255]); + assert_eq!(FrameSeg::union(vec![]).w, 0); + // a part of another size is ignored + let odd = FrameSeg::empty(2, 2); + assert_eq!(FrameSeg::union(vec![seg(&[1]), odd]).matte, vec![0, 255, 0, 0]); +} diff --git a/crates/ui-egui/src/panels/roto_tool.rs b/crates/ui-egui/src/panels/roto_tool.rs index f8cd383c6..bf47b558d 100644 --- a/crates/ui-egui/src/panels/roto_tool.rs +++ b/crates/ui-egui/src/panels/roto_tool.rs @@ -8,7 +8,7 @@ use effectcraft_engine::render::{EvalCtx, RenderOpts, Renderer}; use egui::{Align2, Color32, Pos2, Rect, Sense, Stroke, StrokeKind, pos2, vec2}; use serde_json::json; -use crate::i18n::tr; +use crate::i18n::{tr, tr_args}; use crate::theme::Tokens; use crate::{EffectcraftApp, widgets}; @@ -373,6 +373,31 @@ pub fn bar_buttons(app: &mut EffectcraftApp, ui: &mut egui::Ui, p: &egui::Painte } app.auto.add("layerPanel.roto.freeze", r, label); x += 80.0; + // Objects: strokes belong to the current object; with a tracker each is followed on its own and the mattes are united. + let objects = fx::data(¶ms).objects(); + let current = app.session.state.roto.object; + let known = objects.contains(¤t); + let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(86.0, 20.0)); + let label = tr_args("Object {}", &[&(current + 1)]); + if widgets::text_button(ui, r, &label, false, &t, egui::Id::new(("roto-object", uid))).clicked() { + // cycle through the objects that have strokes (and the empty one being drawn) + let mut all = objects.clone(); + if !known { + all.push(current); + all.sort_unstable(); + } + let next = all.iter().copied().find(|o| *o > current).or_else(|| all.first().copied()).unwrap_or(0); + let _ = app.session.execute("roto.options", json!({"object": next})); + } + app.auto.add("layerPanel.roto.object", r, &label); + x += 90.0; + let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(26.0, 20.0)); + if widgets::text_button(ui, r, "+", false, &t, egui::Id::new(("roto-object-new", uid))).clicked() { + let next = objects.iter().copied().chain([current]).max().map_or(0, |m| m + 1); + let _ = app.session.execute("roto.options", json!({"object": next})); + } + app.auto.add("layerPanel.roto.newObject", r, "New object"); + x += 34.0; if running { let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(60.0, 20.0)); if widgets::text_button(ui, r, tr("Stop"), false, &t, egui::Id::unique("roto-stop")).clicked() { diff --git a/crates/ui-egui/tests/ui_roto_span_confidence.rs b/crates/ui-egui/tests/ui_roto_span_confidence.rs index de561014a..f3a8bcd9b 100644 --- a/crates/ui-egui/tests/ui_roto_span_confidence.rs +++ b/crates/ui-egui/tests/ui_roto_span_confidence.rs @@ -82,6 +82,15 @@ fn the_span_bar_shows_the_trackers_confidence() { h.run_steps(4); assert_eq!(h.state().ui.focused, PanelKind::Layer); + // the object picker is in the bar; "+" starts a new object, which the next stroke would belong to + assert!(h.state().auto.find("layerPanel.roto.object").is_some(), "the object picker is on screen"); + assert!(h.state().auto.find("layerPanel.roto.newObject").is_some()); + assert_eq!(h.state().session.state.roto.object, 0); + h.state_mut().session.execute("roto.options", json!({"object": 1})).unwrap(); + h.run_steps(2); + assert_eq!(h.state().session.state.roto.object, 1); + h.state_mut().session.execute("roto.options", json!({"object": 0})).unwrap(); + h.run_steps(2); let track = h.state().auto.find("layerPanel.rotoSpan").expect("the span bar is on screen").rect; let (comp, layer) = { let comp = h.state().session.active_comp().unwrap().clone(); diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index 8b1931494..656574df6 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -29,6 +29,23 @@ strokes and span they came from are unchanged. Changing a stroke or the span mak Choosing another model, or the classic engine, leaves them in place and `roto.status` says so; a classic Propagate replaces them with classic mattes. Object Matte saves nothing. +## Several objects in one instance +A Roto Brush instance can hold several objects. The **Object N** button in the Layer panel's bar picks the object new +strokes belong to (it cycles through the objects that have strokes), and **+** starts a new one. Agents: `roto.options +{object}` sets the current object, `roto.stroke` takes `object` (0–31; default the current one), `roto.clearStrokes` +takes `object` to remove one object's strokes, and `roto.status` lists `objects` (`{id, strokes}`) and `currentObject`. + +With SAM 2, Propagate follows every object in **one pass** (the frame is encoded once, whatever the number of objects), each +from its own foreground and background strokes, and the mattes are united into the instance's one matte. Each object keeps +its own occlusion: the frame counts as absent only when every object is, and otherwise the least sure object sets the +frame's confidence on the span bar. A background stroke belongs to its object, so it separates that object from what +lies under it without cutting a hole in the others. Refine Edge strokes belong to the whole matte. The classic engine has +no objects: all foreground strokes make one object, and background strokes apply to it as before. Saved projects are +unchanged (strokes without an object are object 0, and nothing is written for object 0). + +This is not several Roto Brush instances sharing a pass: each instance reads the layer after the effects above it, so +instances can see different pixels. Objects inside one instance always see the same ones. + ## Sizes Four checkpoints are in the registry (Settings ▸ Roto Brush); the size is detected from the file, and each is pinned to a Hugging Face commit with its SHA-256. All run the same tracker; a bigger backbone costs time and memory. From f67a17e76555055795143c10e8e4a408969a6015 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Sat, 17 Oct 2026 16:15:21 +0200 Subject: [PATCH 48/77] S2.70: each object keeps its own matte; Show Object outputs one of them Co-Authored-By: Claude Sonnet 5.5 --- crates/effects/src/object_matte.rs | 2 +- crates/effects/src/roto.rs | 156 +++++++++++++++++++++++++++-- docs/effects.md | 2 +- 3 files changed, 152 insertions(+), 8 deletions(-) diff --git a/crates/effects/src/object_matte.rs b/crates/effects/src/object_matte.rs index 3a75bec1b..8a9750282 100644 --- a/crates/effects/src/object_matte.rs +++ b/crates/effects/src/object_matte.rs @@ -227,7 +227,7 @@ pub fn job_frames(params: &Params, limits: [i64; 2]) -> Vec { pub fn chain(params: &Params, layer_size: [f64; 2], fps: f64, f: i64, scale: f64, model: Option>) -> Option { let data = roto_data(params, layer_size, f, model.is_some())?; let keys = data.chain_keys(seed(params, layer_size, fps, model.as_deref())); - Some(Chain { data: Arc::new(data), keys, opts: seg_opts(scale), refine_radius: 0.0, model, saved: None }) + Some(Chain { data: Arc::new(data), keys, opts: seg_opts(scale), refine_radius: 0.0, model, saved: None, show_object: 0 }) } /// Frame `f`'s segmentation at full resolution, computing what is missing (the background diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index ffa41f01e..e3f6da171 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -80,6 +80,7 @@ pub fn specs() -> Vec { p("rotoBrushMatte/contrast", "Contrast", num(0.0), slider(0.0, 100.0, 0.0, 100.0, 0)), p("rotoBrushMatte/shiftEdge", "Shift Edge", num(0.0), slider(-100.0, 100.0, -100.0, 100.0, 0)), p("rotoBrushMatte/reduceChatter", "Reduce Chatter", num(0.0), slider(0.0, 100.0, 0.0, 100.0, 0)), + p("rotoBrushMatte/object", "Show Object", num(0.0), slider(0.0, 32.0, 0.0, 32.0, 0)), p("refineEdgeMatte/smooth", "Smooth", num(1.0), slider(0.0, 100.0, 0.0, 10.0, 1)), p("refineEdgeMatte/feather", "Feather", num(0.0), slider(0.0, 100.0, 0.0, 20.0, 1)), p("refineEdgeMatte/contrast", "Contrast", num(0.0), slider(0.0, 100.0, 0.0, 100.0, 0)), @@ -320,6 +321,11 @@ fn tracked_key(key: u64) -> u64 { key.rotate_left(17) ^ TRACK_SALT } +/// Where one object's matte of a tracked frame lives (kept only when the instance has several objects). +fn object_key(key: u64, object: u32) -> u64 { + tracked_key(key).wrapping_add(0x9e37_79b9_7f4a_7c15u64.wrapping_mul(u64::from(object) + 1)) +} + /// How far to trust a frame the tracker followed. #[derive(Clone, Copy, Debug, PartialEq)] pub struct FrameTrust { @@ -398,6 +404,8 @@ pub struct Chain { pub model: Option>, /// The saved tracking, while it applies (see [`TrackCache`]). pub(crate) saved: Option, + /// Show Object: the matte of this object only (its id plus one; 0 shows all of them together). + pub show_object: u32, } impl Chain { @@ -405,7 +413,8 @@ impl Chain { let data = data(params); let keys = data.chain_keys(seed(params, layer_size, fps)); let saved = saved_for(params, &data, layer_size, fps); - Chain { refine_radius: data.refine_radius(), data, keys, opts: seg_opts(params, scale), model: model_for(params), saved } + let show_object = f(params, "rotoBrushMatte/object").round().clamp(0.0, 32.0) as u32; + Chain { refine_radius: data.refine_radius(), data, keys, opts: seg_opts(params, scale), model: model_for(params), saved, show_object } } /// Frame `f`'s segmentation at `scale` from the saved tracking or the caches; computes nothing. @@ -413,6 +422,49 @@ impl Chain { self.saved_seg(f, scale).or_else(|| self.keys.get(&f).and_then(|k| cached(*k, scale))) } + /// Object `o`'s matte of frame `f` at this chain's scale (saved tracking, else this session's), when the instance + /// has several objects and the tracker followed them there; computes nothing. + pub fn object_seg(&self, f: i64, o: u32) -> Option> { + let sc = self.opts.scale; + let k = *self.keys.get(&f)?; + let ok = object_key(k, o); + if let Some(s) = cached(ok, sc) { + return Some(s); + } + let full = match cached(ok, 1.0) { + Some(s) => s, + None => { + let t = self.saved.as_ref()?; + let s = Arc::new(decode_object(&t.cache, f, o)?); + store(ok, 1.0, s.clone()); + s + } + }; + if sc == 1.0 { + return Some(full); + } + let s = Arc::new(downscale(&full, sc)); + store(ok, sc, s.clone()); + Some(s) + } + + /// What Show Object asks for at frame `f`: `None` shows the whole matte (all objects, a single object, or a frame + /// the tracker did not follow); an object the instance does not have shows nothing. + fn shown_object(&self, f: i64) -> Option> { + let o = self.show_object.checked_sub(1)?; + let objects = self.data.objects(); + if objects.len() < 2 { + return None; + } + if !objects.contains(&o) { + let sc = self.opts.scale; + let [w, h] = + self.saved.as_ref().map(|t| t.cache.size).or_else(|| self.keys.get(&f).and_then(|k| cached(tracked_key(*k), sc)).map(|s| [s.w, s.h]))?; + return Some(Arc::new(FrameSeg::empty(w, h))); + } + self.object_seg(f, o) + } + /// The tracker's trust in frame `f`: from the saved tracking, else from this session's tracking. pub fn trust_at(&self, f: i64) -> Option { if let Some(fr) = self.saved.as_ref().and_then(|t| t.cache.frames.get(&f)) { @@ -455,6 +507,10 @@ impl Chain { pub fn ensure(&self, f: i64, frame: &mut dyn FnMut(i64) -> Option>, cancel: &dyn Fn() -> bool) -> Option> { let sc = self.opts.scale; let lookup = |g: i64| -> Option> { + // Show Object: one object's matte, where the tracker followed the objects. + if let Some(s) = self.shown_object(g) { + return Some(s); + } // A frame the tracker followed wins over the classic engine's: the saved tracking, else // this session's. if let Some(s) = self.saved_seg(g, sc) { @@ -604,13 +660,13 @@ impl Chain { let mut pending: BTreeMap>> = BTreeMap::new(); let finish = |g: i64, parts: Vec>, summary: &mut TrackSummary, progress: &mut dyn FnMut(usize)| { let (Some(key), Some(img)) = (self.keys.get(&g).copied(), (frame.borrow_mut())(g)) else { return }; - let mut segs = vec![]; + let mut segs: Vec<(u32, effectcraft_track::roto::FrameSeg)> = vec![]; let (mut confidence, mut present) = (f32::INFINITY, false); let mut worst = f32::INFINITY; for (obj, t) in objects.iter().zip(parts) { let Some(t) = t else { continue }; let strokes = self.data.strokes_of(g, *obj); - segs.push(rb::segment_tracked(&img, &strokes, t.prob, t.occluded, &self.opts, self.refine_radius)); + segs.push((*obj, rb::segment_tracked(&img, &strokes, t.prob, t.occluded, &self.opts, self.refine_radius))); worst = worst.min(t.confidence); if !t.occluded { present = true; @@ -622,7 +678,16 @@ impl Chain { } // a frame is occluded when every object is absent; otherwise the doubtful object sets the confidence let (occluded, confidence) = (!present, if present { confidence } else { worst }); - store(tracked_key(key), self.opts.scale, Arc::new(effectcraft_track::roto::FrameSeg::union(segs))); + // with several objects each keeps its own matte too (Show Object, the per-object overlay) + let several = objects.len() > 1; + let mut parts = Vec::with_capacity(segs.len()); + for (obj, seg) in segs { + if several { + store(object_key(key, obj), self.opts.scale, Arc::new(seg.clone())); + } + parts.push(seg); + } + store(tracked_key(key), self.opts.scale, Arc::new(effectcraft_track::roto::FrameSeg::union(parts))); trust_cache().lock().unwrap_or_else(|e| e.into_inner()).insert(tracked_key(key), FrameTrust { confidence, occluded }); summary.frames += 1; summary.occluded += usize::from(occluded); @@ -656,13 +721,26 @@ impl Chain { } } +/// Each object's matte of layer frame `f` where the tracker followed several objects (full resolution; computes +/// nothing). Show Object limits it to the one chosen. Empty with fewer than two objects. +pub fn object_mattes(params: &Params, layer_size: [f64; 2], fps: f64, f: i64) -> Vec<(u32, Arc)> { + let mut chain = Chain::new(params, layer_size, fps, 1.0); + let objects = chain.data.objects(); + if objects.len() < 2 || !chain.data.in_span(f) { + return vec![]; + } + let shown = chain.show_object.checked_sub(1); + chain.show_object = 0; + objects.into_iter().filter(|o| shown.is_none_or(|s| s == *o)).filter_map(|o| chain.object_seg(f, o).map(|s| (o, s))).collect() +} + /// Nearest-neighbour resampling of a segmentation to another scale. pub fn downscale(s: &FrameSeg, scale: f64) -> FrameSeg { let (w, h) = (((s.w as f64) * scale).round().max(1.0) as usize, ((s.h as f64) * scale).round().max(1.0) as usize); resample(s, w, h) } -pub(crate) fn resample(s: &FrameSeg, w: usize, h: usize) -> FrameSeg { +pub fn resample(s: &FrameSeg, w: usize, h: usize) -> FrameSeg { if s.w == w && s.h == h { return s.clone(); } @@ -741,6 +819,9 @@ pub fn frozen_key(params: &Params, layer_size: [f64; 2], fps: f64) -> String { fnv(&mut h, &k.to_le_bytes()); } fnv(&mut h, format!("{:?}", matte_params(params)).as_bytes()); + if ch.show_object != 0 { + fnv(&mut h, &ch.show_object.to_le_bytes()); + } format!("{h:016x}") } @@ -788,6 +869,17 @@ pub struct TrackCache { /// The mattes' width and height (the layer's size). pub size: [usize; 2], pub frames: BTreeMap, + /// With several objects: each one's matte of every tracked frame (object id, then frame). + #[serde(skip_serializing_if = "BTreeMap::is_empty")] + pub objects: BTreeMap>, +} + +/// One object's matte and Refine Edge band of a tracked frame (RLE, base64). +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(default)] +pub struct ObjectFrame { + pub m: String, + pub r: String, } /// One tracked frame: its matte and Refine Edge band (RLE, base64), its colour models and the @@ -876,6 +968,16 @@ fn decode_tracked(cache: &TrackCache, f: i64) -> Option { (matte.len() == n && refine.len() == n).then(|| FrameSeg { w, h, matte, refine, fg: fr.fg.clone(), bg: fr.bg.clone() }) } +/// Object `o`'s matte of a tracked frame of saved tracking (full resolution). +fn decode_object(cache: &TrackCache, f: i64, o: u32) -> Option { + let fr = cache.objects.get(&o)?.get(&f)?; + let [w, h] = cache.size; + let n = w.checked_mul(h)?; + let matte = rle::decode(&fr.m)?; + let refine = rle::decode(&fr.r)?; + (matte.len() == n && refine.len() == n).then(|| FrameSeg { w, h, matte, refine, ..Default::default() }) +} + /// The tracking a chain has just made, to save with its instance: every frame of its span that the /// tracker followed, at full resolution, with its trust. `None` when it followed none. pub fn tracks_of(ch: &Chain) -> Option { @@ -890,7 +992,21 @@ pub fn tracks_of(ch: &Chain) -> Option { TrackFrame { m: rle::encode(&seg.matte), r: rle::encode(&seg.refine), fg: seg.fg.clone(), bg: seg.bg.clone(), c: t.confidence, o: t.occluded }, ); } - let cache = TrackCache { key: tracks_key(&ch.keys), model: model.to_string(), size: size?, frames }; + let mut objects: BTreeMap> = BTreeMap::new(); + let ids = ch.data.objects(); + if ids.len() > 1 { + for (f, k) in &ch.keys { + if !frames.contains_key(f) { + continue; + } + for o in &ids { + if let Some(seg) = cached(object_key(*k, *o), 1.0) { + objects.entry(*o).or_default().insert(*f, ObjectFrame { m: rle::encode(&seg.matte), r: rle::encode(&seg.refine) }); + } + } + } + } + let cache = TrackCache { key: tracks_key(&ch.keys), model: model.to_string(), size: size?, frames, objects }; (!cache.frames.is_empty()).then_some(cache) } @@ -1209,6 +1325,34 @@ mod tests { assert!(at(46, 40) != 0, "the moving disk (frame 3 centre x = 46)"); assert!(at(15, 15) != 0, "the fixed square"); assert_eq!(at(60, 70), 0, "nothing elsewhere"); + + // each object keeps its own matte + let (disk, square) = (chain.object_seg(3, 0).expect("object 0"), chain.object_seg(3, 1).expect("object 1")); + assert!(disk.matte[40 * 100 + 46] != 0 && disk.matte[15 * 100 + 15] == 0); + assert!(square.matte[15 * 100 + 15] != 0 && square.matte[40 * 100 + 46] == 0); + // Show Object picks one (its id plus one); an object the instance does not have shows nothing; 0 shows all + let mut shown = |c: &Chain| c.ensure(3, &mut get, &|| false).expect("frame 3"); + chain.show_object = 2; + let one = shown(&chain); + assert!(one.matte[15 * 100 + 15] != 0 && one.matte[40 * 100 + 46] == 0, "only the square"); + chain.show_object = 1; + let zero = shown(&chain); + assert!(zero.matte[40 * 100 + 46] != 0 && zero.matte[15 * 100 + 15] == 0, "only the disk"); + chain.show_object = 7; + assert_eq!(shown(&chain).area(), 0, "no such object"); + chain.show_object = 0; + assert!(shown(&chain).matte[15 * 100 + 15] != 0, "all together"); + + // the objects are saved with the tracking and come back from the project file + let saved = tracks_of(&chain).expect("tracking to save"); + assert_eq!(saved.objects.keys().copied().collect::>(), vec![0, 1]); + assert_eq!(saved.objects.get(&1).map(|f| f.len()), Some(9)); + assert_eq!(TrackCache::from_json(&saved.to_json()), Some(saved.clone())); + // a tracking saved before objects existed has none + let old = TrackCache { objects: Default::default(), ..saved.clone() }; + assert!(!old.to_json().contains("objects")); + let rt = decode_object(&saved, 3, 1).expect("decoded"); + assert!(rt.matte[15 * 100 + 15] != 0 && rt.matte[40 * 100 + 46] == 0); } /// Instances saved before the propagation controls got their own twirl-down (and motion diff --git a/docs/effects.md b/docs/effects.md index b3adb46ec..e233dd47f 100644 --- a/docs/effects.md +++ b/docs/effects.md @@ -259,7 +259,7 @@ Effects are addressed by id (`ec..`) or display name in commands | Object Matte | `ec.matte.objectmatte` | 13 | | 32 | Implemented | | Refine Hard Matte | `ec.matte.refinehard` | 16 | GPU | 32 | Implemented | | Refine Soft Matte | `ec.matte.refinesoft` | 18 | GPU | 32 | Implemented | -| Roto Brush & Refine Edge | `ec.matte.rotobrush` | 24 | | 32 | Implemented | +| Roto Brush & Refine Edge | `ec.matte.rotobrush` | 25 | | 32 | Implemented | | Simple Choker | `ec.matte.simplechoker` | 2 | GPU | 32 | Implemented | ## Noise & Grain From 4e2b2b7e5db61825cec5336d1b615794ae8da0e8 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Sat, 17 Oct 2026 16:15:22 +0200 Subject: [PATCH 49/77] S2.71: the boundary view outlines each object in its own colour Co-Authored-By: Claude Sonnet 5.5 --- crates/engine/src/analysis/roto.rs | 2 +- crates/ui-egui/src/panels/roto_tool.rs | 43 ++++++++++- crates/ui-egui/tests/ui_roto_objects.rs | 97 +++++++++++++++++++++++++ docs/roto-brush-sam2.md | 8 ++ 4 files changed, 148 insertions(+), 2 deletions(-) create mode 100644 crates/ui-egui/tests/ui_roto_objects.rs diff --git a/crates/engine/src/analysis/roto.rs b/crates/engine/src/analysis/roto.rs index 626de3030..399da1869 100644 --- a/crates/engine/src/analysis/roto.rs +++ b/crates/engine/src/analysis/roto.rs @@ -293,7 +293,7 @@ pub(crate) fn params_at(p: &Project, cid: ItemId, layer: &Layer, g: &PropGroup, } /// The layer's source size (layer pixels). -pub(crate) fn layer_size(p: &Project, comp: &effectcraft_project::Comp, layer: &Layer) -> [f64; 2] { +pub fn layer_size(p: &Project, comp: &effectcraft_project::Comp, layer: &Layer) -> [f64; 2] { effectcraft_render::effect_bounds(p, comp, layer).0 } diff --git a/crates/ui-egui/src/panels/roto_tool.rs b/crates/ui-egui/src/panels/roto_tool.rs index bf47b558d..56fa39b0a 100644 --- a/crates/ui-egui/src/panels/roto_tool.rs +++ b/crates/ui-egui/src/panels/roto_tool.rs @@ -21,6 +21,17 @@ fn roto_group(layer: &Layer) -> Option<&effectcraft_engine::project::PropGroup> layer.effects()?.groups().nth(roto_index(layer)?) } +/// The colour of object `o`'s outline: object 0 keeps the chosen boundary colour, the others cycle through a palette. +pub fn object_color(o: u32, base: [f64; 3]) -> Color32 { + const PALETTE: [[u8; 3]; 8] = + [[60, 200, 255], [255, 170, 40], [150, 230, 60], [180, 110, 255], [255, 225, 60], [40, 210, 170], [255, 90, 80], [90, 130, 255]]; + if o == 0 { + return c32(base, 255); + } + let [r, g, b] = PALETTE[(o as usize - 1) % PALETTE.len()]; + Color32::from_rgb(r, g, b) +} + /// A stroke being drawn. #[derive(Clone, Debug)] struct Drag { @@ -94,6 +105,22 @@ pub fn draw_overlay( let alpha: Vec = (0..w * h).map(|i| a_at((i % w) as i64, (i / w) as i64)).collect(); let inside = |x: i64, y: i64| x >= 0 && y >= 0 && x < w as i64 && y < h as i64 && alpha[y as usize * w + x as usize] >= 0.5; let mut px = vec![Color32::TRANSPARENT; w * h]; + // With several objects the boundary view outlines each in its own colour. + let per_object: Vec<(Color32, Vec)> = match (o.view.as_str(), roto_group(layer)) { + ("alphaBoundary", Some(g)) => { + let params = effectcraft_engine::effects::flatten_params(g, &mut |p| p.value.clone()); + let size = effectcraft_engine::roto::layer_size(&app.session.project, comp, layer); + 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()) + }) + .collect() + } + _ => vec![], + }; let bc = c32(o.boundary_color, 255); let oc = c32(o.overlay_color, (o.overlay_opacity / 100.0 * 255.0).clamp(0.0, 255.0) as u8); // View Search Region: the band around the boundary (yellow). @@ -117,7 +144,20 @@ pub fn draw_overlay( _ => { let me = a >= 0.5; let edge = me && [(-1, 0), (1, 0), (0, -1), (0, 1)].iter().any(|(dx, dy)| !inside(x + dx, y + dy)); - if edge { bc } else { Color32::TRANSPARENT } + if per_object.is_empty() { + 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]; + per_object + .iter() + .rev() + .find(|(_, m)| { + inside_of(m, x, y) && [(-1, 0), (1, 0), (0, -1), (0, 1)].iter().any(|(dx, dy)| !inside_of(m, x + dx, y + dy)) + }) + .map_or(Color32::TRANSPARENT, |(c, _)| *c) + } } }; if let Some((sd, r)) = &search { @@ -389,6 +429,7 @@ pub fn bar_buttons(app: &mut EffectcraftApp, ui: &mut egui::Ui, p: &egui::Painte let next = all.iter().copied().find(|o| *o > current).or_else(|| all.first().copied()).unwrap_or(0); let _ = app.session.execute("roto.options", json!({"object": next})); } + p.circle_filled(pos2(r.min.x + 9.0, r.center().y), 4.0, object_color(current, app.session.state.roto.boundary_color)); app.auto.add("layerPanel.roto.object", r, &label); x += 90.0; let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(26.0, 20.0)); diff --git a/crates/ui-egui/tests/ui_roto_objects.rs b/crates/ui-egui/tests/ui_roto_objects.rs new file mode 100644 index 000000000..ee5efc020 --- /dev/null +++ b/crates/ui-egui/tests/ui_roto_objects.rs @@ -0,0 +1,97 @@ +//! A Roto Brush instance with two objects: the Layer panel's boundary view outlines each in its own colour, and +//! Show Object limits the matte to one (egui_kittest, a real render). A stand-in tracker follows two regions, so no +//! weights are needed. `ROTO_OBJECTS_SNAPSHOT=/abs/out.png` also saves the rendered panel. + +use std::sync::Arc; + +use effectcraft_engine::Session; +use effectcraft_segment::{ClipFrame, MaskModel, ModelInfo, ObjectTracker, Prompt, TrackRequest, TrackedFrame}; +use effectcraft_ui_egui::EffectcraftApp; +use effectcraft_ui_egui::dock::PanelKind; +use egui_kittest::Harness; +use serde_json::json; + +/// Follows two fixed regions of the layer: object 0 on the left, object 1 on the right. +struct TwoRegions; + +impl MaskModel for TwoRegions { + fn info(&self) -> &'static ModelInfo { + &effectcraft_segment::SAM2_TINY + } + fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &Prompt) -> effectcraft_segment::Result> { + Err("no weights".into()) + } + fn tracker(&self) -> Option<&dyn ObjectTracker> { + Some(self) + } +} + +impl ObjectTracker for TwoRegions { + fn track( + &self, + req: &TrackRequest, + frame: &mut dyn FnMut(i64) -> Option>, + _: &dyn Fn() -> bool, + out: &mut dyn FnMut(usize, i64, TrackedFrame), + ) -> effectcraft_segment::Result<()> { + for g in req.first..req.first + req.count as i64 { + let img = frame(g).ok_or("no frame")?; + for o in 0..req.objects.len() { + let (x0, x1) = if o == 0 { (img.w / 8, img.w * 3 / 8) } else { (img.w * 5 / 8, img.w * 7 / 8) }; + let prob = (0..img.w * img.h) + .map(|i| { + let (x, y) = (i % img.w, i / img.w); + f32::from(x >= x0 && x < x1 && y > img.h / 4 && y < img.h * 3 / 4) + }) + .collect(); + out(o, g, TrackedFrame { prob, w: img.w, h: img.h, confidence: 0.99, occluded: false }); + } + } + Ok(()) + } +} + +/// Pixels of the viewer (above the panel's bar, whose object swatch has the same colour) within a few levels of `want`. +fn count(image: &image::RgbaImage, want: [u8; 3]) -> usize { + let limit = image.height() * 54 / 100; + image.enumerate_pixels().filter(|(_, y, p)| *y < limit && p.0[..3].iter().zip(want).all(|(a, b)| a.abs_diff(b) <= 12)).count() +} + +#[test] +fn the_boundary_view_outlines_each_object_and_show_object_picks_one() { + let mut s = Session::default(); + s.execute("comp.new", json!({"name": "View", "width": 640, "height": 360, "duration": 4})).unwrap(); + s.execute("layer.newSolid", json!({"name": "Box", "color": "#e04020", "width": 80, "height": 80})).unwrap(); + let id = s.active_comp().unwrap().layers.iter().find(|l| l.name == "Box").unwrap().id.0; + for (object, x) in [(0, 22.0), (1, 58.0)] { + s.execute("roto.stroke", json!({"layer": id, "object": object, "points": [[x, 38.0], [x, 42.0]], "radius": 4})).unwrap(); + } + effectcraft_engine::effects::roto::set_model(Some(Arc::new(TwoRegions))); + s.execute("roto.propagate", json!({"layer": id, "direction": "both", "wait": true})).unwrap(); + let st = s.execute("roto.status", json!({"layer": id})).unwrap(); + assert_eq!(st["objects"], json!([{"id": 0, "strokes": 1}, {"id": 1, "strokes": 1}]), "{st}"); + + // object 0 keeps the boundary colour (pink by default), object 1 is the first palette colour (blue-green) + let (pink, cyan) = ([255, 64, 191], [60, 200, 255]); + let mut h = Harness::builder().with_size(egui::vec2(1600.0, 1000.0)).build_eframe(move |_| EffectcraftApp::new(s)); + h.run_steps(3); + h.state_mut().session.execute("layer.select", json!({"layers": [id]})).unwrap(); + let ctx = h.ctx.clone(); + effectcraft_ui_egui::menus::invoke(h.state_mut(), &ctx, "layer.openLayer", json!({"layer": id})).unwrap(); + h.run_steps(4); + assert_eq!(h.state().ui.focused, PanelKind::Layer); + let both = h.render().expect("the panel renders"); + if let Some(path) = std::env::var_os("ROTO_OBJECTS_SNAPSHOT") { + both.save(path).unwrap(); + } + let (a, b) = (count(&both, pink), count(&both, cyan)); + assert!(a > 50 && b > 50, "both objects are outlined: {a} pink, {b} cyan pixels"); + + // Show Object 2 (object id 1): only its outline remains + h.state_mut().session.execute("prop.set", json!({"layer": id, "path": "effects/#1/rotoBrushMatte/object", "value": 2})).unwrap(); + h.run_steps(4); + let one = h.render().expect("the panel renders"); + let (a, b) = (count(&one, pink), count(&one, cyan)); + assert!(a < 10 && b > 50, "only object 1 is shown: {a} pink, {b} cyan pixels"); + effectcraft_engine::effects::roto::set_model(None); +} diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index 656574df6..4e1792404 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -43,6 +43,14 @@ lies under it without cutting a hole in the others. Refine Edge strokes belong t no objects: all foreground strokes make one object, and background strokes apply to it as before. Saved projects are unchanged (strokes without an object are object 0, and nothing is written for object 0). +**Seeing and choosing objects.** The Layer panel's boundary view outlines each object in its own colour (object 1 in the +chosen boundary colour, the others from a palette), and the picker button carries the current object's colour. **Show +Object** (Roto Brush Matte, a slider: 0 shows all objects together, N shows object N alone) limits the instance's matte to +one object, so an effect after it can treat the pieces separately; an object the instance does not have shows nothing, and +with one object, or on frames the classic engine made, the setting changes nothing. Each object's matte is kept next to +the united one, in the session and in the saved tracking (`tracked` stores them only when there are several objects), so +Show Object renders the same after the project is reopened. Freeze records the Show Object that was set. + This is not several Roto Brush instances sharing a pass: each instance reads the layer after the effects above it, so instances can see different pixels. Objects inside one instance always see the same ones. From ba6fee986306960a02ed26240bb867d39a9b987b Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Sat, 17 Oct 2026 16:26:27 +0200 Subject: [PATCH 50/77] S2.72: the install test covers all four SAM 2 sizes from their files Co-Authored-By: Claude Sonnet 5.5 --- crates/engine/src/tests/models.rs | 30 ++++++++++++++++++++++-------- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index 3dee87e74..1dbc507aa 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -75,22 +75,36 @@ 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. + // The SAM 2.1 sizes, each from its own file (EFFECTCRAFT_SAM2, _SAM2_SMALL, _SAM2_BASE_PLUS, _SAM2_LARGE = + // path/to/model.safetensors of facebook/sam2.1-hiera-*): installed with its notice, loaded (verified, native + // only) and used by Roto Brush 3.0, then dropped and removed. #[cfg(not(target_arch = "wasm32"))] - if let Ok(path) = std::env::var("EFFECTCRAFT_SAM2") { + for (var, id) in [ + ("EFFECTCRAFT_SAM2", "sam2.1-hiera-tiny"), + ("EFFECTCRAFT_SAM2_SMALL", "sam2.1-hiera-small"), + ("EFFECTCRAFT_SAM2_BASE_PLUS", "sam2.1-hiera-base-plus"), + ("EFFECTCRAFT_SAM2_LARGE", "sam2.1-hiera-large"), + ] { + let Ok(path) = std::env::var(var) else { continue }; let r = s.execute_checked("roto.model.install", json!({"path": path})).unwrap(); - assert_eq!(r["id"], "sam2.1-hiera-tiny"); + assert_eq!(r["id"], id, "{var}"); let dir = s.models_dir().unwrap(); - let notice = std::fs::read_to_string(dir.join("sam2.1-hiera-tiny.safetensors.NOTICE.txt")).unwrap(); + let notice = std::fs::read_to_string(dir.join(format!("{id}.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(); + s.execute_checked("roto.model.select", json!({"id": id})).unwrap(); wait_loaded(&mut s, Task::Mask); - assert_eq!(s.models.status(Task::Mask).active.as_deref(), Some("sam2.1-hiera-tiny")); + assert_eq!(s.models.status(Task::Mask).active.as_deref(), Some(id), "{var}"); 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")); + assert_eq!(crate::effects::roto::model_for(¶ms).map(|m| m.info().id), Some(id)); + // the list says it is installed, and exactly one size is recommended + let list = s.execute_checked("roto.models", json!({})).unwrap(); + let entry = list["models"].as_array().unwrap().iter().find(|m| m["id"] == id).unwrap(); + assert_eq!(entry["installed"], true); + assert_eq!(list["models"].as_array().unwrap().iter().filter(|m| m["recommended"] == true).count(), 1); s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); assert!(crate::effects::roto::model_for(¶ms).is_none()); + assert_eq!(s.execute_checked("roto.model.remove", json!({"id": id})).unwrap()["removed"], true); + assert!(!dir.join(format!("{id}.safetensors")).exists(), "{id} removed"); } // EFFECTCRAFT_FACE_LANDMARKER=path/to/face_landmarker.task: installed with its notice, // loaded for face tracking, dropped when the classic tracker is chosen, removed. From a0ec3e2da4d0e15f079d90fb26b98f5f4ef22ab8 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Sat, 17 Oct 2026 17:13:43 +0200 Subject: [PATCH 51/77] S2.73: Tiny is the recommended size (benchmarked per size); stroke count and remove button per object; download test bench --roto per size: no larger size tracked better on the synthetic clips and Large lost the object after it re-entered; the hints say what was measured. The bar shows the current object's stroke count and can remove its strokes. An env-gated test runs the in-app download of a size. Co-Authored-By: Claude Sonnet 5.5 --- crates/engine/src/tests/models.rs | 23 ++++++++++++++++++ crates/sam2/src/mask.rs | 32 ++++++++++++++----------- crates/ui-egui/src/panels/roto_tool.rs | 21 ++++++++++++---- crates/ui-egui/tests/ui_roto_objects.rs | 3 +++ docs/roto-brush-sam2.md | 24 ++++++++++++++++--- 5 files changed, 82 insertions(+), 21 deletions(-) diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index 1dbc507aa..ed0eecf9b 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -160,3 +160,26 @@ fn an_idle_roto_model_is_unloaded() { crate::effects::roto::set_model(None); let _ = std::fs::remove_dir_all(d); } + +/// With `EFFECTCRAFT_DOWNLOAD_SIZE=sam2.1-hiera-small` (or another registry id): the in-app download (Settings ▸ Roto +/// Brush ▸ Download) fetches the pinned file over the network, verifies it, installs it with its notice and loads it. +/// Skipped without the variable: it needs the network and hundreds of megabytes. +#[cfg(not(target_arch = "wasm32"))] +#[test] +fn downloading_a_size_installs_and_loads_it() { + let Ok(id) = std::env::var("EFFECTCRAFT_DOWNLOAD_SIZE") else { return }; + crate::tests_roto::hold_roto_cache_test_lock(); + let (mut s, d) = session("download"); + let r = s.execute_checked("roto.model.download", json!({"id": id, "wait": true})).unwrap(); + assert_eq!(r["id"], id.as_str(), "{r}"); + let dir = s.models_dir().unwrap(); + let info = effectcraft_segment::info(&id).unwrap(); + let file = dir.join(info.file_name); + assert_eq!(std::fs::metadata(&file).unwrap().len(), info.size, "the whole file arrived"); + assert!(dir.join(format!("{}.NOTICE.txt", info.file_name)).exists()); + assert!(!dir.join(format!("{}.part", info.file_name)).exists(), "no partial file is left"); + s.execute_checked("roto.model.select", json!({"id": id})).unwrap(); + wait_loaded(&mut s, Task::Mask); + assert_eq!(s.models.status(Task::Mask).active.as_deref(), Some(id.as_str())); + let _ = std::fs::remove_dir_all(d); +} diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index d840d4052..8470bd4f3 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -148,25 +148,29 @@ pub fn load(bytes: &[u8]) -> std::result::Result, String> { Ok(Arc::new(Sam2Mask::with_info(Arc::new(model), info))) } -/// The size to suggest on this computer: Small when a GPU is usable (no slower to run than Tiny there, and more accurate), -/// Tiny otherwise (the bigger backbones take several seconds a frame on a CPU). +/// The size to suggest: Tiny. On the benchmark clips (`bench --roto`, 24 frames, a GPU) no larger size tracked better, they +/// cost 1.0x (Small), 1.15x (Base+) and 1.7x (Large) the time, and Large lost the object after it re-entered the frame. pub fn recommended() -> &'static ModelInfo { - if crate::gpu::Gpu::shared().is_ok() { &effectcraft_segment::SAM2_SMALL } else { &SAM2_TINY } + &SAM2_TINY } -/// A one-line note for the Settings list: what suits this computer, and what the bigger sizes cost without a GPU. +/// A one-line note for the Settings list: what was measured for this size, and what the bigger ones cost without a GPU. pub fn hint(id: &str) -> Option { - if !SIZES.iter().any(|m| m.id == id) { - return None; - } let gpu = crate::gpu::Gpu::shared().is_ok(); - Some(match (id == recommended().id, gpu, id) { - (true, true, _) => "Recommended: a GPU was found, and this size is no slower than Tiny on it.".to_string(), - (true, false, _) => "Recommended: no GPU was found, and this is the quickest on a CPU.".to_string(), - (false, false, "sam2.1-hiera-base-plus" | "sam2.1-hiera-large") => "Slow without a GPU (several seconds a frame); none was found.".to_string(), - (false, true, "sam2.1-hiera-large") => "For difficult footage; about 0.7 s of GPU time to look at each frame.".to_string(), - _ => return None, - }) + Some( + match id { + "sam2.1-hiera-tiny" => "Recommended: the smallest download, and no larger size tracked better on the benchmark clips.", + "sam2.1-hiera-small" => "As fast as Tiny on a GPU and no better on the benchmark clips; worth a try on difficult footage.", + "sam2.1-hiera-base-plus" if !gpu => "Slow without a GPU (several seconds a frame); none was found.", + "sam2.1-hiera-base-plus" => "About 15% slower than Tiny on a GPU and no better on the benchmark clips.", + "sam2.1-hiera-large" if !gpu => "Slow without a GPU (many seconds a frame); none was found.", + "sam2.1-hiera-large" => { + "About 1.7 times Tiny's time on a GPU; it lost the object after it re-entered in one benchmark clip. Try it on difficult footage." + } + _ => return None, + } + .to_string(), + ) } /// The registry entries this crate loads. diff --git a/crates/ui-egui/src/panels/roto_tool.rs b/crates/ui-egui/src/panels/roto_tool.rs index 56fa39b0a..0b42befbe 100644 --- a/crates/ui-egui/src/panels/roto_tool.rs +++ b/crates/ui-egui/src/panels/roto_tool.rs @@ -417,8 +417,9 @@ pub fn bar_buttons(app: &mut EffectcraftApp, ui: &mut egui::Ui, p: &egui::Painte let objects = fx::data(¶ms).objects(); let current = app.session.state.roto.object; let known = objects.contains(¤t); - let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(86.0, 20.0)); - let label = tr_args("Object {}", &[&(current + 1)]); + let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(104.0, 20.0)); + let strokes = fx::data(¶ms).strokes.iter().filter(|st| st.object == current && st.kind.is_segmentation()).count(); + let label = if strokes == 0 { tr_args("Object {}", &[&(current + 1)]) } else { tr_args("Object {} ({})", &[&(current + 1), &strokes]) }; if widgets::text_button(ui, r, &label, false, &t, egui::Id::new(("roto-object", uid))).clicked() { // cycle through the objects that have strokes (and the empty one being drawn) let mut all = objects.clone(); @@ -431,14 +432,26 @@ pub fn bar_buttons(app: &mut EffectcraftApp, ui: &mut egui::Ui, p: &egui::Painte } p.circle_filled(pos2(r.min.x + 9.0, r.center().y), 4.0, object_color(current, app.session.state.roto.boundary_color)); app.auto.add("layerPanel.roto.object", r, &label); - x += 90.0; + x += 108.0; let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(26.0, 20.0)); if widgets::text_button(ui, r, "+", false, &t, egui::Id::new(("roto-object-new", uid))).clicked() { let next = objects.iter().copied().chain([current]).max().map_or(0, |m| m + 1); let _ = app.session.execute("roto.options", json!({"object": next})); } app.auto.add("layerPanel.roto.newObject", r, "New object"); - x += 34.0; + x += 30.0; + // remove the current object's strokes (the object stays selected; with none left it is simply empty) + if strokes > 0 { + let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(26.0, 20.0)); + let c = widgets::text_button(ui, r, "×", false, &t, egui::Id::new(("roto-object-clear", uid))).on_hover_text(tr("Remove this object's strokes")); + if c.clicked() + && let Err(e) = app.session.execute("roto.clearStrokes", json!({"layer": layer.id.0, "effect": uid, "object": current})) + { + app.ui.status = e.to_string(); + } + app.auto.add("layerPanel.roto.clearObject", r, "Remove object strokes"); + x += 30.0; + } if running { let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(60.0, 20.0)); if widgets::text_button(ui, r, tr("Stop"), false, &t, egui::Id::unique("roto-stop")).clicked() { diff --git a/crates/ui-egui/tests/ui_roto_objects.rs b/crates/ui-egui/tests/ui_roto_objects.rs index ee5efc020..5d16b0ada 100644 --- a/crates/ui-egui/tests/ui_roto_objects.rs +++ b/crates/ui-egui/tests/ui_roto_objects.rs @@ -84,6 +84,9 @@ fn the_boundary_view_outlines_each_object_and_show_object_picks_one() { if let Some(path) = std::env::var_os("ROTO_OBJECTS_SNAPSHOT") { both.save(path).unwrap(); } + // the bar offers to remove the current object's strokes, and says how many it has + assert!(h.state().auto.find("layerPanel.roto.clearObject").is_some(), "the remove button is on screen"); + assert!(h.state().auto.find("layerPanel.roto.object").is_some_and(|e| e.label.contains("(1)")), "the object's stroke count is on the button"); let (a, b) = (count(&both, pink), count(&both, cyan)); assert!(a > 50 && b > 50, "both objects are outlined: {a} pink, {b} cyan pixels"); diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index 4e1792404..72c48e4e1 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -65,9 +65,27 @@ pinned to a Hugging Face commit with its SHA-256. All run the same tracker; a bi | SAM 2.1 Base+ (`sam2.1-hiera-base-plus`) | 323.5 MB | 0.37 s | encoder 2.3e-5 | | SAM 2.1 Large (`sam2.1-hiera-large`) | 897.9 MB | 0.72 s | encoder 5.2e-4 | -Settings marks the size that suits the computer: **Small** when a usable GPU is found (it encodes no slower than -Tiny there and is more accurate), **Tiny** otherwise; Base+ and Large say so when they would be slow on a CPU. The list -carries `recommended` and `hint` per model (`roto.models`). +### How the sizes compare on the benchmark clips +`effectcraft-cli bench --roto --frames 24 --fixes --sam2 FILE`, release build, SAM 2 on the GPU (RX 7900 GRE), one run each. +Mean IoU (after the object returns, in brackets) and seconds per tracked frame: + +| Clip | Tiny | Small | Base+ | Large | +|---|---|---|---|---| +| move | 0.992 | 0.988 | 0.986 | 0.994 | +| light | 0.990 | 0.985 | 0.987 | 0.993 | +| occlusion | 0.989 (0.987) | 0.988 (0.984) | 0.987 (0.984) | 0.989 (0.987) | +| reenter | 0.982 (0.975) | 0.981 (0.970) | 0.978 (0.966) | **0.694 (0.110)** | +| s / tracked frame | 0.69 | 0.70 | 0.79 | 1.17 | + +The classic engine scores 0.98 / 0.98 / 0.49 / 0.62 on the same clips. Peak tracker state is 88 MB for every size (the +weights are not counted). The clips are synthetic and easy, so they cannot tell the sizes apart on accuracy: **no larger +size tracked better here, and Large lost the object after it left and came back** (IoU 0.11 after the return, where +Tiny, Small and Base+ found it again). That is the model's behaviour on this clip, not the port: +Large matches PyTorch stage by stage. Whether a larger size helps on real footage is not measured. + +Settings therefore marks **Tiny** as recommended (the smallest download, and nothing larger did better), and says what was +measured for the others; Base+ and Large also warn when no GPU was found (they take several seconds a frame on a CPU). +The list carries `recommended` and `hint` per model (`roto.models`). ### Half-precision weights (optional) `EFFECTCRAFT_SAM2_F16=1` keeps the weight matrices on the GPU as half precision, halving their memory (Large: about From 8a6d48cdc8910a7da1351757efd1cb9740f67040 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Sat, 17 Oct 2026 20:28:45 +0000 Subject: [PATCH 52/77] S2.74: bound pending multi-object tracking masks --- crates/effects/src/roto.rs | 113 ++++++++++++++++++++--- crates/effects/src/roto/pending.rs | 138 +++++++++++++++++++++++++++++ docs/roto-brush-sam2.md | 5 ++ 3 files changed, 244 insertions(+), 12 deletions(-) create mode 100644 crates/effects/src/roto/pending.rs diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index e3f6da171..e8cae42c3 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -40,6 +40,8 @@ use crate::{Buf, EffectCtx, EffectSpec, Params, num, p, popup, slider}; pub use effectcraft_track::roto::{FrameSeg, Stroke, StrokeKind}; +mod pending; + /// The effect's id. pub const ID: &str = "ec.matte.rotobrush"; /// Hidden parameters. @@ -607,6 +609,16 @@ impl Chain { frame: &mut dyn FnMut(i64) -> Option>, cancel: &dyn Fn() -> bool, progress: &mut dyn FnMut(usize), + ) -> Result, String> { + self.track_with_pending_limit(frame, cancel, progress, pending::MAX_PENDING_BYTES) + } + + fn track_with_pending_limit( + &self, + frame: &mut dyn FnMut(i64) -> Option>, + cancel: &dyn Fn() -> bool, + progress: &mut dyn FnMut(usize), + max_pending_bytes: usize, ) -> Result, String> { let Some(tracker) = self.model.as_ref().and_then(|m| m.tracker()) else { return Ok(None) }; let (Some(_), true) = (self.data.base, self.data.span[1] >= self.data.span[0]) else { return Err("no base frame".into()) }; @@ -657,7 +669,7 @@ impl Chain { let frame = std::cell::RefCell::new(frame); let mut summary = TrackSummary::default(); // A frame is finished when every object has reported it (the objects' mattes are united). - let mut pending: BTreeMap>> = BTreeMap::new(); + let mut pending = pending::PendingMasks::new(objects.len(), max_pending_bytes); let finish = |g: i64, parts: Vec>, summary: &mut TrackSummary, progress: &mut dyn FnMut(usize)| { let (Some(key), Some(img)) = (self.keys.get(&g).copied(), (frame.borrow_mut())(g)) else { return }; let mut segs: Vec<(u32, effectcraft_track::roto::FrameSeg)> = vec![]; @@ -694,27 +706,37 @@ impl Chain { summary.min_confidence = if summary.frames == 1 { confidence } else { summary.min_confidence.min(confidence) }; progress(summary.frames); }; - let n_objects = objects.len(); - tracker.track( + let pending_failed = std::cell::Cell::new(false); + let mut pending_error = None; + let result = tracker.track( &req, &mut |g| { let img = (frame.borrow_mut())(g)?; Some(Arc::new(ClipFrame { rgb: rb::rgb_of(&img), w: img.width as usize, h: img.height as usize })) }, - cancel, + &|| cancel() || pending_failed.get(), &mut |o, g, t| { - let parts = pending.entry(g).or_insert_with(|| vec![None; n_objects]); - if let Some(slot) = parts.get_mut(o) { - *slot = Some(t); + // A tracker that reports object-by-object can otherwise retain a full clip of + // probability maps. Stop buffering immediately, even if it ignores cancellation. + if pending_failed.get() { + return; } - if parts.iter().all(Option::is_some) { - let parts = pending.remove(&g).unwrap_or_default(); - finish(g, parts, &mut summary, progress); + match pending.insert(o, g, t) { + Ok(Some(parts)) => finish(g, parts, &mut summary, progress), + Ok(None) => {} + Err(e) => { + pending_error = Some(e); + pending_failed.set(true); + } } }, - )?; + ); + if let Some(e) = pending_error { + return Err(e); + } + result?; // frames some object never reported (it was not followed there): united from the rest - for (g, parts) in std::mem::take(&mut pending) { + for (g, parts) in pending.into_incomplete() { finish(g, parts, &mut summary, progress); } Ok(Some(summary)) @@ -1285,6 +1307,73 @@ mod tests { } } + /// The ordering SAM 2 uses when objects are first prompted on different frames: each object + /// reports its entire span before the next. No weights or large allocations are needed. + #[derive(Default)] + struct SeparatePasses { + emitted: std::sync::atomic::AtomicUsize, + cancelled: std::sync::atomic::AtomicBool, + } + + impl effectcraft_segment::MaskModel for SeparatePasses { + fn info(&self) -> &'static effectcraft_segment::ModelInfo { + &effectcraft_segment::SAM2_TINY + } + fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &effectcraft_segment::Prompt) -> effectcraft_segment::Result> { + Err("no weights".into()) + } + fn tracker(&self) -> Option<&dyn effectcraft_segment::ObjectTracker> { + Some(self) + } + } + + impl effectcraft_segment::ObjectTracker for SeparatePasses { + fn track( + &self, + req: &effectcraft_segment::TrackRequest, + frames: &mut dyn FnMut(i64) -> Option>, + cancel: &dyn Fn() -> bool, + out: &mut dyn FnMut(usize, i64, effectcraft_segment::TrackedFrame), + ) -> effectcraft_segment::Result<()> { + use std::sync::atomic::Ordering::Relaxed; + assert_ne!(req.objects[0][0].0, req.objects[1][0].0); + for o in 0..req.objects.len() { + for g in req.first..req.first + req.count as i64 { + if cancel() { + self.cancelled.store(true, Relaxed); + return Err("tracking cancelled".into()); + } + let img = frames(g).ok_or("no frame")?; + self.emitted.fetch_add(1, Relaxed); + out(o, g, effectcraft_segment::TrackedFrame { prob: vec![0.5; img.w * img.h], w: img.w, h: img.h, confidence: 0.9, occluded: false }); + } + } + Ok(()) + } + } + + #[test] + fn separate_object_passes_cancel_at_the_pending_mask_budget() { + use std::sync::atomic::Ordering::Relaxed; + let mut d = RotoData::default(); + for object in 0..2 { + d.add(Stroke { kind: StrokeKind::Fg, frame: i64::from(object), radius: 3.0, points: vec![[35.0, 40.0]], object }, 8, [0, 24]); + } + let mut params = params_with(&d); + params.values.insert(INPUT_KEY.into(), Value::Str("pending-budget-regression".into())); + let mut chain = Chain::new(¶ms, [100.0, 80.0], 25.0, 1.0); + let fake = Arc::new(SeparatePasses::default()); + chain.model = Some(fake.clone()); + let mut get = |f: i64| Some(Arc::new(frame(f as f64 / 25.0))); + let mut completed = 0; + // Each map takes 32,000 bytes; two incomplete frames exceed this small test budget. + let error = chain.track_with_pending_limit(&mut get, &|| false, &mut |n| completed = n, 64_000).unwrap_err(); + assert!(error.contains("pending object masks") && error.contains("memory budget"), "{error}"); + assert!(fake.cancelled.load(Relaxed), "the model must see cancellation before allocating the rest of the clip"); + assert_eq!(fake.emitted.load(Relaxed), 2); + assert_eq!(completed, 0, "partial objects are not complete tracked frames"); + } + #[test] fn objects_are_followed_in_one_pass_and_their_mattes_united() { let mut d = RotoData::default(); diff --git a/crates/effects/src/roto/pending.rs b/crates/effects/src/roto/pending.rs new file mode 100644 index 000000000..256789ee9 --- /dev/null +++ b/crates/effects/src/roto/pending.rs @@ -0,0 +1,138 @@ +//! Bounded assembly of the objects a tracker reports separately for each frame. + +use std::collections::BTreeMap; +use std::mem::size_of; + +use effectcraft_segment::TrackedFrame; + +/// Separate from the model's memory banks: probability maps are full-resolution and can otherwise +/// grow with the entire clip when objects start on different frames. +pub(super) const MAX_PENDING_BYTES: usize = 256 << 20; + +pub(super) struct PendingMasks { + frames: BTreeMap>>, + objects: usize, + bytes: usize, + limit: usize, +} + +impl PendingMasks { + pub(super) fn new(objects: usize, limit: usize) -> Self { + Self { frames: BTreeMap::new(), objects, bytes: 0, limit } + } + + /// Probability allocation plus a conservative allowance for the map node and object slots. + fn frame_bytes(&self) -> usize { + self.objects.saturating_mul(size_of::>()).saturating_add(256) + } + + fn mask_bytes(mask: &TrackedFrame) -> usize { + mask.prob.capacity().saturating_mul(size_of::()) + } + + /// Complete frames leave immediately; incomplete frames may never exceed the budget. A rejected + /// insertion leaves the buffer unchanged, including when it replaces an object's previous result. + pub(super) fn insert(&mut self, object: usize, frame: i64, mask: TrackedFrame) -> Result>>, String> { + if object >= self.objects { + return Err("tracker reported an unknown object".into()); + } + if self.objects == 1 { + return Ok(Some(vec![Some(mask)])); + } + let existing = self.frames.get(&frame); + let old_bytes = existing.and_then(|p| p.get(object)).and_then(Option::as_ref).map_or(0, Self::mask_bytes); + let added = Self::mask_bytes(&mask).saturating_add(if existing.is_none() { self.frame_bytes() } else { 0 }); + let bytes = self.bytes.saturating_sub(old_bytes).checked_add(added); + let Some(bytes) = bytes.filter(|n| *n <= self.limit) else { + return Err(format!("pending object masks exceed their {} MiB memory budget", self.limit >> 20)); + }; + let parts = self.frames.entry(frame).or_insert_with(|| vec![None; self.objects]); + if let Some(slot) = parts.get_mut(object) { + *slot = Some(mask); + } + self.bytes = bytes; + if parts.iter().all(Option::is_some) { + let parts = self.frames.remove(&frame).unwrap_or_default(); + let used = parts.iter().flatten().fold(self.frame_bytes(), |n, m| n.saturating_add(Self::mask_bytes(m))); + self.bytes = self.bytes.saturating_sub(used); + return Ok(Some(parts)); + } + Ok(None) + } + + pub(super) fn into_incomplete(self) -> BTreeMap>> { + self.frames + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn mask(value: f32) -> TrackedFrame { + TrackedFrame { prob: vec![value; 16], w: 4, h: 4, confidence: value, occluded: false } + } + + #[test] + fn different_start_passes_stop_before_buffering_the_whole_clip() { + let mut pending = PendingMasks::new(2, 2048); + let mut accepted = 0; + // SAM 2 can report all of object 0 before starting object 1. A long clip must stop + // at the memory limit even though no frame has received all its objects yet. + for frame in 0..10_000 { + match pending.insert(0, frame, mask(0.9)) { + Ok(None) => accepted += 1, + Ok(Some(_)) => panic!("no frame is complete yet"), + Err(e) => { + assert!(e.contains("memory budget"), "{e}"); + break; + } + } + } + assert!(accepted > 0 && accepted < 10_000); + assert_eq!(pending.frames.len(), accepted); + assert!(pending.bytes <= pending.limit); + let before = pending.bytes; + assert!(pending.insert(0, 10_001, mask(0.8)).is_err()); + assert_eq!(pending.bytes, before, "rejection must not retain the new map"); + } + + #[test] + fn interleaved_objects_release_the_budget_and_keep_their_masks() { + let mut pending = PendingMasks::new(2, 2048); + for frame in 0..10_000 { + let first = (frame % 2) as usize; + assert!(pending.insert(first, frame, mask(first as f32)).unwrap().is_none()); + let parts = pending.insert(1 - first, frame, mask((1 - first) as f32)).unwrap().unwrap(); + assert_eq!(parts[0].as_ref().unwrap().prob, vec![0.0; 16]); + assert_eq!(parts[1].as_ref().unwrap().prob, vec![1.0; 16]); + assert!(pending.frames.is_empty()); + assert_eq!(pending.bytes, 0); + } + } + + #[test] + fn replacements_charge_capacity_and_partial_frames_are_preserved() { + let mut pending = PendingMasks::new(2, 2048); + pending.insert(0, 7, mask(0.1)).unwrap(); + let before = pending.bytes; + pending.insert(0, 7, mask(0.2)).unwrap(); + assert_eq!(pending.bytes, before, "a duplicate replaces rather than accumulates"); + let mut oversized = mask(0.3); + oversized.prob.reserve(2048); + assert!(pending.insert(0, 7, oversized).is_err(), "capacity, not length, owns the memory"); + assert_eq!(pending.bytes, before); + let frames = pending.into_incomplete(); + assert_eq!(frames[&7][0].as_ref().unwrap().prob, vec![0.2; 16]); + assert!(frames[&7][1].is_none()); + } + + #[test] + fn single_objects_need_no_pending_budget_and_unknown_objects_are_rejected() { + let mut pending = PendingMasks::new(1, 0); + assert!(pending.insert(0, 0, mask(0.9)).unwrap().is_some()); + assert_eq!(pending.bytes, 0); + assert!(pending.insert(1, 0, mask(0.9)).is_err()); + assert!(pending.frames.is_empty()); + } +} diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index 72c48e4e1..a80052bab 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -120,6 +120,11 @@ stored frame, pruned to the frames later ones can still look at) and refuses to one is set, with a clear error instead of running out of memory. Two runs of the same prompts give bit-identical masks on the CPU. +Pending per-object probability maps have a separate 256 MiB budget, including their allocated capacity and frame +slots. This bounds the full-resolution results waiting to be united when objects first prompted on different frames +report in separate passes. Reaching the budget cancels tracking and reports the reason through the existing +`trackerNote`; propagation falls back to the classic engine. A single object's results need no pending buffer. + ## How it works `crates/sam2` is a Rust port (on [candle](https://github.com/huggingface/candle), native only) of the Hugging Face `transformers` SAM 2 implementation, Apache-2.0 (see its `NOTICE`). Every stage is checked against PyTorch by From 6d10a6e932410a38681d5b8c1f1ef19ca4f26c14 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Sat, 17 Oct 2026 20:48:55 +0000 Subject: [PATCH 53/77] S2.75: propagate saved tracking with the newly selected model --- crates/effects/src/roto.rs | 58 ++++++++++++++++++++++++++++++ crates/engine/src/analysis/roto.rs | 5 ++- crates/engine/src/tests/roto.rs | 29 +++++++++++++++ 3 files changed, 91 insertions(+), 1 deletion(-) diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index e8cae42c3..5bcb7386d 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -598,6 +598,14 @@ impl Chain { self.model.as_ref().is_some_and(|m| m.tracker().is_some()) } + /// Explicit propagation replaces tracking made by another model. Rendering and Freeze keep + /// using saved mattes even when that model is unavailable, so this is only called for Propagate. + pub fn prepare_propagation(&mut self) { + if self.saved_model().is_some_and(|saved| self.model.as_ref().is_none_or(|model| model.info().id != saved)) { + self.saved = None; + } + } + /// Follow the object through the span with the model's tracker, from the strokes on their /// frames (only the frames without a tracked result yet: after a correction stroke, the frames it /// changes, tracked on from the nearest prompted frame), and cache every frame's segmentation (under the tracker's keys, with its @@ -1374,6 +1382,56 @@ mod tests { assert_eq!(completed, 0, "partial objects are not complete tracked frames"); } + #[test] + fn propagation_replaces_saved_tracking_from_a_different_model() { + let mut d = RotoData::default(); + d.add(Stroke { kind: StrokeKind::Fg, frame: 0, radius: 3.0, points: vec![[35.0, 40.0]], object: 0 }, 2, [0, 24]); + let mut params = params_with(&d); + params.values.insert(INPUT_KEY.into(), Value::Str("replace-saved-tracking-model".into())); + let size = [100.0, 80.0]; + let old_model = effectcraft_segment::SAM2_SMALL.id; + let old_keys = d.chain_keys(seed_with(¶ms, size, 25.0, old_model)); + let empty = vec![0; 100 * 80]; + let old_frame = TrackFrame { m: rle::encode(&empty), r: rle::encode(&empty), c: 0.2, ..Default::default() }; + let saved = TrackCache { + key: tracks_key(&old_keys), + model: old_model.into(), + size: [100, 80], + frames: old_keys.keys().map(|f| (*f, old_frame.clone())).collect(), + ..Default::default() + }; + params.values.insert(TRACKED.into(), Value::Str(saved.to_json())); + let mut chain = Chain::new(¶ms, size, 25.0, 1.0); + let fake = Arc::new(TwoObjects(Default::default())); + chain.model = Some(fake.clone()); + chain.keys = d.chain_keys(seed_with(¶ms, size, 25.0, fake.info().id)); + let mut get = |f: i64| Some(Arc::new(frame(f as f64 / 25.0))); + assert_eq!(chain.saved_model(), Some(old_model)); + assert_eq!(chain.ensure(0, &mut get, &|| false).unwrap().area(), 0, "rendering keeps the saved matte until propagation"); + chain.prepare_propagation(); + assert_eq!(chain.saved_model(), None); + assert_eq!(chain.track(&mut get, &|| false, &mut |_| {}).unwrap().unwrap().frames, 3); + assert_eq!(fake.0.lock().unwrap().len(), 1, "the newly selected tracker must actually run"); + assert!(chain.ensure(0, &mut get, &|| false).unwrap().area() > 0); + let replacement = tracks_of(&chain).expect("new tracking is saved"); + assert_eq!(replacement.model, effectcraft_segment::SAM2_TINY.id); + assert_eq!(replacement.frames.len(), 3); + params.values.insert(TRACKED.into(), Value::Str(replacement.to_json())); + let mut reopened = Chain::new(¶ms, size, 25.0, 1.0); + assert_eq!(reopened.saved_model(), Some(effectcraft_segment::SAM2_TINY.id)); + assert!(reopened.ensure(0, &mut get, &|| false).unwrap().area() > 0, "reopening uses the replacement without loading its model"); + reopened.model = Some(fake.clone()); + reopened.keys = chain.keys.clone(); + reopened.prepare_propagation(); + assert_eq!(reopened.saved_model(), Some(effectcraft_segment::SAM2_TINY.id)); + assert_eq!(reopened.track(&mut get, &|| false, &mut |_| {}).unwrap().unwrap().frames, 0, "matching saved tracking is reused"); + assert_eq!(fake.0.lock().unwrap().len(), 1); + let mut classic = Chain::new(¶ms, size, 25.0, 1.0); + classic.model = None; + classic.prepare_propagation(); + assert_eq!(classic.saved_model(), None, "Classic propagation also replaces saved model tracking"); + } + #[test] fn objects_are_followed_in_one_pass_and_their_mattes_united() { let mut d = RotoData::default(); diff --git a/crates/engine/src/analysis/roto.rs b/crates/engine/src/analysis/roto.rs index 399da1869..7a240f9df 100644 --- a/crates/engine/src/analysis/roto.rs +++ b/crates/engine/src/analysis/roto.rs @@ -446,7 +446,10 @@ async fn run_work(w: Work, shared: &RotoShared) { s.object_matte = w.object_matte; } let cancel = || shared.cancel.load(Ordering::Relaxed); - let chain = fx::Chain::new(&w.params, w.size, w.fps, 1.0); + let mut chain = fx::Chain::new(&w.params, w.size, w.fps, 1.0); + if w.task == RotoTask::Propagate && !w.object_matte { + chain.prepare_propagation(); + } let model = fx::installed_model(); let key = if w.object_matte { om::frozen_key(&w.params, w.size, w.fps, model.as_deref()) } else { fx::frozen_key(&w.params, w.size, w.fps) }; let mut frozen = FrozenCache { key, ..Default::default() }; diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index e4db67647..b28490a57 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -827,6 +827,9 @@ fn saved_tracking_is_kept_in_the_project_and_applies_when_it_reopens() { assert_eq!(st["source"], "tracker", "{st}"); assert!(st["notice"].as_str().unwrap().contains("Propagate again"), "{st}"); // A classic Propagate replaces the saved tracking with nothing. + // The synthetic source is not a file on disk; give the reopened session its provider before + // asking it to recompute frames. The checks above intentionally rendered without that provider. + s2.footage = s.footage.clone(); s2.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})).unwrap(); let st = s2.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); assert_eq!(st["source"], "classical", "{st}"); @@ -834,6 +837,32 @@ fn saved_tracking_is_kept_in_the_project_and_applies_when_it_reopens() { let _ = std::fs::remove_dir_all(std::path::Path::new(&path).parent().unwrap()); } +/// Propagate replaces a fully saved span when a different tracking model is active. +#[test] +fn propagating_saved_tracking_uses_the_new_model() { + struct OtherModel; + impl effectcraft_segment::MaskModel for OtherModel { + fn info(&self) -> &'static effectcraft_segment::ModelInfo { + &effectcraft_segment::SAM2_SMALL + } + fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &effectcraft_segment::Prompt) -> effectcraft_segment::Result> { + Err("no weights".into()) + } + fn tracker(&self) -> Option<&dyn effectcraft_segment::ObjectTracker> { + Some(&Exact) + } + } + let (mut s, clip, _uid, path) = saved_project("model-switch"); + effectcraft_effects::roto::set_model(Some(std::sync::Arc::new(OtherModel))); + let result = s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})); + effectcraft_effects::roto::set_model(None); + result.unwrap(); + let status = s.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); + assert_eq!(status["tracker"]["saved"], "sam2.1-hiera-small", "{status}"); + assert_eq!(status["tracker"]["frames"].as_object().unwrap().len(), 21, "{status}"); + let _ = std::fs::remove_dir_all(std::path::Path::new(&path).parent().unwrap()); +} + /// A tracked clip saved to a project file: the session that tracked it, its layer and instance, and the /// file. fn saved_project(name: &str) -> (Session, LayerId, u64, String) { From 66d8e5e94871719e1c3bc238c0747a332dbfd7f9 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 13:15:09 +0000 Subject: [PATCH 54/77] SAM 2 roto tracker: review fixes (direction-limited tracking, non-blocking model reload, per-object start passes, loader dispatch, CLI errors, hints) Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01RDvxCS8xcR28VArjzn1x7U --- apps/effectcraft-cli/src/main.rs | 40 ++++-- apps/effectcraft-cli/tests/it/cli.rs | 64 ++++++++++ crates/effects/src/catalog.rs | 3 +- crates/effects/src/roto.rs | 42 ++++++- crates/engine/src/analysis/models.rs | 10 +- crates/engine/src/analysis/roto.rs | 2 +- crates/engine/src/tests/models.rs | 28 +++++ crates/engine/src/tests/roto.rs | 3 + crates/sam2/src/gpu/mod.rs | 11 +- crates/sam2/src/mask.rs | 39 ++++-- crates/sam2/src/session.rs | 115 +++++++++++++++++- crates/sam2/src/track.rs | 31 ++--- crates/segment/README.md | 3 + crates/segment/src/lib.rs | 41 ++++++- crates/ui-egui/tests/it/main.rs | 1 + crates/ui-egui/tests/it/roto_model_guard.rs | 26 ++++ crates/ui-egui/tests/ui_roto_objects.rs | 3 +- .../ui-egui/tests/ui_roto_span_confidence.rs | 3 +- docs/agents.md | 2 +- 19 files changed, 410 insertions(+), 57 deletions(-) create mode 100644 crates/ui-egui/tests/it/roto_model_guard.rs diff --git a/apps/effectcraft-cli/src/main.rs b/apps/effectcraft-cli/src/main.rs index 1c528c8fb..dc47671e2 100644 --- a/apps/effectcraft-cli/src/main.rs +++ b/apps/effectcraft-cli/src/main.rs @@ -1013,26 +1013,46 @@ fn bench_cmd(args: &Args) -> Result<(), Failure> { Ok(()) } +/// The model a bench option asks for: the file named on the command line, where failing to read or +/// build it is an error saying why (not "not installed"); else the installed file, where a file that +/// cannot be used just leaves the model out. +fn bench_weights( + flag: &str, + named: Option<&str>, + installed: Option, + build: impl Fn(&[u8]) -> Result, +) -> Result, Failure> { + match named { + Some(path) => { + let bytes = std::fs::read(path).map_err(|e| Failure::Error(format!("{flag} {path}: {e}")))?; + build(&bytes).map(Some).map_err(|e| Failure::Error(format!("{flag} {path}: {e}"))) + } + None => Ok(installed.and_then(|p| std::fs::read(p).ok()).and_then(|bytes| build(&bytes).ok())), + } +} + /// `bench --roto`: the Roto Brush's trackers on synthetic clips whose true mattes are known. fn bench_roto(args: &Args) -> Result<(), Failure> { use effectcraft_segment::{Loaded, MOBILE_SAM, MaskModel}; use effectcraft_track::roto::bench; let frames = args.num("--frames")?.unwrap_or(24.0).clamp(12.0, 240.0) as usize; let session = effectcraft_host::session(); - // `--mobilesam FILE` uses the weights at FILE; otherwise the installed ones. - let mobilesam_file = args.opt("--mobilesam").map(std::path::PathBuf::from).or_else(|| session.model_path(&MOBILE_SAM)); - let mobilesam: Option> = - mobilesam_file.and_then(|p| std::fs::read(p).ok()).and_then(|bytes| effectcraft_segment::load(MOBILE_SAM.id, &bytes).ok()).and_then(|m| match m { - Loaded::Mask(m) => Some(m), - _ => None, - }); + // `--mobilesam FILE` uses the weights at FILE (a file that cannot be used is an error); + // otherwise the installed ones. + let mobilesam: Option> = bench_weights("--mobilesam", args.opt("--mobilesam"), session.model_path(&MOBILE_SAM), |bytes| { + match effectcraft_segment::load(MOBILE_SAM.id, bytes)? { + Loaded::Mask(m) => Ok(m), + _ => Err("not a mask model".to_string()), + } + })?; if mobilesam.is_none() { note!("bench --roto: MobileSAM is not installed; comparing the classical engine only"); } - // `--sam2 FILE` uses the weights at FILE; otherwise the installed ones. - let sam2_file = args.opt("--sam2").map(std::path::PathBuf::from).or_else(|| effectcraft_sam2::mask::SIZES.iter().find_map(|m| session.model_path(m))); + // `--sam2 FILE` uses the weights at FILE (a file that cannot be used is an error); otherwise + // the installed ones. + let installed_sam2 = effectcraft_sam2::mask::SIZES.iter().find_map(|m| session.model_path(m)); let sam2: Option> = - sam2_file.and_then(|p| std::fs::read(p).ok()).and_then(|bytes| effectcraft_sam2::mask::load(&bytes).ok()).map(|m| m as std::sync::Arc); + bench_weights("--sam2", args.opt("--sam2"), installed_sam2, |bytes| effectcraft_sam2::mask::load(bytes).map(|m| m as std::sync::Arc))?; if sam2.is_none() { note!("bench --roto: no SAM 2.1 size is installed; no SAM 2 column"); } diff --git a/apps/effectcraft-cli/tests/it/cli.rs b/apps/effectcraft-cli/tests/it/cli.rs index 5348f2a3a..ed6d98086 100644 --- a/apps/effectcraft-cli/tests/it/cli.rs +++ b/apps/effectcraft-cli/tests/it/cli.rs @@ -678,6 +678,70 @@ fn render_bitrate_applies_to_vp9_webm() { let _ = std::fs::remove_dir_all(&root); } +/// `bench --roto --clip FRAMES --truth MASKS` runs the trackers on the user's own image sequences: here a bright disc +/// drifting over a dark textured frame, with its true masks. Mismatched sequences are refused with a message. +#[test] +fn bench_roto_runs_on_a_users_clip() { + let (frames, masks) = (tmp("clip-frames"), tmp("clip-masks")); + for d in [&frames, &masks] { + let _ = std::fs::remove_dir_all(d); + std::fs::create_dir_all(d).unwrap(); + } + let (w, h) = (96u32, 72u32); + for f in 0..6u32 { + let (cx, cy) = (30.0 + 4.0 * f as f32, 36.0f32); + let (mut img, mut mask) = (image::RgbImage::new(w, h), image::GrayImage::new(w, h)); + for y in 0..h { + for x in 0..w { + let inside = ((x as f32 - cx).powi(2) + (y as f32 - cy).powi(2)).sqrt() < 14.0; + let tex = 40 + ((x * 7 + y * 13) % 30) as u8; + img.put_pixel(x, y, if inside { image::Rgb([230, 120, 40 + (x % 20) as u8]) } else { image::Rgb([tex / 2, tex, tex + 20]) }); + mask.put_pixel(x, y, image::Luma([if inside { 255 } else { 0 }])); + } + } + img.save(frames.join(format!("f{f:03}.png"))).unwrap(); + mask.save(masks.join(format!("f{f:03}.png"))).unwrap(); + } + let (f, m) = (frames.to_str().unwrap(), masks.to_str().unwrap()); + let out = bin().args(["bench", "--roto", "--clip", f, "--truth", m, "--frames", "6"]).output().unwrap(); + let text = String::from_utf8_lossy(&out.stdout).to_string(); + assert!(out.status.success(), "{text}\n{}", String::from_utf8_lossy(&out.stderr)); + let row = text.lines().find(|l| l.contains("classical")).unwrap_or_else(|| panic!("no classical row: {text}")); + let mean: f64 = row.split_whitespace().nth(2).unwrap().parse().unwrap(); + assert!(mean > 0.8, "classical mean IoU on the user's clip: {mean}\n{text}"); + // one mask short: refused + std::fs::remove_file(masks.join("f005.png")).unwrap(); + let out = bin().args(["bench", "--roto", "--clip", f, "--truth", m]).output().unwrap(); + assert!(!out.status.success()); + assert!(String::from_utf8_lossy(&out.stderr).contains("must match"), "{}", String::from_utf8_lossy(&out.stderr)); + // --clip without --truth + let out = bin().args(["bench", "--roto", "--clip", f]).output().unwrap(); + assert!(!out.status.success()); +} + +/// `bench --roto --sam2 FILE` / `--mobilesam FILE` say why a file named on the command line cannot be +/// used (a missing file, one that is not the published weights), instead of "not installed". +#[test] +fn bench_roto_reports_why_a_named_model_file_failed() { + let dir = tmp("bench-weights"); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + let bogus = dir.join("bogus.bin"); + std::fs::write(&bogus, b"not the weights").unwrap(); + for flag in ["--sam2", "--mobilesam"] { + let missing = dir.join("missing.bin"); + let out = bin().args(["bench", "--roto", flag, missing.to_str().unwrap()]).output().unwrap(); + let err = String::from_utf8_lossy(&out.stderr).to_string(); + assert!(!out.status.success(), "{flag}: {err}"); + assert!(err.contains(flag) && err.contains("missing.bin"), "{flag} names the file that could not be read: {err}"); + let out = bin().args(["bench", "--roto", flag, bogus.to_str().unwrap()]).output().unwrap(); + let err = String::from_utf8_lossy(&out.stderr).to_string(); + assert!(!out.status.success(), "{flag}: {err}"); + assert!(err.contains("does not match"), "{flag} says the file is not the published weights: {err}"); + } + let _ = std::fs::remove_dir_all(&dir); +} + /// Render uses the same empty/project/demo selection as the other CLI commands. #[test] fn render_empty_does_not_load_the_demo() { diff --git a/crates/effects/src/catalog.rs b/crates/effects/src/catalog.rs index 332968baf..7fcb024bf 100644 --- a/crates/effects/src/catalog.rs +++ b/crates/effects/src/catalog.rs @@ -128,7 +128,8 @@ pub fn effects_markdown() -> String { s.push_str("# Effects\n\n"); s.push_str("\n\n"); - let params: usize = r.iter().map(|e| e.params.len()).sum(); + // The parameters the table's rows list: the shown ones, not the hidden bookkeeping ones. + let params: usize = r.iter().map(|e| e.params.iter().filter(|p| !matches!(p.ui, effectcraft_project::ParamUi::Hidden)).count()).sum(); let gpu = r.iter().filter(|e| e.gpu).count(); let float = r.iter().filter(|e| e.float).count(); let partial = r.iter().filter(|e| PARTIAL.iter().any(|(id, _)| *id == e.id)).count(); diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 5bcb7386d..235c26afa 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -588,9 +588,12 @@ impl Chain { cur } - /// Whether frame `f` is computed (at full resolution or at `scale`). + /// Whether frame `f` is computed (at full resolution or at `scale`): in this session's caches, + /// or in the tracking saved with the project, which needs no work to be used even before it has + /// been decoded (and whichever model Settings choose now). pub fn is_cached(&self, f: i64, scale: f64) -> bool { self.keys.get(&f).is_some_and(|k| [*k, tracked_key(*k)].iter().any(|k| cached(*k, 1.0).is_some() || cached(*k, scale).is_some())) + || self.saved.as_ref().is_some_and(|t| t.cache.frames.contains_key(&f) && t.keys.contains_key(&f)) } /// Whether the chosen model can follow an object through a clip (SAM 2). @@ -618,11 +621,24 @@ impl Chain { cancel: &dyn Fn() -> bool, progress: &mut dyn FnMut(usize), ) -> Result, String> { - self.track_with_pending_limit(frame, cancel, progress, pending::MAX_PENDING_BYTES) + self.track_frames(None, frame, cancel, progress) + } + + /// [`Chain::track`] limited to `wanted` frames (a one-direction Propagate): the window, the work and the + /// progress count cover only those frames of the span. + pub fn track_frames( + &self, + wanted: Option<&[i64]>, + frame: &mut dyn FnMut(i64) -> Option>, + cancel: &dyn Fn() -> bool, + progress: &mut dyn FnMut(usize), + ) -> Result, String> { + self.track_with_pending_limit(wanted, frame, cancel, progress, pending::MAX_PENDING_BYTES) } fn track_with_pending_limit( &self, + wanted: Option<&[i64]>, frame: &mut dyn FnMut(i64) -> Option>, cancel: &dyn Fn() -> bool, progress: &mut dyn FnMut(usize), @@ -635,6 +651,7 @@ impl Chain { // frames from it on (both directions), so only those come back here. let (lo_span, hi_span) = (self.data.span[0], self.data.span[1]); let needed: Vec = (lo_span..=hi_span) + .filter(|f| wanted.is_none_or(|w| w.contains(f))) .filter(|f| self.keys.get(f).is_some_and(|k| self.saved_seg(*f, self.opts.scale).is_none() && cached(tracked_key(*k), self.opts.scale).is_none())) .collect(); let (Some(&need_lo), Some(&need_hi)) = (needed.first(), needed.last()) else { @@ -679,6 +696,9 @@ impl Chain { // A frame is finished when every object has reported it (the objects' mattes are united). let mut pending = pending::PendingMasks::new(objects.len(), max_pending_bytes); let finish = |g: i64, parts: Vec>, summary: &mut TrackSummary, progress: &mut dyn FnMut(usize)| { + if !wanted.is_none_or(|w| w.contains(&g)) { + return; + } let (Some(key), Some(img)) = (self.keys.get(&g).copied(), (frame.borrow_mut())(g)) else { return }; let mut segs: Vec<(u32, effectcraft_track::roto::FrameSeg)> = vec![]; let (mut confidence, mut present) = (f32::INFINITY, false); @@ -1375,13 +1395,29 @@ mod tests { let mut get = |f: i64| Some(Arc::new(frame(f as f64 / 25.0))); let mut completed = 0; // Each map takes 32,000 bytes; two incomplete frames exceed this small test budget. - let error = chain.track_with_pending_limit(&mut get, &|| false, &mut |n| completed = n, 64_000).unwrap_err(); + let error = chain.track_with_pending_limit(None, &mut get, &|| false, &mut |n| completed = n, 64_000).unwrap_err(); assert!(error.contains("pending object masks") && error.contains("memory budget"), "{error}"); assert!(fake.cancelled.load(Relaxed), "the model must see cancellation before allocating the rest of the clip"); assert_eq!(fake.emitted.load(Relaxed), 2); assert_eq!(completed, 0, "partial objects are not complete tracked frames"); } + #[test] + fn a_one_direction_pass_tracks_and_reports_only_its_frames() { + let mut d = RotoData::default(); + d.add(Stroke { kind: StrokeKind::Fg, frame: 12, radius: 3.0, points: vec![[35.0, 40.0]], object: 0 }, 12, [0, 24]); + let params = params_with(&d); + let mut chain = Chain::new(¶ms, [100.0, 80.0], 25.0, 1.0); + chain.model = Some(Arc::new(TwoObjects(Default::default()))); + let mut get = |f: i64| Some(Arc::new(frame(f as f64 / 25.0))); + let forward: Vec = (12..=24).collect(); + let mut last = 0; + let summary = chain.track_frames(Some(&forward), &mut get, &|| false, &mut |n| last = n).unwrap().unwrap(); + assert_eq!((summary.frames, last), (forward.len(), forward.len()), "progress never exceeds the frames asked for"); + let _ = chain.ensure(5, &mut get, &|| false); + assert!(tracks_of(&chain).is_none_or(|t| t.frames.keys().all(|f| *f >= 12)), "earlier frames were not tracked"); + } + #[test] fn propagation_replaces_saved_tracking_from_a_different_model() { let mut d = RotoData::default(); diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index d166702ac..ae15fe1d3 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -311,11 +311,15 @@ impl Session { } } - /// Load the chosen Roto Brush model again if it was unloaded for being idle (before a stroke uses it). + /// Start loading the chosen Roto Brush model again if it was unloaded for being idle (before a stroke uses it). + /// The load runs on the model worker, so the stroke is never blocked on it (the classic engine segments it + /// meanwhile and `poll_models` installs the model when it arrives), and a model that already failed to load is + /// not retried on every stroke: choosing it again clears the error. pub fn reload_roto_model(&mut self) { let chosen = self.chosen_model(Task::Mask).to_string(); - if chosen != seg::CLASSICAL && self.models.status(Task::Mask).active.is_none() && self.models.status(Task::Mask).busy.is_none() { - self.apply_model(Task::Mask, true); + let st = self.models.status(Task::Mask); + if chosen != seg::CLASSICAL && st.active.is_none() && st.busy.is_none() && st.error.is_none() { + self.apply_model(Task::Mask, false); } } diff --git a/crates/engine/src/analysis/roto.rs b/crates/engine/src/analysis/roto.rs index 7a240f9df..a07b0010e 100644 --- a/crates/engine/src/analysis/roto.rs +++ b/crates/engine/src/analysis/roto.rs @@ -465,7 +465,7 @@ async fn run_work(w: Work, shared: &RotoShared) { s.fps = if s.elapsed > 0.0 { s.done as f64 / s.elapsed } else { 0.0 }; crate::offload::report(s.done, s.total); }; - let result = chain.track(&mut get, &cancel, &mut report); + let result = chain.track_frames(Some(&w.frames), &mut get, &cancel, &mut report); if cancel() { let mut s = lock(&shared.state); s.cancelled = true; diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index ed0eecf9b..84bbd2381 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -161,6 +161,34 @@ fn an_idle_roto_model_is_unloaded() { let _ = std::fs::remove_dir_all(d); } +/// A stroke reloading an unloaded model never waits for the load and never retries one that already failed. +#[cfg(not(target_arch = "wasm32"))] +#[test] +fn reloading_for_a_stroke_does_not_block_or_retry_a_failed_load() { + super::roto::hold_roto_cache_test_lock(); + let (mut s, d) = session("reload-failed"); + let info = effectcraft_segment::info("mobilesam").unwrap(); + let dir = s.models_dir().unwrap(); + std::fs::create_dir_all(&dir).unwrap(); + // the right size, the wrong content: found as installed, then fails to load + std::fs::File::create(dir.join(info.file_name)).unwrap().set_len(info.size).unwrap(); + s.execute_checked("roto.model.select", json!({"id": "mobilesam"})).unwrap(); + let t0 = std::time::Instant::now(); + while s.models.status(Task::Mask).busy.is_some() && t0.elapsed().as_secs() < 30 { + std::thread::sleep(std::time::Duration::from_millis(5)); + } + let failed = s.models.status(Task::Mask).error; + assert!(failed.is_some(), "the unreadable file fails to load"); + s.reload_roto_model(); + let st = s.models.status(Task::Mask); + assert!(st.busy.is_none(), "a failed model is not loaded again by every stroke"); + assert_eq!(st.error, failed); + // an unloaded, not failed model loads on the worker: the call returns with the load still pending or done + s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); + crate::effects::roto::set_model(None); + let _ = std::fs::remove_dir_all(d); +} + /// With `EFFECTCRAFT_DOWNLOAD_SIZE=sam2.1-hiera-small` (or another registry id): the in-app download (Settings ▸ Roto /// Brush ▸ Download) fetches the pinned file over the network, verifies it, installs it with its notice and loads it. /// Skipped without the variable: it needs the network and hundreds of megabytes. diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index b28490a57..6c57d5c29 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -705,6 +705,7 @@ fn sam2_tracks_the_clip_through_roto_propagate() { let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); assert_eq!(st["source"], "tracker", "{st}"); assert!(st["notice"].as_str().unwrap().contains("Propagate again"), "{st}"); + assert_eq!(st["computed"].as_array().unwrap().len(), 21, "still computed under another model: {st}"); // A classic Propagate: its own segmentation, no tracker data. s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})).unwrap(); let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); @@ -808,6 +809,8 @@ fn saved_tracking_is_kept_in_the_project_and_applies_when_it_reopens() { assert_eq!(st["source"], "tracker", "{st}"); assert_eq!(st["tracker"]["saved"], sam2, "{st}"); assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 21, "{st}"); + // The span bar's "computed" frames include the saved ones before anything has been decoded. + assert_eq!(st["computed"].as_array().unwrap().len(), 21, "{st}"); let j = iou_at(&mut s2, clip, 12, center(12)); assert!(j > 0.9, "IoU of a saved frame: {j}"); // The render path (Chain::ensure) reads the same frame. diff --git a/crates/sam2/src/gpu/mod.rs b/crates/sam2/src/gpu/mod.rs index 152c67c50..fab8573e1 100644 --- a/crates/sam2/src/gpu/mod.rs +++ b/crates/sam2/src/gpu/mod.rs @@ -96,15 +96,24 @@ pub fn half_weights() -> bool { std::env::var("EFFECTCRAFT_SAM2_F16").is_ok_and(|v| matches!(v.trim(), "1" | "on" | "true")) } +/// The process's one GPU device, once something has asked for it. +static DEVICE: std::sync::OnceLock, String>> = std::sync::OnceLock::new(); + impl Gpu { /// The process's one GPU device, created on first use. Several devices at once upset some Vulkan loaders /// (a crash when two threads create or drop them together), and one is all SAM 2 needs: models come and go, /// the device stays. An error (no GPU) is remembered too. pub fn shared() -> Result> { - static DEVICE: std::sync::OnceLock, String>> = std::sync::OnceLock::new(); DEVICE.get_or_init(|| Gpu::new().map(Arc::new).map_err(|e| e.to_string())).clone().map_err(Error::Compute) } + /// Whether the shared device has been looked for yet, and found: `None` before anything has asked + /// for it (this never creates it, so a caller that only wants to describe the GPU does not stall + /// on the adapter and shader compilation), then `Some(true)` or `Some(false)` (no GPU). + pub fn probe() -> Option { + DEVICE.get().map(|d| d.is_ok()) + } + /// A device on the best real GPU (a software adapter is refused: the CPU path is faster). fn new() -> Result { let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle()); diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index 8470bd4f3..911ff2e00 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -156,17 +156,29 @@ pub fn recommended() -> &'static ModelInfo { /// A one-line note for the Settings list: what was measured for this size, and what the bigger ones cost without a GPU. pub fn hint(id: &str) -> Option { - let gpu = crate::gpu::Gpu::shared().is_ok(); + hint_for(id, crate::gpu::Gpu::probe()) +} + +/// [`hint`] given whether a GPU is known (`None`: not looked for yet). The list is built on the +/// Settings path, so it never creates the GPU device itself: that happens when a model is loaded, +/// and after that the notes say whether one was found. +fn hint_for(id: &str, gpu: Option) -> Option { Some( match id { "sam2.1-hiera-tiny" => "Recommended: the smallest download, and no larger size tracked better on the benchmark clips.", "sam2.1-hiera-small" => "As fast as Tiny on a GPU and no better on the benchmark clips; worth a try on difficult footage.", - "sam2.1-hiera-base-plus" if !gpu => "Slow without a GPU (several seconds a frame); none was found.", - "sam2.1-hiera-base-plus" => "About 15% slower than Tiny on a GPU and no better on the benchmark clips.", - "sam2.1-hiera-large" if !gpu => "Slow without a GPU (many seconds a frame); none was found.", - "sam2.1-hiera-large" => { - "About 1.7 times Tiny's time on a GPU; it lost the object after it re-entered in one benchmark clip. Try it on difficult footage." - } + "sam2.1-hiera-base-plus" => match gpu { + Some(false) => "Slow without a GPU (several seconds a frame); none was found.", + Some(true) => "About 15% slower than Tiny on a GPU and no better on the benchmark clips.", + None => "About 15% slower than Tiny on a GPU and no better on the benchmark clips; several seconds a frame without one.", + }, + "sam2.1-hiera-large" => match gpu { + Some(false) => "Slow without a GPU (many seconds a frame); none was found.", + Some(true) => { + "About 1.7 times Tiny's time on a GPU; it lost the object after it re-entered in one benchmark clip. Try it on difficult footage." + } + None => "About 1.7 times Tiny's time on a GPU, many seconds a frame without one; it lost the object after it re-entered in one benchmark clip.", + }, _ => return None, } .to_string(), @@ -191,6 +203,19 @@ mod tests { assert!(v.iter().all(|x| x.is_finite())); } + /// The Settings list asks for hints without creating the GPU device: before anything has loaded a + /// model there is no answer yet, and the notes say so rather than guess. + #[test] + fn hints_do_not_need_the_gpu_to_have_been_created() { + for id in ["sam2.1-hiera-base-plus", "sam2.1-hiera-large"] { + let (unknown, none, found) = (hint_for(id, None), hint_for(id, Some(false)), hint_for(id, Some(true))); + assert!(unknown.is_some() && none.is_some() && found.is_some()); + assert!(none.as_deref().is_some_and(|h| h.contains("none was found")), "{id}: {none:?}"); + assert!(unknown.as_deref().is_some_and(|h| h.contains("without one") && !h.contains("none was found")), "{id}: {unknown:?}"); + } + assert!(hint_for("mobilesam", None).is_none()); + } + #[test] fn a_wrong_file_is_refused_before_parsing() { let e = load(b"not the weights").err().unwrap_or_default(); diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 336518bd5..54a39d212 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -615,15 +615,118 @@ impl<'a> VideoTracker<'a> { if let Some(limit) = self.budget.max_seconds.filter(|l| started.elapsed().as_secs_f64() > *l) { return Err(Error::Budget(format!("tracking took more than {limit} s"))); } - let cached = self.objects.get(id).is_some_and(|o| o.cond.contains_key(&frame)); - let stored = self.objects.get(id).and_then(|o| o.cond.get(&frame)).and_then(|o| o.result.clone()); - let mask = match stored { - Some(m) if cached => m, - _ => self.step(*id, frame, None, reverse, false)?, - }; + let mask = self.mask_at(*id, frame, reverse)?; out(mask); } } Ok(()) } + + /// Object `id` on `frame`: the result stored with its prompt there, else the tracker's step. + fn mask_at(&mut self, id: ObjectId, frame: usize, reverse: bool) -> Result { + let stored = self.objects.get(&id).and_then(|o| o.cond.get(&frame)).and_then(|o| o.result.clone()); + match stored { + Some(m) => Ok(m), + None => self.step(id, frame, None, reverse, false), + } + } + + /// Track every listed object from its own prompted frame (`starts`: object, frame) in one pass + /// in `direction`, frame by frame with all the objects that have started: a frame is finished + /// for all of them before the next begins, so a caller that assembles the objects of a frame + /// does not hold frames waiting for a pass that has not got there. A backward pass leaves out + /// each object's own start frame (the forward pass reported it). + pub fn propagate_each(&mut self, starts: &[(ObjectId, usize)], direction: Direction, cancel: &Cancel, out: &mut dyn FnMut(FrameMask)) -> Result<()> { + let n = self.frames.num_frames(); + for (id, start) in starts { + self.check_frame(*start)?; + if self.objects.get(id).is_none_or(|o| o.cond.is_empty()) { + return Err(Error::Input("no object has been prompted".into())); + } + } + if starts.is_empty() { + return Err(Error::Input("no object has been prompted".into())); + } + let reverse = direction == Direction::Backward; + let started = std::time::Instant::now(); + for (id, frame) in pass_order(starts, n, reverse) { + if cancel.is_cancelled() { + return Err(Error::Cancelled); + } + if let Some(limit) = self.budget.max_seconds.filter(|l| started.elapsed().as_secs_f64() > *l) { + return Err(Error::Budget(format!("tracking took more than {limit} s"))); + } + let mask = self.mask_at(id, frame, reverse)?; + out(mask); + } + Ok(()) + } +} + +/// The (object, frame) steps of [`VideoTracker::propagate_each`], in order: frame by frame from the +/// first start (the last, backwards), each frame for the objects whose pass has reached it. +fn pass_order(starts: &[(ObjectId, usize)], frames: usize, reverse: bool) -> impl Iterator + '_ { + let first = if reverse { starts.iter().map(|s| s.1).max() } else { starts.iter().map(|s| s.1).min() }; + let span = match first { + Some(f) if reverse => f + 1, + Some(f) => frames.saturating_sub(f), + None => 0, + }; + (0..span).flat_map(move |k| { + let frame = if reverse { first.unwrap_or(0).saturating_sub(k) } else { first.unwrap_or(0) + k }; + starts.iter().filter_map(move |(id, start)| { + let reached = if reverse { *start > frame } else { *start <= frame }; + reached.then_some((*id, frame)) + }) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn order(starts: &[(u32, usize)], frames: usize, reverse: bool) -> Vec<(u32, usize)> { + let starts: Vec<(ObjectId, usize)> = starts.iter().map(|(o, s)| (ObjectId(*o), *s)).collect(); + pass_order(&starts, frames, reverse).map(|(o, f)| (o.0, f)).collect() + } + + /// Forward, every frame is finished for all the objects that have started before the next frame + /// begins; an object joins the pass at its own start frame. + #[test] + fn the_forward_pass_goes_frame_by_frame_over_the_objects_that_have_started() { + assert_eq!(order(&[(1, 0), (2, 3)], 6, false), [(1, 0), (1, 1), (1, 2), (1, 3), (2, 3), (1, 4), (2, 4), (1, 5), (2, 5)]); + // Prompts on the same frame go together; the pass starts at the earliest start. + assert_eq!(order(&[(1, 2), (2, 2)], 4, false), [(1, 2), (2, 2), (1, 3), (2, 3)]); + assert_eq!(order(&[(1, 2)], 4, false), [(1, 2), (1, 3)]); + } + + /// Backward, from the last start down to frame 0, leaving out each object's own start frame. + #[test] + fn the_backward_pass_goes_frame_by_frame_and_skips_each_objects_own_start() { + assert_eq!(order(&[(1, 0), (2, 3)], 6, true), [(2, 2), (2, 1), (2, 0)]); + assert_eq!(order(&[(1, 2), (2, 3)], 6, true), [(2, 2), (1, 1), (2, 1), (1, 0), (2, 0)]); + assert!(order(&[(1, 0)], 6, true).is_empty(), "nothing before frame 0"); + } + + /// A frame is finished for every object before the next begins, however far apart the prompts + /// are: the number of frames waiting for an object never grows with the clip. + #[test] + fn no_frame_waits_for_an_object_that_has_not_reached_it() { + let steps = order(&[(1, 0), (2, 1000)], 2000, false); + let mut waiting = 0usize; + let mut worst = 0usize; + let mut seen = std::collections::BTreeMap::::new(); + for (_, f) in steps { + let n = seen.entry(f).or_default(); + *n += 1; + // From frame 1000 on, both objects report each frame back to back. + if f >= 1000 && *n == 2 { + waiting -= 1; + } else if f >= 1000 { + waiting += 1; + } + worst = worst.max(waiting); + } + assert!(worst <= 1, "{worst} frames waiting"); + } } diff --git a/crates/sam2/src/track.rs b/crates/sam2/src/track.rs index 1aff06d9d..ad920efe7 100644 --- a/crates/sam2/src/track.rs +++ b/crates/sam2/src/track.rs @@ -83,7 +83,6 @@ impl Sam2Mask { let mut tracker = VideoTracker::new(self.video_model().clone(), Box::new(source))?; let object_id = |oi: usize| ObjectId(oi as u32 + 1); // each object's first prompted frame: its pass starts there, forwards and backwards - let mut starts: std::collections::BTreeMap> = Default::default(); let mut first_of: HashMap = HashMap::new(); for (oi, idx, p) in prompt_frames { if cancel() { @@ -110,9 +109,8 @@ impl Sam2Mask { let e = first_of.entry(oi).or_insert(idx); *e = (*e).min(idx); } - for (oi, idx) in first_of { - starts.entry(idx).or_default().push(object_id(oi)); - } + let mut starts: Vec<(ObjectId, usize)> = first_of.into_iter().map(|(oi, idx)| (object_id(oi), idx)).collect(); + starts.sort_by_key(|(id, _)| id.0); let flag = Cancel::new(); let mut failure: Option = None; let mut emit = |m: FrameMask| { @@ -128,19 +126,22 @@ impl Sam2Mask { } } }; - for (start, ids) in starts { - tracker.propagate_objects(Some(&ids), start, Direction::Forward, None, &flag, &mut emit)?; - if start > 0 { - // the start frame was reported by the forward pass; the backward pass passes it again - tracker.propagate_objects(Some(&ids), start, Direction::Backward, None, &flag, &mut |m| { - if m.frame != start { - emit(m); - } - })?; + // One pass each way over all the objects, frame by frame, so the host can finish a frame (and + // let go of its masks) as soon as every object has it, however far apart the prompts are. + let passes = (|| { + tracker.propagate_each(&starts, Direction::Forward, &flag, &mut emit)?; + if starts.iter().any(|(_, start)| *start > 0) { + tracker.propagate_each(&starts, Direction::Backward, &flag, &mut emit)?; } - } + Ok(()) + })(); self.peak.store(tracker.peak_state_bytes() as u64, std::sync::atomic::Ordering::Relaxed); - failure.map_or(Ok(()), Err) + // A frame that could not be converted cancelled the passes; that is the error to report, not + // the cancellation it caused. + match failure { + Some(e) => Err(e), + None => passes, + } } } diff --git a/crates/segment/README.md b/crates/segment/README.md index 4962f8ef8..e6e4c6c98 100644 --- a/crates/segment/README.md +++ b/crates/segment/README.md @@ -37,6 +37,9 @@ kernels at run time). No ML framework, no unsafe code. |---|---|---|---|---|---|---| | `mobilesam` | Roto Brush | [MobileSAM](https://github.com/ChaoningZhang/MobileSAM): Segment Anything with a 5M-parameter TinyViT encoder | Chaoning Zhang et al., building on Meta AI's Segment Anything and Microsoft's TinyViT | Apache-2.0 | `mobile_sam.pt` from the official repository | 40.7 MB | | `sam2.1-hiera-tiny` | Roto Brush | [SAM 2.1 Tiny](https://github.com/facebookresearch/sam2): Hiera-Tiny video object tracker with streaming memory | Meta AI (Nikhila Ravi et al.) | Apache-2.0 | `model.safetensors` of `facebook/sam2.1-hiera-tiny` (Hugging Face, pinned to a commit); loaded by `effectcraft-sam2` | 155.9 MB | +| `sam2.1-hiera-small` | Roto Brush | [SAM 2.1 Small](https://github.com/facebookresearch/sam2): Hiera-Small video object tracker with streaming memory | Meta AI (Nikhila Ravi et al.) | Apache-2.0 | `model.safetensors` of `facebook/sam2.1-hiera-small` (Hugging Face, pinned to a commit); loaded by `effectcraft-sam2` | 184.3 MB | +| `sam2.1-hiera-base-plus` | Roto Brush | [SAM 2.1 Base+](https://github.com/facebookresearch/sam2): Hiera-Base+ video object tracker with streaming memory | Meta AI (Nikhila Ravi et al.) | Apache-2.0 | `model.safetensors` of `facebook/sam2.1-hiera-base-plus` (Hugging Face, pinned to a commit); loaded by `effectcraft-sam2` | 323.5 MB | +| `sam2.1-hiera-large` | Roto Brush | [SAM 2.1 Large](https://github.com/facebookresearch/sam2): Hiera-Large video object tracker with streaming memory | Meta AI (Nikhila Ravi et al.) | Apache-2.0 | `model.safetensors` of `facebook/sam2.1-hiera-large` (Hugging Face, pinned to a commit); loaded by `effectcraft-sam2` | 897.9 MB | | `mediapipe-face` | Face tracking | [MediaPipe Face Landmarker](https://developers.google.com/edge/mediapipe/solutions/vision/face_landmarker): BlazeFace detector and Face Mesh V2 | Google LLC | Apache-2.0 ([model card](https://storage.googleapis.com/mediapipe-assets/Model%20Card%20MediaPipe%20Face%20Mesh%20V2.pdf)) | `face_landmarker.task` (float16, version 1) from Google's model storage | 3.8 MB | How they were checked: diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index fb16f01ce..04894cb2b 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -281,14 +281,34 @@ pub fn load(id: &str, bytes: &[u8]) -> Result { if got != info.sha256 { return Err(format!("{}: the file does not match (SHA-256 {got}, expected {})", info.name, info.sha256)); } - match id { - "mobilesam" => { + match loader(id) { + Some(Loader::MobileSam) => { let w = pt::read(bytes)?; Ok(Loaded::Mask(Arc::new(Sam { net: mobilesam::MobileSam::from_weights(&w)? }))) } - "mediapipe-face" => Ok(Loaded::Face(Arc::new(mediapipe::FaceLandmarker::from_task(bytes)?))), - "sam2.1-hiera-tiny" => Err(format!("{} is loaded by the desktop application (effectcraft-sam2)", info.name)), - _ => Err(format!("no loader for `{id}`")), + Some(Loader::MediaPipe) => Ok(Loaded::Face(Arc::new(mediapipe::FaceLandmarker::from_task(bytes)?))), + Some(Loader::Desktop) => Err(format!("{} is loaded by the desktop application (effectcraft-sam2)", info.name)), + None => Err(format!("no loader for `{id}`")), + } +} + +/// What builds a registered model. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Loader { + MobileSam, + MediaPipe, + /// Built outside this crate, by the desktop application (every SAM 2.1 size). + Desktop, +} + +/// The loader for model `id`: one answer for every size of a family, so a new size cannot fall +/// through to "no loader" while its sibling gets the helpful error. +fn loader(id: &str) -> Option { + match id { + "mobilesam" => Some(Loader::MobileSam), + "mediapipe-face" => Some(Loader::MediaPipe), + id if id.starts_with("sam2.1-hiera-") => Some(Loader::Desktop), + _ => None, } } @@ -412,4 +432,15 @@ mod tests { assert!(load("mobilesam", b"not the weights").err().is_some_and(|e| e.contains("does not match"))); assert!(load("nope", b"").is_err()); } + + /// Every registered model has a loader, so none falls through to "no loader" after its checksum + /// matches; every SAM 2.1 size is loaded by the desktop application. + #[test] + fn every_registered_model_has_a_loader() { + for m in MODELS { + let l = loader(m.id); + assert!(l.is_some(), "{} has no loader", m.id); + assert_eq!(l == Some(Loader::Desktop), m.id.starts_with("sam2.1-"), "{}", m.id); + } + } } diff --git a/crates/ui-egui/tests/it/main.rs b/crates/ui-egui/tests/it/main.rs index 4fc34d9ff..9ef4c4b54 100644 --- a/crates/ui-egui/tests/it/main.rs +++ b/crates/ui-egui/tests/it/main.rs @@ -3,6 +3,7 @@ mod gpu_test; mod native_menu_text_keys; mod render_queue_snapshot; +mod roto_model_guard; mod shared_docking_ui; mod ui_3d; mod ui_accessibility; diff --git a/crates/ui-egui/tests/it/roto_model_guard.rs b/crates/ui-egui/tests/it/roto_model_guard.rs new file mode 100644 index 000000000..a65805b55 --- /dev/null +++ b/crates/ui-egui/tests/it/roto_model_guard.rs @@ -0,0 +1,26 @@ +//! The Roto Brush model is process-global (`effects::roto::set_model`), and every test in this binary shares the +//! process: a test that installs a scripted model holds this guard, so another cannot swap or clear the model +//! mid-run, and the model is cleared again even when the test panics. + +use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; + +use effectcraft_segment::MaskModel; + +static LOCK: Mutex<()> = Mutex::new(()); + +pub struct InstalledModel { + _exclusive: MutexGuard<'static, ()>, +} + +/// Install `model` for the roto effect until the returned guard drops. +pub fn install(model: Arc) -> InstalledModel { + let guard = LOCK.lock().unwrap_or_else(PoisonError::into_inner); + effectcraft_engine::effects::roto::set_model(Some(model)); + InstalledModel { _exclusive: guard } +} + +impl Drop for InstalledModel { + fn drop(&mut self) { + effectcraft_engine::effects::roto::set_model(None); + } +} diff --git a/crates/ui-egui/tests/ui_roto_objects.rs b/crates/ui-egui/tests/ui_roto_objects.rs index 5d16b0ada..252ecfb3d 100644 --- a/crates/ui-egui/tests/ui_roto_objects.rs +++ b/crates/ui-egui/tests/ui_roto_objects.rs @@ -66,7 +66,7 @@ fn the_boundary_view_outlines_each_object_and_show_object_picks_one() { for (object, x) in [(0, 22.0), (1, 58.0)] { s.execute("roto.stroke", json!({"layer": id, "object": object, "points": [[x, 38.0], [x, 42.0]], "radius": 4})).unwrap(); } - effectcraft_engine::effects::roto::set_model(Some(Arc::new(TwoRegions))); + let _model = crate::roto_model_guard::install(Arc::new(TwoRegions)); s.execute("roto.propagate", json!({"layer": id, "direction": "both", "wait": true})).unwrap(); let st = s.execute("roto.status", json!({"layer": id})).unwrap(); assert_eq!(st["objects"], json!([{"id": 0, "strokes": 1}, {"id": 1, "strokes": 1}]), "{st}"); @@ -96,5 +96,4 @@ fn the_boundary_view_outlines_each_object_and_show_object_picks_one() { let one = h.render().expect("the panel renders"); let (a, b) = (count(&one, pink), count(&one, cyan)); assert!(a < 10 && b > 50, "only object 1 is shown: {a} pink, {b} cyan pixels"); - effectcraft_engine::effects::roto::set_model(None); } diff --git a/crates/ui-egui/tests/ui_roto_span_confidence.rs b/crates/ui-egui/tests/ui_roto_span_confidence.rs index f3a8bcd9b..04379f1ae 100644 --- a/crates/ui-egui/tests/ui_roto_span_confidence.rs +++ b/crates/ui-egui/tests/ui_roto_span_confidence.rs @@ -70,7 +70,7 @@ fn the_span_bar_shows_the_trackers_confidence() { let id = s.active_comp().unwrap().layers.iter().find(|l| l.name == "Box").unwrap().id.0; s.execute("roto.stroke", json!({"layer": id, "points": [[20.0, 40.0], [60.0, 40.0]], "radius": 6})).unwrap(); s.execute("roto.stroke", json!({"layer": id, "kind": "bg", "points": [[5.0, 5.0], [75.0, 5.0]], "radius": 4})).unwrap(); - effectcraft_engine::effects::roto::set_model(Some(Arc::new(Scripted))); + let _model = crate::roto_model_guard::install(Arc::new(Scripted)); // the model stays loaded while the panel draws (its id is part of the cache keys) s.execute("roto.propagate", json!({"layer": id, "direction": "both", "wait": true})).unwrap(); @@ -115,5 +115,4 @@ fn the_span_bar_shows_the_trackers_confidence() { let got = sample(f); assert!(close(got, want), "frame {f} ({what}): drew {got:?}, wanted about {want:?}"); } - effectcraft_engine::effects::roto::set_model(None); } diff --git a/docs/agents.md b/docs/agents.md index ddd730557..ef88acccb 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -506,7 +506,7 @@ Trained model (Roto Brush 2.0 / 3.0): licence, size, URL, SHA-256 and whether it is installed, selected and active. - `roto.model.download {id}` fetches and verifies the official weights in the background; `roto.model.install {path, id?}` installs a file you already have. -- `roto.model.select {id: classical|mobilesam|sam2.1-hiera-tiny}` chooses which one Roto Brush uses (SAM 2 also tracks the object through the clip, see [roto-brush-sam2.md](roto-brush-sam2.md); `roto.status` then reports `source` and per-frame `tracker.frames` with `confidence` and `occluded`). Loading runs +- `roto.model.select {id: classical|mobilesam|sam2.1-hiera-tiny|sam2.1-hiera-small|sam2.1-hiera-base-plus|sam2.1-hiera-large}` chooses which one Roto Brush uses (SAM 2 also tracks the object through the clip, see [roto-brush-sam2.md](roto-brush-sam2.md); `roto.status` then reports `source` and per-frame `tracker.frames` with `confidence` and `occluded`). Loading runs in the background; `"wait": true` (also on download and install) returns once it is in use. - `roto.propagate` never falls back to the classic engine silently: a chosen model that is not installed, or failed to load, is an error, and without `"wait": true` so is a model still loading. `roto.status` reports `model` From 020d5cf980aa18131c348eaa7f01d489b111efae Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 13:44:28 +0000 Subject: [PATCH 55/77] Roto model loading: drop a load for an old choice, key viewer frames by the active model Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01RDvxCS8xcR28VArjzn1x7U --- crates/engine/src/analysis/models.rs | 49 +++++++++++++++++++++++-- crates/engine/src/tests/models.rs | 30 +++++++++++++++ crates/ui-egui/src/app/viewer_frames.rs | 12 +++++- 3 files changed, 87 insertions(+), 4 deletions(-) diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index ae15fe1d3..b66680507 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -36,6 +36,8 @@ pub struct Status { /// A finished job's message and whether it failed, for a toast (taken by /// [`Session::poll_models`]). pub done: Option<(String, bool)>, + /// Counts selection changes: a load started under an older count is stale and its model is discarded. + pub generation: u64, } /// The session's trained models: per task, the worker's status; and the face model in use. @@ -63,6 +65,12 @@ impl Models { self.face_model.lock().ok().and_then(|m| m.clone()) } + /// Activate a model that finished loading for the choice counted as `generation`; `false` (and the model is + /// dropped) when the choice changed meanwhile. + pub(crate) fn set_active_if_current(&self, task: Task, model: Loaded, generation: u64) -> bool { + activate_if_current(self.lane(task), &self.face_model, task, model, generation) + } + pub(crate) fn set_active(&self, task: Task, model: Option) { activate(self.lane(task), &self.face_model, task, model); } @@ -92,6 +100,30 @@ fn activate(lane: &Arc>, face: &FaceSlot, task: Task, model: Optio set(lane, |s| s.active = id); } +/// [`activate`] a freshly loaded model unless the choice changed since the load started (`generation`); whether it +/// was activated. The check and the activation happen under the lane's lock, so a selection can't slip between them. +fn activate_if_current(lane: &Arc>, face: &FaceSlot, task: Task, model: Loaded, generation: u64) -> bool { + let Ok(mut s) = lane.lock() else { return false }; + if s.generation != generation { + return false; + } + let id = match &model { + Loaded::Mask(m) => m.info().id.to_string(), + Loaded::Face(m) => m.info().id.to_string(), + }; + match (task, model) { + (Task::Mask, Loaded::Mask(m)) => crate::effects::roto::set_model(Some(m)), + (Task::Face, Loaded::Face(f)) => { + if let Ok(mut slot) = face.lock() { + *slot = Some(f); + } + } + _ => return false, + } + s.active = Some(id); + true +} + fn set(status: &Arc>, f: impl FnOnce(&mut Status)) { if let Ok(mut s) = status.lock() { f(&mut s); @@ -160,7 +192,11 @@ impl Session { Task::Mask => &mut self.prefs.roto.model, Task::Face => &mut self.prefs.face.model, }; - *slot = id.to_string(); + if slot != id { + *slot = id.to_string(); + // a load still running for the previous choice must not activate when it finishes + set(self.models.lane(task), |s| s.generation = s.generation.wrapping_add(1)); + } } /// Load (or drop) the models Settings chooses; called whenever settings change. @@ -190,9 +226,13 @@ impl Session { return; }; let (face_slot, label) = (self.models.face_model.clone(), kind(task).label); + let generation = st.generation; spawn(&lane, label, format!("Loading {}…", info.name), wait, move |status| { let m = load_file(&path, info.id)?; - activate(status, &face_slot, task, Some(m)); + // the choice changed while this loaded: drop the model, `poll_models` loads the new choice + if !activate_if_current(status, &face_slot, task, m, generation) { + return Ok(String::new()); + } Ok(format!("{label}: {} loaded", info.name)) }); } @@ -329,7 +369,10 @@ impl Session { for task in [Task::Mask, Task::Face] { let done = self.models.lane(task).lock().ok().and_then(|mut s| s.done.take()); if let Some((message, error)) = done { - self.events.push(crate::Event::Toast { message, error }); + // a discarded (stale) load has no message + if !message.is_empty() { + self.events.push(crate::Event::Toast { message, error }); + } self.apply_model(task, false); } } diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index 84bbd2381..567b51666 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -161,6 +161,36 @@ fn an_idle_roto_model_is_unloaded() { let _ = std::fs::remove_dir_all(d); } +/// A model that finishes loading after the choice changed is dropped, not activated. +#[test] +fn a_load_for_an_old_choice_does_not_activate() { + use effectcraft_segment::{Loaded, MaskModel, ModelInfo, Prompt}; + struct Fake; + impl MaskModel for Fake { + fn info(&self) -> &'static ModelInfo { + &effectcraft_segment::SAM2_TINY + } + fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &Prompt) -> effectcraft_segment::Result> { + Err("no weights".into()) + } + } + super::roto::hold_roto_cache_test_lock(); + let (mut s, d) = session("stale-load"); + let started = s.models.status(Task::Mask).generation; + // the user picks another engine while the load runs + s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); + s.execute_checked("roto.model.select", json!({"id": "sam2.1-hiera-tiny"})).unwrap(); + assert_ne!(s.models.status(Task::Mask).generation, started, "a changed choice is counted"); + let m: Arc = Arc::new(Fake); + assert!(!s.models.set_active_if_current(Task::Mask, Loaded::Mask(m.clone()), started), "the stale load is dropped"); + assert!(s.models.status(Task::Mask).active.is_none() && crate::effects::roto::installed_model().is_none()); + let now = s.models.status(Task::Mask).generation; + assert!(s.models.set_active_if_current(Task::Mask, Loaded::Mask(m), now), "a load for the current choice activates"); + assert_eq!(s.models.status(Task::Mask).active.as_deref(), Some("sam2.1-hiera-tiny")); + crate::effects::roto::set_model(None); + let _ = std::fs::remove_dir_all(d); +} + /// A stroke reloading an unloaded model never waits for the load and never retries one that already failed. #[cfg(not(target_arch = "wasm32"))] #[test] diff --git a/crates/ui-egui/src/app/viewer_frames.rs b/crates/ui-egui/src/app/viewer_frames.rs index c878fafa4..f87db615c 100644 --- a/crates/ui-egui/src/app/viewer_frames.rs +++ b/crates/ui-egui/src/app/viewer_frames.rs @@ -84,10 +84,20 @@ impl EffectcraftApp { frame: 0, scale: (scale * 1000.0).round() as u32, view: self.view_hash(comp), - opts: frames::opts_hash(&self.frame_opts(comp, scale)), + opts: frames::opts_hash(&self.frame_opts(comp, scale)) ^ self.roto_model_stamp(), } } + /// Changes when the Roto Brush model in use does: a frame drawn with the classic fallback while the model loaded + /// is not the frame the model gives, so the two must not share a cache entry. + fn roto_model_stamp(&self) -> u64 { + use std::hash::{Hash, Hasher}; + let Some(id) = self.session.models_roto_active() else { return 0 }; + let mut h = std::collections::hash_map::DefaultHasher::new(); + id.hash(&mut h); + h.finish() | 1 + } + /// The frame series the viewer of `comp` shows now (its scale, view and options, for the /// comp's current content): what the cache bars count. pub fn shown_series(&self, comp: ItemId) -> FrameKey { From 92cd5554401cd01901a72e7035617bfb2adf9db0 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 13:45:39 +0000 Subject: [PATCH 56/77] Roto model loading: keep the stale-load helper test-only Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01RDvxCS8xcR28VArjzn1x7U --- crates/engine/src/analysis/models.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index b66680507..9ff8a5ca7 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -67,6 +67,7 @@ impl Models { /// Activate a model that finished loading for the choice counted as `generation`; `false` (and the model is /// dropped) when the choice changed meanwhile. + #[cfg(test)] pub(crate) fn set_active_if_current(&self, task: Task, model: Loaded, generation: u64) -> bool { activate_if_current(self.lane(task), &self.face_model, task, model, generation) } From 36ed77820d11d4390b1a1bcb0a348bdd095fe8e7 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 16:09:57 +0000 Subject: [PATCH 57/77] Roto Brush: keep saved tracking when nothing is retracked, read saved object mattes under their own keys, let Freeze use complete saved tracking without its model Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01RDvxCS8xcR28VArjzn1x7U --- crates/effects/src/roto.rs | 79 ++++++++++++++++++++++++++---- crates/engine/src/analysis/roto.rs | 18 ++++++- crates/engine/src/tests/roto.rs | 39 +++++++++++++++ 3 files changed, 126 insertions(+), 10 deletions(-) diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 235c26afa..615ea8114 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -428,7 +428,13 @@ impl Chain { /// has several objects and the tracker followed them there; computes nothing. pub fn object_seg(&self, f: i64, o: u32) -> Option> { let sc = self.opts.scale; - let k = *self.keys.get(&f)?; + // A frame the saved tracking has is read from it, under the keys it was made with: the model chosen now + // may not be the one that tracked it, and its own cache must not hold another model's matte. + let saved = self.saved.as_ref().filter(|t| t.cache.frames.contains_key(&f)).and_then(|t| t.keys.get(&f).map(|k| (t, *k))); + let k = match saved { + Some((_, saved_key)) => saved_key, + None => *self.keys.get(&f)?, + }; let ok = object_key(k, o); if let Some(s) = cached(ok, sc) { return Some(s); @@ -436,7 +442,7 @@ impl Chain { let full = match cached(ok, 1.0) { Some(s) => s, None => { - let t = self.saved.as_ref()?; + let (t, _) = saved?; let s = Arc::new(decode_object(&t.cache, f, o)?); store(ok, 1.0, s.clone()); s @@ -596,6 +602,11 @@ impl Chain { || self.saved.as_ref().is_some_and(|t| t.cache.frames.contains_key(&f) && t.keys.contains_key(&f)) } + /// Whether the tracking saved with the project has frame `f` (and applies to this chain). + pub fn saved_has(&self, f: i64) -> bool { + self.saved.as_ref().is_some_and(|t| t.cache.frames.contains_key(&f) && t.keys.contains_key(&f)) + } + /// Whether the chosen model can follow an object through a clip (SAM 2). pub fn has_tracker(&self) -> bool { self.model.as_ref().is_some_and(|m| m.tracker().is_some()) @@ -1032,15 +1043,21 @@ fn decode_object(cache: &TrackCache, f: i64, o: u32) -> Option { /// tracker followed, at full resolution, with its trust. `None` when it followed none. pub fn tracks_of(ch: &Chain) -> Option { let model = ch.model.as_ref()?.info().id; - let mut size = None; + // The saved tracking applies to this chain's frames when the same model made it: what it has and this session did + // not track again stays (a project reopened and propagated with nothing new to track must keep its tracking). + let saved = ch.saved.as_ref().filter(|t| t.cache.model == model); + let mut size = saved.map(|t| t.cache.size); let mut frames = BTreeMap::new(); for (f, k) in &ch.keys { - let (Some(seg), Some(t)) = (cached(tracked_key(*k), 1.0), trust(*k)) else { continue }; - size.get_or_insert([seg.w, seg.h]); - frames.insert( - *f, - TrackFrame { m: rle::encode(&seg.matte), r: rle::encode(&seg.refine), fg: seg.fg.clone(), bg: seg.bg.clone(), c: t.confidence, o: t.occluded }, - ); + if let (Some(seg), Some(t)) = (cached(tracked_key(*k), 1.0), trust(*k)) { + size.get_or_insert([seg.w, seg.h]); + frames.insert( + *f, + TrackFrame { m: rle::encode(&seg.matte), r: rle::encode(&seg.refine), fg: seg.fg.clone(), bg: seg.bg.clone(), c: t.confidence, o: t.occluded }, + ); + } else if let Some(kept) = saved.filter(|t| t.keys.get(f) == Some(k)).and_then(|t| t.cache.frames.get(f)) { + frames.insert(*f, kept.clone()); + } } let mut objects: BTreeMap> = BTreeMap::new(); let ids = ch.data.objects(); @@ -1052,6 +1069,8 @@ pub fn tracks_of(ch: &Chain) -> Option { for o in &ids { if let Some(seg) = cached(object_key(*k, *o), 1.0) { objects.entry(*o).or_default().insert(*f, ObjectFrame { m: rle::encode(&seg.matte), r: rle::encode(&seg.refine) }); + } else if let Some(kept) = saved.filter(|t| t.keys.get(f) == Some(k)).and_then(|t| t.cache.objects.get(o)).and_then(|m| m.get(f)) { + objects.entry(*o).or_default().insert(*f, kept.clone()); } } } @@ -1060,6 +1079,12 @@ pub fn tracks_of(ch: &Chain) -> Option { (!cache.frames.is_empty()).then_some(cache) } +/// Whether the tracking saved with `params` covers every one of `frames`: they need no model to be read. +pub fn saved_covers(params: &Params, layer_size: [f64; 2], fps: f64, frames: &[i64]) -> bool { + let chain = Chain::new(params, layer_size, fps, 1.0); + !frames.is_empty() && frames.iter().all(|f| chain.saved_has(*f)) +} + /// Encode one frame for [`FrozenCache`]. pub fn freeze_frame(alpha: &[f32], band: &[u8]) -> FrozenFrame { let a: Vec = alpha.iter().map(|v| (v.clamp(0.0, 1.0) * 255.0).round() as u8).collect(); @@ -1402,6 +1427,42 @@ mod tests { assert_eq!(completed, 0, "partial objects are not complete tracked frames"); } + /// A cold chain with saved tracking and the same model keeps the saved frames when asked what to save (nothing + /// was tracked this session), and Show Object reads a frame the saved tracking has under its own keys. + #[test] + fn saved_tracking_is_kept_and_its_objects_stay_out_of_another_models_keys() { + let mut d = RotoData::default(); + d.add(Stroke { kind: StrokeKind::Fg, frame: 0, radius: 3.0, points: vec![[35.0, 40.0]], object: 0 }, 2, [0, 24]); + let mut params = params_with(&d); + params.values.insert(INPUT_KEY.into(), Value::Str("keep-saved-and-objects".into())); + let size = [100.0, 80.0]; + let model = effectcraft_segment::SAM2_TINY.id; + let keys = d.chain_keys(seed_with(¶ms, size, 25.0, model)); + let (matte, empty) = (vec![1u8; 100 * 80], vec![0u8; 100 * 80]); + let frame = TrackFrame { m: rle::encode(&matte), r: rle::encode(&empty), c: 0.9, ..Default::default() }; + let object = ObjectFrame { m: rle::encode(&matte), r: rle::encode(&empty) }; + let saved = TrackCache { + key: tracks_key(&keys), + model: model.into(), + size: [100, 80], + frames: keys.keys().map(|f| (*f, frame.clone())).collect(), + objects: [(1, keys.keys().map(|f| (*f, object.clone())).collect())].into_iter().collect(), + }; + params.values.insert(TRACKED.into(), Value::Str(saved.to_json())); + let mut chain = Chain::new(¶ms, size, 25.0, 1.0); + chain.model = Some(Arc::new(TwoObjects(Default::default()))); + chain.keys = keys.clone(); + let kept = tracks_of(&chain).expect("the saved tracking is kept"); + assert_eq!(kept.frames.keys().collect::>(), saved.frames.keys().collect::>()); + assert_eq!(kept.frames.get(&1), saved.frames.get(&1)); + // another model is chosen now: its keys differ, and the saved object matte is not stored under them + let other = d.chain_keys(seed_with(¶ms, size, 25.0, effectcraft_segment::SAM2_SMALL.id)); + chain.keys = other.clone(); + let seg = chain.object_seg(1, 1).expect("the saved object matte"); + assert_eq!(seg.matte.len(), 100 * 80); + assert!(cached(object_key(other[&1], 1), 1.0).is_none(), "the other model's cache holds none of it"); + } + #[test] fn a_one_direction_pass_tracks_and_reports_only_its_frames() { let mut d = RotoData::default(); diff --git a/crates/engine/src/analysis/roto.rs b/crates/engine/src/analysis/roto.rs index a07b0010e..03b287d6e 100644 --- a/crates/engine/src/analysis/roto.rs +++ b/crates/engine/src/analysis/roto.rs @@ -642,6 +642,21 @@ impl Session { } } + /// Whether the tracking saved with Roto Brush `effect` covers every frame a Freeze would write. + fn freeze_is_saved(&self, comp: ItemId, layer: LayerId, effect: Uid) -> bool { + let Some(c) = self.project.comp(comp) else { return false }; + let Some(l) = c.layer(layer) else { return false }; + let Some(g) = l.effects().and_then(|f| f.groups().find(|g| g.uid == effect)).filter(|g| is_roto(g)) else { return false }; + let pstat = params_static(g); + let d = fx::data(&pstat); + if d.is_empty() { + return false; + } + let t = comp_time_of(c, l, d.base.unwrap_or(0)); + let params = params_at(&self.project, comp, l, g, t, self.expr.as_deref()); + fx::saved_covers(¶ms, layer_size(&self.project, c, l), c.frame_rate.as_f64(), &job_frames(&d, Direction::Both)) + } + /// Whether `effect` on `layer` segments with the model Settings chose: Object Matte always, /// Roto Brush unless it is version 1.0 (the classic engine). pub(crate) fn uses_roto_model(&self, comp: ItemId, layer: LayerId, effect: Uid) -> bool { @@ -654,7 +669,8 @@ impl Session { pub fn start_roto(&mut self, comp: ItemId, layer: LayerId, effect: Uid, task: RotoTask, dir: Direction, wait: bool) -> Result { // The model must be in use before the job starts: otherwise it would run on the classic // engine without saying so. - if self.uses_roto_model(comp, layer, effect) { + // A Freeze whose frames all have saved tracking reads them without the model. + if self.uses_roto_model(comp, layer, effect) && !(task == RotoTask::Freeze && self.freeze_is_saved(comp, layer, effect)) { self.ensure_roto_model(wait)?; } if self.is_roto_running() { diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index 6c57d5c29..afc3adad9 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -840,6 +840,45 @@ fn saved_tracking_is_kept_in_the_project_and_applies_when_it_reopens() { let _ = std::fs::remove_dir_all(std::path::Path::new(&path).parent().unwrap()); } +/// A reopened project propagated again with the same model and nothing new to track keeps the tracking it has: the +/// saved frames are not replaced by an empty one, so a later reopen still finds them. +#[test] +fn propagating_with_nothing_new_keeps_the_saved_tracking() { + let (s, clip, uid, path) = saved_project("keep-saved"); + forget(&s, clip, uid); + let mut s2 = Session::default(); + s2.execute_checked("file.open", json!({"path": path})).unwrap(); + s2.footage = s.footage.clone(); + effectcraft_effects::roto::set_model(Some(std::sync::Arc::new(Exact))); + let result = s2.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})); + effectcraft_effects::roto::set_model(None); + result.unwrap(); + let st = s2.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); + assert_eq!(st["tracker"]["saved"], "sam2.1-hiera-tiny", "{st}"); + assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 21, "{st}"); + // and it is still in the file after saving and reopening + s2.execute_checked("file.save", json!({})).unwrap(); + let mut s3 = Session::default(); + s3.execute_checked("file.open", json!({"path": path})).unwrap(); + let st = s3.execute("roto.status", json!({"layer": clip.0, "frame": 12})).unwrap(); + assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 21, "{st}"); + let _ = std::fs::remove_dir_all(std::path::Path::new(&path).parent().unwrap()); +} + +/// Freeze reads saved tracking without its model: choosing a model that is not installed does not stop it when +/// the saved frames cover every frame it freezes. +#[test] +fn freeze_uses_complete_saved_tracking_without_the_model() { + let (s, clip, uid, path) = saved_project("freeze-saved"); + forget(&s, clip, uid); + let mut s2 = Session::default(); + s2.execute_checked("file.open", json!({"path": path})).unwrap(); + s2.footage = s.footage.clone(); + s2.execute_checked("roto.model.select", json!({"id": "sam2.1-hiera-tiny"})).unwrap(); + s2.execute_checked("roto.freeze", json!({"layer": clip.0, "wait": true})).unwrap(); + let _ = std::fs::remove_dir_all(std::path::Path::new(&path).parent().unwrap()); +} + /// Propagate replaces a fully saved span when a different tracking model is active. #[test] fn propagating_saved_tracking_uses_the_new_model() { From 31307d15dcbdd39d30351e7e8a27c730a317176f Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 29 Oct 2026 19:37:38 +0100 Subject: [PATCH 58/77] #1539: bound saved tracking decoders by validated plane dimensions --- crates/effects/src/roto.rs | 43 ++++++++++++++++++++++++++++++++++---- 1 file changed, 39 insertions(+), 4 deletions(-) diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 615ea8114..649060be4 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -1024,8 +1024,11 @@ fn decode_tracked(cache: &TrackCache, f: i64) -> Option { // The size comes from the project file: checked, so a corrupt one is ignored rather than overflowing. let [w, h] = cache.size; let n = w.checked_mul(h)?; - let matte = rle::decode(&fr.m)?; - let refine = rle::decode(&fr.r)?; + if n == 0 || n > rle::MAX_PLANE_BYTES { + return None; + } + let matte = rle::decode_bounded(&fr.m, n)?; + let refine = rle::decode_bounded(&fr.r, n)?; (matte.len() == n && refine.len() == n).then(|| FrameSeg { w, h, matte, refine, fg: fr.fg.clone(), bg: fr.bg.clone() }) } @@ -1034,8 +1037,11 @@ fn decode_object(cache: &TrackCache, f: i64, o: u32) -> Option { let fr = cache.objects.get(&o)?.get(&f)?; let [w, h] = cache.size; let n = w.checked_mul(h)?; - let matte = rle::decode(&fr.m)?; - let refine = rle::decode(&fr.r)?; + if n == 0 || n > rle::MAX_PLANE_BYTES { + return None; + } + let matte = rle::decode_bounded(&fr.m, n)?; + let refine = rle::decode_bounded(&fr.r, n)?; (matte.len() == n && refine.len() == n).then(|| FrameSeg { w, h, matte, refine, ..Default::default() }) } @@ -1463,6 +1469,35 @@ mod tests { assert!(cached(object_key(other[&1], 1), 1.0).is_none(), "the other model's cache holds none of it"); } + #[test] + fn saved_tracking_decoders_validate_dimensions_and_plane_lengths() { + let valid = rle::encode(&[255, 0, 255, 0]); + let mut cache = TrackCache { + size: [2, 2], + frames: [(0, TrackFrame { m: valid.clone(), r: valid.clone(), ..Default::default() })].into_iter().collect(), + objects: [(1, [(0, ObjectFrame { m: valid.clone(), r: valid.clone() })].into_iter().collect())].into_iter().collect(), + ..Default::default() + }; + assert_eq!(decode_tracked(&cache, 0).unwrap().matte, [255, 0, 255, 0]); + assert_eq!(decode_object(&cache, 0, 1).unwrap().refine, [255, 0, 255, 0]); + for (size, matte, refine) in [ + ([0, 2], rle::encode(&[]), rle::encode(&[])), + ([2, 0], rle::encode(&[]), rle::encode(&[])), + ([usize::MAX, 2], valid.clone(), valid.clone()), + ([rle::MAX_PLANE_BYTES + 1, 1], valid.clone(), valid.clone()), + ([2, 2], rle::encode(&[0; 5]), valid.clone()), + ([2, 2], valid.clone(), rle::encode(&[0; 5])), + ([2, 2], rle::encode(&[0; 3]), valid.clone()), + ([2, 2], valid, rle::encode(&[0; 3])), + ] { + cache.size = size; + cache.frames.insert(0, TrackFrame { m: matte.clone(), r: refine.clone(), ..Default::default() }); + cache.objects.insert(1, [(0, ObjectFrame { m: matte, r: refine })].into_iter().collect()); + assert!(decode_tracked(&cache, 0).is_none(), "tracked plane {size:?}"); + assert!(decode_object(&cache, 0, 1).is_none(), "object plane {size:?}"); + } + } + #[test] fn a_one_direction_pass_tracks_and_reports_only_its_frames() { let mut d = RotoData::default(); From b148397d4e9db1c38975979600dc422b4542a7ea Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 29 Oct 2026 19:55:42 +0100 Subject: [PATCH 59/77] #1539: keep Show Object empty when its tracked matte is missing --- crates/effects/src/roto.rs | 62 +++++++++++++++++++++++++++++++++----- 1 file changed, 55 insertions(+), 7 deletions(-) diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 649060be4..8046c4773 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -457,20 +457,24 @@ impl Chain { } /// What Show Object asks for at frame `f`: `None` shows the whole matte (all objects, a single object, or a frame - /// the tracker did not follow); an object the instance does not have shows nothing. + /// the tracker did not follow); a missing object on a tracked frame shows nothing. fn shown_object(&self, f: i64) -> Option> { let o = self.show_object.checked_sub(1)?; let objects = self.data.objects(); if objects.len() < 2 { return None; } - if !objects.contains(&o) { - let sc = self.opts.scale; - let [w, h] = - self.saved.as_ref().map(|t| t.cache.size).or_else(|| self.keys.get(&f).and_then(|k| cached(tracked_key(*k), sc)).map(|s| [s.w, s.h]))?; - return Some(Arc::new(FrameSeg::empty(w, h))); + if objects.contains(&o) + && let Some(seg) = self.object_seg(f, o) + { + return Some(seg); } - self.object_seg(f, o) + let sc = self.opts.scale; + let tracked = self.saved_seg(f, sc).or_else(|| { + let key = tracked_key(*self.keys.get(&f)?); + cached(key, sc).or_else(|| cached(key, 1.0).map(|seg| Arc::new(downscale(&seg, sc)))) + })?; + Some(Arc::new(FrameSeg::empty(tracked.w, tracked.h))) } /// The tracker's trust in frame `f`: from the saved tracking, else from this session's tracking. @@ -1469,6 +1473,50 @@ mod tests { assert!(cached(object_key(other[&1], 1), 1.0).is_none(), "the other model's cache holds none of it"); } + #[test] + fn showing_a_missing_tracked_object_never_falls_back_to_other_objects() { + let mut d = RotoData::default(); + for object in 0..2 { + d.add(Stroke { kind: StrokeKind::Fg, frame: 0, radius: 1.0, points: vec![[1.0, 1.0]], object }, 1, [0, 1]); + } + let union = Arc::new(FrameSeg { w: 4, h: 2, matte: vec![255; 8], refine: vec![0; 8], ..Default::default() }); + for saved in [false, true] { + for scale in [1.0, 0.5] { + let mut params = params_with(&d); + params.values.insert("version".into(), Value::Scalar(0.0)); + params.values.insert(INPUT_KEY.into(), Value::Str(format!("missing-object-{saved}-{scale}"))); + let keys = d.chain_keys(seed(¶ms, [4.0, 2.0], 25.0)); + if saved { + let matte = rle::encode(&union.matte); + let refine = rle::encode(&union.refine); + let cache = TrackCache { + key: tracks_key(&keys), + model: effectcraft_segment::CLASSICAL.into(), + size: [4, 2], + frames: [(0, TrackFrame { m: matte.clone(), r: refine.clone(), ..Default::default() })].into_iter().collect(), + objects: [(0, [(0, ObjectFrame { m: matte, r: refine })].into_iter().collect())].into_iter().collect(), + }; + params.values.insert(TRACKED.into(), Value::Str(cache.to_json())); + } else { + store(tracked_key(keys[&0]), 1.0, union.clone()); + store(object_key(keys[&0], 0), 1.0, union.clone()); + } + let mut chain = Chain::new(¶ms, [4.0, 2.0], 25.0, scale); + let expected = downscale(&union, scale); + for object in [2, 7] { + chain.show_object = object; + let shown = chain.ensure(0, &mut |_| None, &|| false).expect("tracked frame"); + assert_eq!((shown.w, shown.h, shown.area()), (expected.w, expected.h, 0), "saved={saved}, scale={scale}, object={object}"); + } + assert!(chain.shown_object(1).is_none(), "untracked frames keep their existing fallback"); + for object in [0, 1] { + chain.show_object = object; + assert_eq!(chain.ensure(0, &mut |_| None, &|| false).unwrap().matte, expected.matte); + } + } + } + } + #[test] fn saved_tracking_decoders_validate_dimensions_and_plane_lengths() { let valid = rle::encode(&[255, 0, 255, 0]); From 365abf0f2d39c8722d3347c81f22bb87c54e8da9 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 29 Oct 2026 20:59:50 +0100 Subject: [PATCH 60/77] Roto Brush: discard failed tracks and bound SAM2 caches (#1539) --- apps/effectcraft-cli/tests/it/cli.rs | 11 ++-- crates/effects/src/roto.rs | 32 +++++++++++ crates/engine/src/analysis/roto.rs | 3 ++ crates/engine/src/tests/roto.rs | 48 ++++++++++++----- crates/sam2/src/mask.rs | 11 ++-- crates/sam2/src/neck.rs | 81 +++++++++++++++++++++++----- crates/sam2/src/resize.rs | 23 +++++--- crates/sam2/src/track.rs | 30 +++++++++-- docs/architecture.md | 3 +- 9 files changed, 194 insertions(+), 48 deletions(-) diff --git a/apps/effectcraft-cli/tests/it/cli.rs b/apps/effectcraft-cli/tests/it/cli.rs index ed6d98086..8bdfb3c00 100644 --- a/apps/effectcraft-cli/tests/it/cli.rs +++ b/apps/effectcraft-cli/tests/it/cli.rs @@ -683,6 +683,7 @@ fn render_bitrate_applies_to_vp9_webm() { #[test] fn bench_roto_runs_on_a_users_clip() { let (frames, masks) = (tmp("clip-frames"), tmp("clip-masks")); + let config = frames.join("config"); for d in [&frames, &masks] { let _ = std::fs::remove_dir_all(d); std::fs::create_dir_all(d).unwrap(); @@ -703,7 +704,7 @@ fn bench_roto_runs_on_a_users_clip() { mask.save(masks.join(format!("f{f:03}.png"))).unwrap(); } let (f, m) = (frames.to_str().unwrap(), masks.to_str().unwrap()); - let out = bin().args(["bench", "--roto", "--clip", f, "--truth", m, "--frames", "6"]).output().unwrap(); + let out = bin().env("EFFECTCRAFT_CONFIG_DIR", &config).args(["bench", "--roto", "--clip", f, "--truth", m, "--frames", "6"]).output().unwrap(); let text = String::from_utf8_lossy(&out.stdout).to_string(); assert!(out.status.success(), "{text}\n{}", String::from_utf8_lossy(&out.stderr)); let row = text.lines().find(|l| l.contains("classical")).unwrap_or_else(|| panic!("no classical row: {text}")); @@ -711,11 +712,11 @@ fn bench_roto_runs_on_a_users_clip() { assert!(mean > 0.8, "classical mean IoU on the user's clip: {mean}\n{text}"); // one mask short: refused std::fs::remove_file(masks.join("f005.png")).unwrap(); - let out = bin().args(["bench", "--roto", "--clip", f, "--truth", m]).output().unwrap(); + let out = bin().env("EFFECTCRAFT_CONFIG_DIR", &config).args(["bench", "--roto", "--clip", f, "--truth", m]).output().unwrap(); assert!(!out.status.success()); assert!(String::from_utf8_lossy(&out.stderr).contains("must match"), "{}", String::from_utf8_lossy(&out.stderr)); // --clip without --truth - let out = bin().args(["bench", "--roto", "--clip", f]).output().unwrap(); + let out = bin().env("EFFECTCRAFT_CONFIG_DIR", &config).args(["bench", "--roto", "--clip", f]).output().unwrap(); assert!(!out.status.success()); } @@ -730,11 +731,11 @@ fn bench_roto_reports_why_a_named_model_file_failed() { std::fs::write(&bogus, b"not the weights").unwrap(); for flag in ["--sam2", "--mobilesam"] { let missing = dir.join("missing.bin"); - let out = bin().args(["bench", "--roto", flag, missing.to_str().unwrap()]).output().unwrap(); + let out = bin().env("EFFECTCRAFT_CONFIG_DIR", dir.join("config")).args(["bench", "--roto", flag, missing.to_str().unwrap()]).output().unwrap(); let err = String::from_utf8_lossy(&out.stderr).to_string(); assert!(!out.status.success(), "{flag}: {err}"); assert!(err.contains(flag) && err.contains("missing.bin"), "{flag} names the file that could not be read: {err}"); - let out = bin().args(["bench", "--roto", flag, bogus.to_str().unwrap()]).output().unwrap(); + let out = bin().env("EFFECTCRAFT_CONFIG_DIR", dir.join("config")).args(["bench", "--roto", flag, bogus.to_str().unwrap()]).output().unwrap(); let err = String::from_utf8_lossy(&out.stderr).to_string(); assert!(!out.status.success(), "{flag}: {err}"); assert!(err.contains("does not match"), "{flag} says the file is not the published weights: {err}"); diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 8046c4773..e66400416 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -624,6 +624,20 @@ impl Chain { } } + /// Discard a failed track completely before falling back to the classic engine. + pub fn forget_tracking(&mut self) { + let mut objects = self.data.objects(); + if let Some(saved) = &self.saved { + objects.extend(saved.cache.objects.keys().copied()); + objects.sort_unstable(); + objects.dedup(); + } + let keys = self.keys.values().chain(self.saved.iter().flat_map(|saved| saved.keys.values())).copied(); + forget(keys.flat_map(|key| std::iter::once(key).chain(objects.iter().copied().map(move |object| object_key(key, object))))); + self.saved = None; + self.model = None; + } + /// Follow the object through the span with the model's tracker, from the strokes on their /// frames (only the frames without a tracked result yet: after a correction stroke, the frames it /// changes, tracked on from the nearest prompted frame), and cache every frame's segmentation (under the tracker's keys, with its @@ -1471,6 +1485,24 @@ mod tests { let seg = chain.object_seg(1, 1).expect("the saved object matte"); assert_eq!(seg.matte.len(), 100 * 80); assert!(cached(object_key(other[&1], 1), 1.0).is_none(), "the other model's cache holds none of it"); + for key in [keys[&1], other[&1]] { + for scale in [1.0, 0.5] { + store(tracked_key(key), scale, seg.clone()); + store(object_key(key, 1), scale, seg.clone()); + } + trust_cache().lock().unwrap().insert(tracked_key(key), FrameTrust { confidence: 0.9, occluded: false }); + } + chain.forget_tracking(); + assert!(chain.model.is_none()); + assert!(!chain.saved_has(1)); + for key in [keys[&1], other[&1]] { + for scale in [1.0, 0.5] { + assert!(cached(tracked_key(key), scale).is_none()); + assert!(cached(object_key(key, 1), scale).is_none()); + } + assert!(trust(key).is_none()); + } + assert!(chain.ensure(1, &mut |_| None, &|| false).is_none(), "neither saved nor in-memory tracking survives fallback"); } #[test] diff --git a/crates/engine/src/analysis/roto.rs b/crates/engine/src/analysis/roto.rs index 03b287d6e..67b0afc68 100644 --- a/crates/engine/src/analysis/roto.rs +++ b/crates/engine/src/analysis/roto.rs @@ -477,6 +477,9 @@ async fn run_work(w: Work, shared: &RotoShared) { Ok(_) => fx::tracks_of(&chain).map(|t| t.to_json()).unwrap_or_default(), Err(_) => String::new(), }); + if result.is_err() { + chain.forget_tracking(); + } let mut s = lock(&shared.state); match result { Ok(summary) => s.tracked = summary, diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index afc3adad9..e686962e1 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -937,9 +937,9 @@ fn a_stroke_change_makes_the_saved_tracking_stale() { let _ = std::fs::remove_dir_all(std::path::Path::new(&path).parent().unwrap()); } -/// A model whose tracker always fails: Roto Brush must carry on with the classic engine and say so. +/// A tracker that emits a completed frame before failing must leave only classic results behind. #[test] -fn a_failing_tracker_falls_back_to_the_classic_engine() { +fn a_partially_failing_tracker_falls_back_to_the_classic_engine() { use effectcraft_segment::{ClipFrame, MaskModel, ModelInfo, ObjectTracker, Prompt, TrackRequest, TrackedFrame}; use std::sync::Arc; struct Broken; @@ -957,29 +957,51 @@ fn a_failing_tracker_falls_back_to_the_classic_engine() { impl ObjectTracker for Broken { fn track( &self, - _: &TrackRequest, - _: &mut dyn FnMut(i64) -> Option>, + req: &TrackRequest, + frame: &mut dyn FnMut(i64) -> Option>, _: &dyn Fn() -> bool, - _: &mut dyn FnMut(usize, i64, TrackedFrame), + out: &mut dyn FnMut(usize, i64, TrackedFrame), ) -> effectcraft_segment::Result<()> { - Err("the tracker broke".into()) + let img = frame(req.first).ok_or("missing test frame")?; + for object in 0..req.objects.len() { + out(object, req.first, TrackedFrame { prob: vec![1.0; img.w * img.h], w: img.w, h: img.h, confidence: 0.9, occluded: false }); + } + Err("the tracker broke after emitting a frame".into()) } } hold_roto_cache_test_lock(); - let (mut s, clip, _) = painted(center, "disk-broken-tracker.mov"); + let (mut s, clip, uid) = painted(center, "disk-partially-broken-tracker.mov"); + s.execute_checked("roto.stroke", json!({"layer": clip.0, "object": 1, "points": [center(0)], "radius": 4})).unwrap(); effectcraft_effects::roto::set_model(Some(Arc::new(Broken))); let r = s.execute_checked("roto.propagate", json!({"layer": clip.0, "direction": "both", "wait": true})); + let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 0})); + let comp = s.active_comp().unwrap(); + let layer = comp.layer(clip).unwrap(); + let params = crate::roto::params_static(layer.props.find_group(uid).unwrap()); + let size = crate::roto::layer_size(&s.project, comp, layer); + let mut chain = effectcraft_effects::roto::Chain::new(¶ms, size, comp.frame_rate.as_f64(), 1.0); + let index = layer.effects().unwrap().groups().position(|g| g.uid == uid).unwrap(); + let mut source = + crate::roto::FrameSource::new(&s.project, s.footage.as_ref(), s.expr.as_deref(), &s.layer_cache, s.active_comp_id().unwrap(), comp, layer, index, size); + let mut get = |f: i64| source.get(f); + let actual: Option> = (0..=20).map(|f| chain.ensure(f, &mut get, &|| false)).collect(); + let retained_object = chain.object_seg(0, 0).is_some() || chain.object_seg(0, 1).is_some(); + let retained_trust = chain.trust_at(0); + // Build the reference only after capturing the fallback, with no model or cached union. + effectcraft_effects::roto::forget(chain.keys.values().copied()); + chain.model = None; + let expected: Option> = (0..=20).map(|f| chain.ensure(f, &mut get, &|| false)).collect(); effectcraft_effects::roto::set_model(None); let r = r.unwrap(); assert_eq!(r["frames"], 21, "{r}"); - let st = s.execute("roto.status", json!({"layer": clip.0, "frame": 10})).unwrap(); - assert_eq!(st["source"], "classical", "{st}"); + let st = st.unwrap(); + assert_ne!(st["source"], "tracker", "{st}"); assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 0); - let mut worst: f64 = 1.0; - for f in 0..=20 { - worst = worst.min(iou_at(&mut s, clip, f, center(f))); + assert!(!retained_object, "partial per-object mattes must be discarded"); + assert_eq!(retained_trust, None, "partial tracker confidence must be discarded"); + for (frame, (actual, expected)) in actual.unwrap().into_iter().zip(expected.unwrap()).enumerate() { + assert!(actual == expected, "fallback frame {frame} differs from the classic engine"); } - assert!(worst > 0.9, "the classic engine should have carried the clip: {worst}"); } /// Without a tracking model the Roto Brush chain is exactly the classic engine: every frame diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index 911ff2e00..382f9b514 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -19,10 +19,7 @@ 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()))); - }; + let pixels = validate_frame_size(w, 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]; @@ -40,6 +37,12 @@ pub fn preprocess(rgb: &[[f32; 3]], w: usize, h: usize) -> Result { Ok(x.broadcast_sub(&mean)?.broadcast_div(&std)?) } +pub(crate) fn validate_frame_size(w: usize, h: usize, len: usize) -> Result { + w.checked_mul(h) + .filter(|n| *n > 0 && *n == len && w <= MAX_SIDE && h <= MAX_SIDE && *n <= MAX_PIXELS) + .ok_or_else(|| Error::Input(format!("a {w} x {h} frame needs {} pixels (at most {MAX_PIXELS}), got {len}", w.saturating_mul(h)))) +} + 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 { diff --git a/crates/sam2/src/neck.rs b/crates/sam2/src/neck.rs index c897e9818..3a80ede9e 100644 --- a/crates/sam2/src/neck.rs +++ b/crates/sam2/src/neck.rs @@ -13,23 +13,47 @@ use crate::weights::Weights; use crate::{Error, Result}; /// Sine position encoding `[1, 2 * num_pos_features, h, w]`, normalised (`Sam2SinePositionEmbedding`). -/// It depends only on the sizes, so each is built once (it was 150 ms of a 2.7 s frame). +/// Only small tables are retained globally; the neck owns its larger feature-pyramid tables. pub fn sine_position(num_pos_features: usize, h: usize, w: usize) -> Result { - static CACHE: OnceLock>> = OnceLock::new(); - let cache = CACHE.get_or_init(Default::default); - if let Some(t) = cache.lock().unwrap_or_else(PoisonError::into_inner).get(&(num_pos_features, h, w)) { - return Ok(t.clone()); + static CACHE: OnceLock = OnceLock::new(); + CACHE.get_or_init(|| PositionCache::new(1024 * 1024 / size_of::())).get(num_pos_features, h, w) +} + +const MAX_POSITION_ELEMENTS: usize = 128 * 1024 * 1024 / size_of::(); + +struct PositionCache { + entries: Mutex>, + max_elements: usize, +} + +impl PositionCache { + fn new(max_elements: usize) -> Self { + Self { entries: Mutex::new(HashMap::new()), max_elements } } - let t = build_sine_position(num_pos_features, h, w)?; - let mut map = cache.lock().unwrap_or_else(PoisonError::into_inner); - if map.len() >= 16 { - map.clear(); + + fn get(&self, features: usize, h: usize, w: usize) -> Result { + let key = (features, h, w); + if let Some(t) = self.entries.lock().unwrap_or_else(PoisonError::into_inner).get(&key) { + return Ok(t.clone()); + } + let t = build_sine_position(features, h, w)?; + let count = t.elem_count(); + if count <= self.max_elements { + let mut map = self.entries.lock().unwrap_or_else(PoisonError::into_inner); + let retained = map.values().fold(0usize, |n, t| n.saturating_add(t.elem_count())); + if map.len() >= 16 || retained.saturating_add(count) > self.max_elements { + map.clear(); + } + map.insert(key, t.clone()); + } + Ok(t) } - map.insert((num_pos_features, h, w), t.clone()); - Ok(t) } fn build_sine_position(num_pos_features: usize, h: usize, w: usize) -> Result { + let count = num_pos_features.checked_mul(2).and_then(|n| n.checked_mul(h)).and_then(|n| n.checked_mul(w)); + let count = + count.filter(|n| *n > 0 && *n <= MAX_POSITION_ELEMENTS).ok_or_else(|| Error::Input("sine position dimensions are empty or exceed 128 MiB".into()))?; let eps = 1e-6f32; let scale = 2.0 * std::f32::consts::PI; let dim_t: Vec = (0..num_pos_features).map(|i| 10000f32.powf(2.0 * (i / 2) as f32 / num_pos_features as f32)).collect(); @@ -39,7 +63,7 @@ fn build_sine_position(num_pos_features: usize, h: usize, w: usize) -> Result, Vec<_>) = ((0..h).map(|k| axis(h, k)).collect(), (0..w).map(|k| axis(w, k)).collect()); - let mut out = vec![0f32; 2 * num_pos_features * h * w]; + let mut out = vec![0f32; count]; for y in 0..h { for x in 0..w { for f in 0..num_pos_features { @@ -75,6 +99,7 @@ pub struct Neck { convs: Vec, top_down: Vec, pos_features: usize, + positions: PositionCache, } /// What the neck produces, both in the neck's own order (lowest resolution first, all four @@ -97,7 +122,7 @@ impl Neck { }) }) .collect::>>()?; - Ok(Neck { convs, top_down: top_down.to_vec(), pos_features: fpn / 2 }) + Ok(Neck { convs, top_down: top_down.to_vec(), pos_features: fpn / 2, positions: PositionCache::new(MAX_POSITION_ELEMENTS) }) } /// `stages`: Hiera stage outputs, NHWC, highest resolution first. @@ -119,7 +144,7 @@ impl Neck { } _ => lateral, }; - positions.push(sine_position(self.pos_features, cur.dim(2)?, cur.dim(3)?)?); + positions.push(self.positions.get(self.pos_features, cur.dim(2)?, cur.dim(3)?)?); features.push(cur.clone()); prev = Some(cur); } @@ -160,3 +185,31 @@ impl VisionEncoder { Ok(ImageFeatures { features: keep(out.features), positions: keep(out.positions) }) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn position_tables_reject_empty_overflowing_and_oversized_dimensions() { + for (features, h, w) in [(0, 1, 1), (1, 0, 1), (1, 1, 0), (usize::MAX, 2, 2), (128, usize::MAX, 2), (128, 1024, 1024)] { + assert!(matches!(sine_position(features, h, w), Err(Error::Input(_))), "{features}, {h}, {w}"); + } + } + + #[test] + fn position_cache_bounds_retained_elements_and_preserves_the_encoding() { + let cache = PositionCache::new(64); + let first = cache.get(2, 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(); + 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(); + 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/resize.rs b/crates/sam2/src/resize.rs index 67e96e9af..fdbcbab48 100644 --- a/crates/sam2/src/resize.rs +++ b/crates/sam2/src/resize.rs @@ -5,8 +5,9 @@ use candle_core::{Device, Tensor}; +use crate::Result; +use crate::mask::validate_frame_size; use crate::prompt::INPUT_SIZE; -use crate::{Error, Result}; /// For each output index: the first input index and the normalised triangle-filter weights. fn axis_weights(n_in: usize, n_out: usize) -> Vec<(usize, Vec)> { @@ -35,13 +36,13 @@ fn axis_weights(n_in: usize, n_out: usize) -> Vec<(usize, Vec)> { /// Antialiased bilinear resize of a row-major `h x w` image to `oh x ow`. pub fn antialias_bilinear(src: &[f32], (h, w): (usize, usize), (oh, ow): (usize, usize)) -> Result> { - if h == 0 || w == 0 || oh == 0 || ow == 0 || src.len() != h * w { - return Err(Error::Input(format!("cannot resize {} values as {w}x{h} to {ow}x{oh}", src.len()))); - } + validate_frame_size(w, h, src.len())?; + let out_len = validate_frame_size(ow, oh, oh.saturating_mul(ow))?; + let mid_len = validate_frame_size(ow, h, h.saturating_mul(ow))?; let xs = axis_weights(w, ow); let ys = axis_weights(h, oh); // horizontal pass: h x ow - let mut mid = vec![0f32; h * ow]; + let mut mid = vec![0f32; mid_len]; for y in 0..h { let row = src.get(y * w..(y + 1) * w).unwrap_or(&[]); for (x, (x0, wt)) in xs.iter().enumerate() { @@ -52,7 +53,7 @@ pub fn antialias_bilinear(src: &[f32], (h, w): (usize, usize), (oh, ow): (usize, } } // vertical pass: oh x ow - let mut out = vec![0f32; oh * ow]; + let mut out = vec![0f32; out_len]; for (y, (y0, wt)) in ys.iter().enumerate() { for x in 0..ow { let acc: f32 = wt.iter().enumerate().map(|(k, a)| a * mid.get((y0 + k) * ow + x).copied().unwrap_or(0.0)).sum(); @@ -67,6 +68,7 @@ pub fn antialias_bilinear(src: &[f32], (h, w): (usize, usize), (oh, ow): (usize, /// A binary mask of any size as the model's mask prompt: `[1, 1, 1024, 1024]` of 0 and 1 (resized /// like the reference: antialiased bilinear, then thresholded at one half). pub fn mask_to_model(mask: &[u8], w: usize, h: usize) -> Result { + validate_frame_size(w, h, mask.len())?; let side = INPUT_SIZE; let f: Vec = mask.iter().map(|v| f32::from(*v != 0)).collect(); let data = if (w, h) == (side, side) { f } else { antialias_bilinear(&f, (h, w), (side, side))?.into_iter().map(|v| f32::from(v >= 0.5)).collect() }; @@ -77,6 +79,15 @@ pub fn mask_to_model(mask: &[u8], w: usize, h: usize) -> Result { mod tests { use super::*; + #[test] + fn resize_rejects_overflowing_oversized_and_mismatched_grids() { + for (input, output) in [((usize::MAX, 2), (1, 1)), ((1, 1), (usize::MAX, 2)), ((1, 1), (16_385, 1)), ((0, 1), (1, 1))] { + assert!(antialias_bilinear(&[1.0], input, output).is_err(), "{input:?} to {output:?}"); + } + assert!(mask_to_model(&[1], usize::MAX, 2).is_err()); + assert!(mask_to_model(&[1], INPUT_SIZE, INPUT_SIZE).is_err()); + } + #[test] fn a_constant_stays_constant_and_weights_sum_to_one() { let out = antialias_bilinear(&[3.0; 64 * 64], (64, 64), (16, 16)).unwrap(); diff --git a/crates/sam2/src/track.rs b/crates/sam2/src/track.rs index ad920efe7..7e73a1c74 100644 --- a/crates/sam2/src/track.rs +++ b/crates/sam2/src/track.rs @@ -11,7 +11,7 @@ use std::sync::Arc; use candle_core::Tensor; use effectcraft_segment::{ClipFrame, ObjectTracker, TrackRequest, TrackedFrame}; -use crate::mask::{Sam2Mask, preprocess}; +use crate::mask::{Sam2Mask, preprocess, validate_frame_size}; use crate::prompt::{INPUT_SIZE, Prompt}; use crate::resize::mask_to_model; use crate::session::{Cancel, Direction, FrameMask, FrameSource, ObjectId, VideoTracker}; @@ -36,8 +36,12 @@ impl FrameSource for Source<'_> { } fn frame(&self, index: usize) -> Result { - let f = (self.fetch.borrow_mut())(self.first + index as i64) - .ok_or_else(|| Error::Input(format!("frame {} is not available", self.first + index as i64)))?; + let number = i64::try_from(index) + .ok() + .filter(|_| index < self.count) + .and_then(|offset| self.first.checked_add(offset)) + .ok_or_else(|| Error::Input(format!("frame index {index} is outside the clip")))?; + let f = (self.fetch.borrow_mut())(number).ok_or_else(|| Error::Input(format!("frame {number} is not available")))?; self.sizes.borrow_mut().insert(index, (f.w, f.h)); preprocess(&f.rgb, f.w, f.h) } @@ -59,6 +63,7 @@ impl Sam2Mask { if req.count == 0 || req.count > MAX_CLIP { return Err(Error::Input(format!("a tracking request needs 1 to {MAX_CLIP} frames, got {}", req.count))); } + req.first.checked_add((req.count - 1) as i64).ok_or_else(|| Error::Input("the clip's last frame overflows its frame number".into()))?; if req.objects.is_empty() || req.objects.len() > MAX_OBJECTS || req.objects.iter().any(|o| o.is_empty()) { return Err(Error::Input(format!("tracking needs 1 to {MAX_OBJECTS} objects, each with a prompt on at least one frame"))); } @@ -69,12 +74,14 @@ impl Sam2Mask { let mut prompt_frames = vec![]; for (oi, prompts) in req.objects.iter().enumerate() { for (frame, p) in prompts { - let idx = usize::try_from(frame - req.first) - .ok() + let idx = frame + .checked_sub(req.first) + .and_then(|offset| usize::try_from(offset).ok()) .filter(|i| *i < req.count) .ok_or_else(|| Error::Input(format!("prompt frame {frame} is outside the clip")))?; if !sizes.borrow().contains_key(&idx) { let f = (source.fetch.borrow_mut())(*frame).ok_or_else(|| Error::Input(format!("frame {frame} is not available")))?; + validate_frame_size(f.w, f.h, f.rgb.len())?; sizes.borrow_mut().insert(idx, (f.w, f.h)); } prompt_frames.push((oi, idx, p)); @@ -159,3 +166,16 @@ impl ObjectTracker for Sam2Mask { r } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn source_rejects_overflowing_and_out_of_range_indices_before_fetching() { + let mut fetch = |_| panic!("invalid indices must not fetch a frame"); + let source = Source { first: i64::MAX, count: 2, fetch: RefCell::new(&mut fetch), sizes: Default::default() }; + assert!(matches!(source.frame(1), Err(Error::Input(_)))); + assert!(matches!(source.frame(usize::MAX), Err(Error::Input(_)))); + } +} diff --git a/docs/architecture.md b/docs/architecture.md index 0bfe4023e..c4bebd9ed 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -1,6 +1,6 @@ # Architecture -> **Last reviewed:** 2026-10-10 · **Last updated:** 2026-10-10 · **Change:** minor (crate table checked against `crates/`: effect count 306, VP9 inter frames) · **Target:** Adobe After Effects 2026 +> **Last reviewed:** 2026-10-10 · **Last updated:** 2026-10-29 · **Change:** minor (SAM 2 crate added to the layer table) · **Target:** Adobe After Effects 2026 EffectCraft is a stack of small Rust crates. The engine knows nothing about the user interface; the egui frontend, the command-line tool and the MCP server all sit on top of the same `Session` @@ -39,6 +39,7 @@ cpal or muda. Everything in L0 to L4, the egui UI and the web app also build for | L3 | `expr` | The expression engine (JavaScript via boa) with the After Effects object model | | L3 | `export` | Render queue encoding: H.264, HEVC and AV1 MP4, ProRes, WebM (VP9 + alpha or AV1, Opus), PNG/JPEG/TIFF/EXR/Targa sequences, GIF, WAV/AIFF | | L3 | `gpu` | The GPU compositor and GPU effects on wgpu compute shaders (Metal, Vulkan, Direct3D 12, WebGPU), checked against the CPU renderer | +| L3 | `sam2` | SAM 2.1 image segmentation and video object tracking for Roto Brush, with candle CPU inference and optional wgpu acceleration; native only, empty on wasm32, with user-installed weights | | L3 | `interchange` | Premiere Pro interop through timeline interchange (FCP7 XML, FCPXML, OTIO, EDL, AAF, OMF via FilmCraft's `filmcraft-interchange`): sequences ↔ compositions (§6a) | | L3 | `lottie` | Lottie JSON / dotLottie import and export (layers, precomps, eased and spatial keyframes, shapes, masks, mattes) with a warnings list for what Lottie cannot express | | L3 | `plugin` | Effect plug-ins: loads sandboxed WebAssembly effect modules (plug-in API v1; the wasmi interpreter behind the `wasm` feature, on for native hosts) into the effect registry ([plugins.md](plugins.md)) | From c4dc558a9830dc54d10c33e5bce0007843bf7d16 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 29 Oct 2026 21:33:33 +0100 Subject: [PATCH 61/77] 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 From 5a977297030665b80a3ee598f3e989226206da8a Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 29 Oct 2026 21:59:10 +0100 Subject: [PATCH 62/77] Roto Brush: exclude native SAM2 models from browser choices (#1539) --- crates/engine/src/analysis/models.rs | 32 +++++++++++++++++++++++++--- crates/engine/src/tests/models.rs | 28 +++++++++++++++++++++++- 2 files changed, 56 insertions(+), 4 deletions(-) diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index 9ff8a5ca7..8d45b5473 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -449,9 +449,17 @@ fn install_bytes(dir: &Path, m: &ModelInfo, bytes: &[u8]) -> std::result::Result Ok(path) } +fn models_for_platform(task: Task, wasm: bool) -> impl Iterator { + seg::models(task).filter(move |m| !wasm || !m.id.starts_with("sam2.1-hiera-")) +} + +fn target_models(task: Task) -> impl Iterator { + models_for_platform(task, cfg!(target_arch = "wasm32")) +} + fn model_p(p: &Value, cmd: &str, task: Task) -> Result<&'static ModelInfo> { let id = str_p(p, "id").ok_or_else(|| bad(cmd, "missing `id`"))?; - seg::info(id).filter(|m| m.task == task).ok_or_else(|| bad(cmd, format!("unknown model `{id}` (see {}.models)", kind(task).prefix))) + target_models(task).find(|m| m.id == id).ok_or_else(|| bad(cmd, format!("unknown model `{id}` (see {}.models)", kind(task).prefix))) } fn list(s: &mut Session, task: Task) -> Result { @@ -467,7 +475,7 @@ fn list(s: &mut Session, task: Task) -> Result { "selected": chosen == seg::CLASSICAL, "active": st.active.is_none(), })]; - for m in seg::models(task) { + for m in target_models(task) { #[cfg(not(target_arch = "wasm32"))] let (recommended, hint) = (effectcraft_sam2::mask::recommended().id == m.id, effectcraft_sam2::mask::hint(m.id)); #[cfg(target_arch = "wasm32")] @@ -528,7 +536,9 @@ fn install(s: &mut Session, p: &Value, task: Task) -> Result { Some(_) => model_p(p, &cmd, task)?, None => { let h = seg::sha256::hex(&bytes); - seg::models(task).find(|m| m.sha256 == h).ok_or_else(|| bad(&cmd, "this file is not a published model in the registry (SHA-256 differs)"))? + target_models(task) + .find(|m| m.sha256 == h) + .ok_or_else(|| bad(&cmd, "this file is not a supported published model in the registry (SHA-256 differs)"))? } }; let dest = install_bytes(&dir, m, &bytes).map_err(|e| bad(&cmd, e))?; @@ -652,3 +662,19 @@ pub fn specs() -> Vec { cmd!("face.model.remove", "Remove Face Tracking Model", [], None, "{id}", always, |s, p| remove(s, p, Task::Face)), ] } + +#[cfg(test)] +mod availability_tests { + use super::*; + + #[test] + fn browser_models_exclude_native_trackers_and_keep_portable_models() { + let browser: Vec<_> = models_for_platform(Task::Mask, true).map(|m| m.id).collect(); + assert_eq!(browser, ["mobilesam"]); + let desktop: Vec<_> = models_for_platform(Task::Mask, false).map(|m| m.id).collect(); + assert_eq!(desktop, ["mobilesam", "sam2.1-hiera-tiny", "sam2.1-hiera-small", "sam2.1-hiera-base-plus", "sam2.1-hiera-large"]); + for wasm in [false, true] { + assert_eq!(models_for_platform(Task::Face, wasm).map(|m| m.id).collect::>(), ["mediapipe-face"]); + } + } +} diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index 567b51666..28ddebf55 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -15,6 +15,32 @@ fn session(tag: &str) -> (Session, std::path::PathBuf) { (Session { config: Some(store), ..Default::default() }, d) } +#[test] +fn model_list_and_direct_selection_agree_on_target_availability() { + super::roto::hold_roto_cache_test_lock(); + let mut s = Session::default(); + let list = s.execute_checked("roto.models", json!({})).unwrap(); + let models = list["models"].as_array().unwrap(); + for id in ["sam2.1-hiera-tiny", "sam2.1-hiera-small", "sam2.1-hiera-base-plus", "sam2.1-hiera-large"] { + let chosen = s.chosen_model(Task::Mask).to_string(); + let generation = s.models.status(Task::Mask).generation; + let selected = s.execute_checked("roto.model.select", json!({"id": id})); + if cfg!(target_arch = "wasm32") { + assert!(!models.iter().any(|m| m["id"] == id)); + assert!(selected.is_err()); + assert_eq!(s.chosen_model(Task::Mask), chosen); + assert_eq!(s.models.status(Task::Mask).generation, generation, "rejection must not alter model state"); + } else { + assert!(models.iter().any(|m| m["id"] == id)); + selected.unwrap(); + assert_eq!(s.chosen_model(Task::Mask), id); + } + } + s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); + assert_eq!(s.chosen_model(Task::Mask), "classical"); + assert!(models.iter().any(|m| m["id"] == "mobilesam")); +} + /// Poll until `task`'s model is in use (or a few seconds pass). fn wait_loaded(s: &mut Session, task: Task) { for _ in 0..400 { @@ -37,7 +63,7 @@ fn models_are_listed_chosen_and_verified() { let list = s.execute_checked(&format!("{prefix}.models"), json!({})).unwrap(); let ids: Vec<&str> = list["models"].as_array().unwrap().iter().map(|m| m["id"].as_str().unwrap()).collect(); // Roto Brush also offers SAM 2.1 in four sizes (video trackers; see crates/sam2). - let expect: Vec<&str> = if task == Task::Mask { + let expect: Vec<&str> = if task == Task::Mask && !cfg!(target_arch = "wasm32") { vec!["classical", model, "sam2.1-hiera-tiny", "sam2.1-hiera-small", "sam2.1-hiera-base-plus", "sam2.1-hiera-large"] } else { vec!["classical", model] From 47e5fa436bf475fa6f563a980bd8a48a7d5a4f42 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 29 Oct 2026 22:37:14 +0100 Subject: [PATCH 63/77] SAM2: isolate model worker failures by selection generation --- crates/engine/src/analysis/models.rs | 100 ++++++++++++++++++++++++--- crates/engine/src/tests/models.rs | 4 ++ 2 files changed, 95 insertions(+), 9 deletions(-) diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index 8d45b5473..a641e01e1 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -193,11 +193,18 @@ impl Session { Task::Mask => &mut self.prefs.roto.model, Task::Face => &mut self.prefs.face.model, }; - if slot != id { + let changed = slot != id; + if changed { *slot = id.to_string(); - // a load still running for the previous choice must not activate when it finishes - set(self.models.lane(task), |s| s.generation = s.generation.wrapping_add(1)); } + set(self.models.lane(task), |s| { + // A load still running for the previous choice must not activate or report an error. + if changed { + s.generation = s.generation.wrapping_add(1); + } + // Selecting the same model explicitly retries a failed load. + s.error = None; + }); } /// Load (or drop) the models Settings chooses; called whenever settings change. @@ -219,7 +226,7 @@ impl Session { } return; }; - if st.active.as_deref() == Some(info.id) || st.busy.is_some() { + if st.active.as_deref() == Some(info.id) || st.busy.is_some() || st.error.is_some() { return; } let Some(path) = self.model_path(info) else { @@ -228,7 +235,7 @@ impl Session { }; let (face_slot, label) = (self.models.face_model.clone(), kind(task).label); let generation = st.generation; - spawn(&lane, label, format!("Loading {}…", info.name), wait, move |status| { + spawn(&lane, label, format!("Loading {}…", info.name), wait, Some(generation), move |status| { let m = load_file(&path, info.id)?; // the choice changed while this loaded: drop the model, `poll_models` loads the new choice if !activate_if_current(status, &face_slot, task, m, generation) { @@ -325,12 +332,12 @@ impl Session { self.apply_model(Task::Mask, wait); } let st = self.models.status(Task::Mask); - if let Some(e) = st.error { - return Err(format!("{} failed to load: {e}", info.name)); - } if st.active.as_deref() == Some(info.id) { return Ok(()); } + if let Some(e) = st.error { + return Err(format!("{} failed to load: {e}", info.name)); + } Err(format!("{} is still loading: try again when it is ready, or pass wait: true", info.name)) } @@ -397,6 +404,7 @@ fn spawn( label: &'static str, what: String, wait: bool, + generation: Option, job: impl FnOnce(&Arc>) -> std::result::Result + Send + 'static, ) { if cfg!(target_arch = "wasm32") { @@ -412,6 +420,11 @@ fn spawn( let r = job(&st); set(&st, |s| { s.busy = None; + if generation.is_some_and(|g| g != s.generation) { + // Keep a completion marker so polling starts the newly selected model. + s.done = Some((String::new(), false)); + return; + } match r { Ok(msg) => s.done = Some((msg, false)), Err(e) => { @@ -543,6 +556,8 @@ fn install(s: &mut Session, p: &Value, task: Task) -> Result { }; let dest = install_bytes(&dir, m, &bytes).map_err(|e| bad(&cmd, e))?; let wait = b_p(p, "wait").unwrap_or(false); + // Installing verified weights is an explicit retry, like selecting the model again. + set(s.models.lane(task), |s| s.error = None); s.apply_model(task, wait); if wait { waited(s, task, &cmd)?; @@ -560,7 +575,7 @@ fn download(s: &mut Session, p: &Value, task: Task) -> Result { } let size = if m.size >= 10_000_000 { format!("{} MB", m.size / 1_000_000) } else { format!("{:.1} MB", m.size as f64 / 1e6) }; let wait = b_p(p, "wait").unwrap_or(false); - spawn(s.models.lane(task), k.label, format!("Downloading {} ({size})…", m.name), wait, move |_| { + spawn(s.models.lane(task), k.label, format!("Downloading {} ({size})…", m.name), wait, None, move |_| { std::fs::create_dir_all(&dir).map_err(|e| format!("cannot create {}: {e}", dir.display()))?; let tmp = dir.join(format!("{}.part", m.file_name)); let out = std::process::Command::new("curl") @@ -667,6 +682,73 @@ pub fn specs() -> Vec { mod availability_tests { use super::*; + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn stale_model_worker_completions_do_not_poison_the_current_choice() { + for failure in [false, true] { + let lane = Arc::new(Mutex::new(Status::default())); + spawn(&lane, "Roto Brush", "Loading old choice".into(), true, Some(0), move |status| { + set(status, |s| s.generation += 1); + if failure { Err("old choice failed".into()) } else { Ok("old choice loaded".into()) } + }); + let status = lane.lock().unwrap(); + assert!(status.busy.is_none()); + assert!(status.error.is_none(), "an obsolete worker must not block the current model"); + assert_eq!(status.done, Some((String::new(), false)), "polling must still retry the current choice"); + } + } + + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn current_load_and_download_failures_remain_actionable() { + for generation in [Some(0), None] { + let lane = Arc::new(Mutex::new(Status::default())); + spawn(&lane, "Roto Brush", "Working".into(), true, generation, move |status| { + if generation.is_none() { + set(status, |s| s.generation += 1); + } + Err("current operation failed".into()) + }); + let status = lane.lock().unwrap(); + assert!(status.busy.is_none()); + assert_eq!(status.error.as_deref(), Some("current operation failed")); + assert_eq!(status.done, Some(("Roto Brush model: current operation failed".into(), true))); + } + } + + #[test] + fn choosing_a_model_clears_its_previous_load_error_even_on_retry() { + let mut session = Session::default(); + for id in ["mobilesam", "mobilesam", seg::CLASSICAL] { + let previous = session.chosen_model(Task::Mask).to_string(); + let generation = session.models.status(Task::Mask).generation; + set(session.models.lane(Task::Mask), |s| s.error = Some("previous load failed".into())); + session.choose_model(Task::Mask, id); + let status = session.models.status(Task::Mask); + assert!(status.error.is_none()); + assert_eq!(status.generation, generation + u64::from(previous != id)); + } + } + + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn selected_active_model_is_ready_despite_an_unrelated_download_error() { + let dir = std::env::temp_dir().join(format!("ec-model-ready-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let info = seg::info("mobilesam").unwrap(); + let file = std::fs::File::create(dir.join(info.file_name)).unwrap(); + file.set_len(info.size).unwrap(); + drop(file); + let mut session = Session { models_dir: Some(dir.clone()), ..Default::default() }; + session.choose_model(Task::Mask, info.id); + set(session.models.lane(Task::Mask), |s| { + s.active = Some(info.id.into()); + s.error = Some("another model could not be downloaded".into()); + }); + assert_eq!(session.ensure_roto_model(false), Ok(())); + std::fs::remove_dir_all(dir).unwrap(); + } + #[test] fn browser_models_exclude_native_trackers_and_keep_portable_models() { let browser: Vec<_> = models_for_platform(Task::Mask, true).map(|m| m.id).collect(); diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index 28ddebf55..e0347ffe1 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -239,6 +239,10 @@ fn reloading_for_a_stroke_does_not_block_or_retry_a_failed_load() { let st = s.models.status(Task::Mask); assert!(st.busy.is_none(), "a failed model is not loaded again by every stroke"); assert_eq!(st.error, failed); + s.poll_models(); + let st = s.models.status(Task::Mask); + assert_eq!(st.error, failed, "polling must keep the actionable error until an explicit retry"); + assert!(st.busy.is_none(), "a failed load must not restart on every poll"); // an unloaded, not failed model loads on the worker: the call returns with the load still pending or done s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); crate::effects::roto::set_model(None); From e34cea52cb08e9c82b5dfc0c6bec6ba30c80feff Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 30 Oct 2026 02:12:39 +0100 Subject: [PATCH 64/77] SAM2: separate model load failures from download failures --- crates/engine/src/analysis/models.rs | 117 +++++++++++++++++++++++---- 1 file changed, 103 insertions(+), 14 deletions(-) diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index a641e01e1..7aa50b6b9 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -33,10 +33,12 @@ pub struct Status { pub active: Option, /// The last failure, for the Settings page. pub error: Option, + /// Failure of the selected model's load; downloads cannot block or retry that load. + pub load_error: Option, /// A finished job's message and whether it failed, for a toast (taken by /// [`Session::poll_models`]). pub done: Option<(String, bool)>, - /// Counts selection changes: a load started under an older count is stale and its model is discarded. + /// Counts selection and selected-weight changes: an older load is stale and its model is discarded. pub generation: u64, } @@ -204,6 +206,7 @@ impl Session { } // Selecting the same model explicitly retries a failed load. s.error = None; + s.load_error = None; }); } @@ -226,7 +229,7 @@ impl Session { } return; }; - if st.active.as_deref() == Some(info.id) || st.busy.is_some() || st.error.is_some() { + if st.active.as_deref() == Some(info.id) || st.busy.is_some() || st.load_error.is_some() { return; } let Some(path) = self.model_path(info) else { @@ -250,6 +253,21 @@ impl Session { self.models.status(Task::Mask).active } + /// Only installing the chosen model is a retry of that model's failed load. + fn retry_installed_model(&mut self, task: Task, id: &str, wait: bool) -> bool { + if self.chosen_model(task) != id { + return false; + } + set(self.models.lane(task), |s| { + s.error = None; + s.load_error = None; + // A load that read the old file must not report failure after these weights replaced it. + s.generation = s.generation.wrapping_add(1); + }); + self.apply_model(task, wait); + true + } + /// The Roto Brush model Settings chose, when it is one (not the classic engine). fn roto_model_info(&self) -> Option<&'static ModelInfo> { seg::info(self.chosen_model(Task::Mask)).filter(|m| m.task == Task::Mask) @@ -268,7 +286,7 @@ impl Session { "installed": self.roto_model_info().is_none_or(|m| self.model_path(m).is_some()), "active": st.active, "loading": st.busy.is_some(), - "error": st.error, + "error": st.load_error.or(st.error), }) } @@ -284,7 +302,7 @@ impl Session { None } else if self.model_path(info).is_none() { Some(format!("{} is selected but not installed: propagate refuses until it is installed or Classic is chosen", info.name)) - } else if let Some(e) = st.error { + } else if let Some(e) = st.load_error { Some(format!("{} is selected but failed to load ({e}): propagate refuses until it is selected again", info.name)) } else { Some(format!("{} is selected but not loaded yet: propagate loads it first with wait: true, and refuses without it", info.name)) @@ -328,14 +346,14 @@ impl Session { return Err(format!("{} is selected for Roto Brush but not installed: install or download it, or choose {}", info.name, ROTO.classic_name)); } let st = self.models.status(Task::Mask); - if st.active.as_deref() != Some(info.id) && st.busy.is_none() && st.error.is_none() { + if st.active.as_deref() != Some(info.id) && st.busy.is_none() && st.load_error.is_none() { self.apply_model(Task::Mask, wait); } let st = self.models.status(Task::Mask); if st.active.as_deref() == Some(info.id) { return Ok(()); } - if let Some(e) = st.error { + if let Some(e) = st.load_error { return Err(format!("{} failed to load: {e}", info.name)); } Err(format!("{} is still loading: try again when it is ready, or pass wait: true", info.name)) @@ -366,7 +384,7 @@ impl Session { pub fn reload_roto_model(&mut self) { let chosen = self.chosen_model(Task::Mask).to_string(); let st = self.models.status(Task::Mask); - if chosen != seg::CLASSICAL && st.active.is_none() && st.busy.is_none() && st.error.is_none() { + if chosen != seg::CLASSICAL && st.active.is_none() && st.busy.is_none() && st.load_error.is_none() { self.apply_model(Task::Mask, false); } } @@ -408,7 +426,13 @@ fn spawn( job: impl FnOnce(&Arc>) -> std::result::Result + Send + 'static, ) { if cfg!(target_arch = "wasm32") { - set(status, |s| s.error = Some("trained models are available in the desktop app".into())); + set(status, |s| { + let message = "trained models are available in the desktop app".to_string(); + if generation.is_some() { + s.load_error = Some(message.clone()); + } + s.error = Some(message); + }); return; } set(status, |s| { @@ -429,6 +453,9 @@ fn spawn( Ok(msg) => s.done = Some((msg, false)), Err(e) => { s.done = Some((format!("{label} model: {e}"), true)); + if generation.is_some() { + s.load_error = Some(e.clone()); + } s.error = Some(e); } } @@ -439,7 +466,11 @@ fn spawn( } else if let Err(e) = std::thread::Builder::new().name("ec-model".into()).spawn(run) { set(status, |s| { s.busy = None; - s.error = Some(format!("cannot start: {e}")); + let message = format!("cannot start: {e}"); + if generation.is_some() { + s.load_error = Some(message.clone()); + } + s.error = Some(message); }); } } @@ -556,10 +587,7 @@ fn install(s: &mut Session, p: &Value, task: Task) -> Result { }; let dest = install_bytes(&dir, m, &bytes).map_err(|e| bad(&cmd, e))?; let wait = b_p(p, "wait").unwrap_or(false); - // Installing verified weights is an explicit retry, like selecting the model again. - set(s.models.lane(task), |s| s.error = None); - s.apply_model(task, wait); - if wait { + if s.retry_installed_model(task, m.id, wait) && wait { waited(s, task, &cmd)?; } Ok(json!({"id": m.id, "path": dest.display().to_string()})) @@ -694,6 +722,7 @@ mod availability_tests { let status = lane.lock().unwrap(); assert!(status.busy.is_none()); assert!(status.error.is_none(), "an obsolete worker must not block the current model"); + assert!(status.load_error.is_none()); assert_eq!(status.done, Some((String::new(), false)), "polling must still retry the current choice"); } } @@ -712,6 +741,7 @@ mod availability_tests { let status = lane.lock().unwrap(); assert!(status.busy.is_none()); assert_eq!(status.error.as_deref(), Some("current operation failed")); + assert_eq!(status.load_error.as_deref(), generation.map(|_| "current operation failed")); assert_eq!(status.done, Some(("Roto Brush model: current operation failed".into(), true))); } } @@ -722,10 +752,14 @@ mod availability_tests { for id in ["mobilesam", "mobilesam", seg::CLASSICAL] { let previous = session.chosen_model(Task::Mask).to_string(); let generation = session.models.status(Task::Mask).generation; - set(session.models.lane(Task::Mask), |s| s.error = Some("previous load failed".into())); + set(session.models.lane(Task::Mask), |s| { + s.error = Some("previous load failed".into()); + s.load_error = s.error.clone(); + }); session.choose_model(Task::Mask, id); let status = session.models.status(Task::Mask); assert!(status.error.is_none()); + assert!(status.load_error.is_none()); assert_eq!(status.generation, generation + u64::from(previous != id)); } } @@ -749,6 +783,61 @@ mod availability_tests { std::fs::remove_dir_all(dir).unwrap(); } + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn downloads_do_not_block_an_inactive_model_or_retry_a_failed_load() { + crate::tests::roto::hold_roto_cache_test_lock(); + let dir = std::env::temp_dir().join(format!("ec-model-download-state-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let info = seg::info("mobilesam").unwrap(); + std::fs::File::create(dir.join(info.file_name)).unwrap().set_len(info.size).unwrap(); + let mut session = Session { models_dir: Some(dir.clone()), ..Default::default() }; + session.choose_model(Task::Mask, info.id); + let lane = session.models.lane(Task::Mask).clone(); + spawn(&lane, "Roto Brush", "Downloading another model".into(), true, None, |_| Err("download failed".into())); + assert_eq!(session.models.status(Task::Mask).error.as_deref(), Some("download failed")); + assert!(session.models.status(Task::Mask).load_error.is_none()); + let error = session.ensure_roto_model(true).unwrap_err(); + assert!(!error.contains("download failed"), "the installed model must be loaded despite a download failure"); + let failure = session.models.status(Task::Mask).load_error; + assert!(failure.is_some(), "the synthetic weights were actually loaded and rejected"); + for download_failed in [true, false] { + spawn(&lane, "Roto Brush", "Downloading another model".into(), true, None, move |_| { + if download_failed { Err("download failed".into()) } else { Ok("downloaded".into()) } + }); + session.poll_models(); + session.reload_roto_model(); + let status = session.models.status(Task::Mask); + assert_eq!(status.load_error, failure, "a download must not clear the selected model's load failure"); + assert!(status.busy.is_none(), "polling and strokes must not retry a failed model"); + assert!(session.ensure_roto_model(true).unwrap_err().contains(failure.as_deref().unwrap())); + assert!(session.roto_model_notice(true).unwrap().contains(failure.as_deref().unwrap())); + } + crate::effects::roto::set_model(None); + std::fs::remove_dir_all(dir).unwrap(); + } + + #[test] + fn installing_an_unselected_model_keeps_the_selected_models_load_failure() { + crate::tests::roto::hold_roto_cache_test_lock(); + let mut session = Session::default(); + session.choose_model(Task::Mask, "mobilesam"); + set(session.models.lane(Task::Mask), |s| { + s.error = Some("selected model failed".into()); + s.load_error = s.error.clone(); + }); + let before = session.models.status(Task::Mask); + for wait in [false, true] { + assert!(!session.retry_installed_model(Task::Mask, "sam2.1-hiera-tiny", wait)); + assert_eq!(session.models.status(Task::Mask), before); + } + // Installing the selected weights is an explicit retry, even if the installed path + // has since gone away (no load is started from this session's missing model folder). + assert!(session.retry_installed_model(Task::Mask, "mobilesam", false)); + assert!(session.models.status(Task::Mask).load_error.is_none()); + assert!(session.models.status(Task::Mask).error.is_none()); + } + #[test] fn browser_models_exclude_native_trackers_and_keep_portable_models() { let browser: Vec<_> = models_for_platform(Task::Mask, true).map(|m| m.id).collect(); From 3fa954478d04151d386fbe80652c2d91aec7132d Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 30 Oct 2026 02:36:01 +0100 Subject: [PATCH 65/77] SAM2: wait for replaced model loads before completing installation --- crates/engine/src/analysis/models.rs | 70 ++++++++++++++++++++++------ 1 file changed, 56 insertions(+), 14 deletions(-) diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index 7aa50b6b9..cc4e8506b 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -254,9 +254,9 @@ impl Session { } /// Only installing the chosen model is a retry of that model's failed load. - fn retry_installed_model(&mut self, task: Task, id: &str, wait: bool) -> bool { + fn retry_installed_model(&mut self, task: Task, id: &str, wait: bool) -> std::result::Result { if self.chosen_model(task) != id { - return false; + return Ok(false); } set(self.models.lane(task), |s| { s.error = None; @@ -264,8 +264,22 @@ impl Session { // A load that read the old file must not report failure after these weights replaced it. s.generation = s.generation.wrapping_add(1); }); + if wait { + self.wait_model_worker(task)?; + } self.apply_model(task, wait); - true + Ok(true) + } + + fn wait_model_worker(&self, task: Task) -> std::result::Result<(), String> { + let started = Instant::now(); + while self.models.status(task).busy.is_some() { + if started.elapsed() >= ROTO_MODEL_WAIT { + return Err(format!("{} is still loading after {} minutes", self.chosen_model(task), ROTO_MODEL_WAIT.as_secs() / 60)); + } + std::thread::sleep(Duration::from_millis(5)); + } + Ok(()) } /// The Roto Brush model Settings chose, when it is one (not the classic engine). @@ -330,14 +344,8 @@ impl Session { let Some(info) = self.roto_model_info() else { return Ok(()) }; // A download or load under way (from Settings, or `select` without `wait`) finishes first. // Nothing is in flight in the browser build, so this never sleeps there. - if wait && self.models.status(Task::Mask).busy.is_some() { - let started = Instant::now(); - while self.models.status(Task::Mask).busy.is_some() { - if started.elapsed() >= ROTO_MODEL_WAIT { - return Err(format!("{} is still loading after {} minutes", info.name, ROTO_MODEL_WAIT.as_secs() / 60)); - } - std::thread::sleep(Duration::from_millis(5)); - } + if wait { + self.wait_model_worker(Task::Mask)?; } if self.model_path(info).is_none() { if self.models.status(Task::Mask).busy.is_some() { @@ -587,7 +595,7 @@ fn install(s: &mut Session, p: &Value, task: Task) -> Result { }; let dest = install_bytes(&dir, m, &bytes).map_err(|e| bad(&cmd, e))?; let wait = b_p(p, "wait").unwrap_or(false); - if s.retry_installed_model(task, m.id, wait) && wait { + if s.retry_installed_model(task, m.id, wait).map_err(|error| bad(&cmd, error))? && wait { waited(s, task, &cmd)?; } Ok(json!({"id": m.id, "path": dest.display().to_string()})) @@ -828,16 +836,50 @@ mod availability_tests { }); let before = session.models.status(Task::Mask); for wait in [false, true] { - assert!(!session.retry_installed_model(Task::Mask, "sam2.1-hiera-tiny", wait)); + assert!(!session.retry_installed_model(Task::Mask, "sam2.1-hiera-tiny", wait).unwrap()); assert_eq!(session.models.status(Task::Mask), before); } // Installing the selected weights is an explicit retry, even if the installed path // has since gone away (no load is started from this session's missing model folder). - assert!(session.retry_installed_model(Task::Mask, "mobilesam", false)); + assert!(session.retry_installed_model(Task::Mask, "mobilesam", false).unwrap()); assert!(session.models.status(Task::Mask).load_error.is_none()); assert!(session.models.status(Task::Mask).error.is_none()); } + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn installing_selected_weights_with_wait_finishes_the_replacement_load() { + crate::tests::roto::hold_roto_cache_test_lock(); + let dir = std::env::temp_dir().join(format!("ec-model-install-running-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let info = seg::info("mobilesam").unwrap(); + std::fs::File::create(dir.join(info.file_name)).unwrap().set_len(info.size).unwrap(); + let mut session = Session { models_dir: Some(dir.clone()), ..Default::default() }; + session.choose_model(Task::Mask, info.id); + let generation = session.models.status(Task::Mask).generation; + let lane = session.models.lane(Task::Mask).clone(); + spawn(&lane, "Roto Brush", "Loading the previous file".into(), false, Some(generation), move |status| { + let started = Instant::now(); + while status.lock().unwrap().generation == generation { + if started.elapsed() > Duration::from_secs(3) { + return Err("test did not start the replacement".into()); + } + std::thread::sleep(Duration::from_millis(1)); + } + Err("previous file failed".into()) + }); + assert!(session.retry_installed_model(Task::Mask, info.id, true).unwrap()); + let status = session.models.status(Task::Mask); + assert!(status.busy.is_none()); + assert!( + status.load_error.as_ref().is_some_and(|error| error.contains("SHA-256")), + "wait must finish loading the replacement weights and report their actual validation error" + ); + assert_ne!(status.error.as_deref(), Some("previous file failed")); + crate::effects::roto::set_model(None); + std::fs::remove_dir_all(dir).unwrap(); + } + #[test] fn browser_models_exclude_native_trackers_and_keep_portable_models() { let browser: Vec<_> = models_for_platform(Task::Mask, true).map(|m| m.id).collect(); From 0fba3429d90033a7167dcba53bdb45a2af7f6149 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 30 Oct 2026 03:20:07 +0100 Subject: [PATCH 66/77] #1539: wait only for the selected model worker --- crates/engine/src/analysis/models.rs | 107 ++++++++++++++++++++++----- 1 file changed, 88 insertions(+), 19 deletions(-) diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index cc4e8506b..a64e880d1 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -29,10 +29,14 @@ use crate::{EngineError, Result, Session, cmd, query}; pub struct Status { /// What is happening now ("Downloading MobileSAM…"). pub busy: Option, + /// Model whose load/download owns this lane, for selection-specific waits. + pub busy_model: Option<&'static str>, /// The model in use. pub active: Option, /// The last failure, for the Settings page. pub error: Option, + /// Model whose operation failed, so Roto status does not report another model's download. + pub error_model: Option<&'static str>, /// Failure of the selected model's load; downloads cannot block or retry that load. pub load_error: Option, /// A finished job's message and whether it failed, for a toast (taken by @@ -206,6 +210,7 @@ impl Session { } // Selecting the same model explicitly retries a failed load. s.error = None; + s.error_model = None; s.load_error = None; }); } @@ -238,7 +243,7 @@ impl Session { }; let (face_slot, label) = (self.models.face_model.clone(), kind(task).label); let generation = st.generation; - spawn(&lane, label, format!("Loading {}…", info.name), wait, Some(generation), move |status| { + spawn(&lane, label, format!("Loading {}…", info.name), wait, Some(generation), info.id, move |status| { let m = load_file(&path, info.id)?; // the choice changed while this loaded: drop the model, `poll_models` loads the new choice if !activate_if_current(status, &face_slot, task, m, generation) { @@ -260,26 +265,40 @@ impl Session { } set(self.models.lane(task), |s| { s.error = None; + s.error_model = None; s.load_error = None; // A load that read the old file must not report failure after these weights replaced it. s.generation = s.generation.wrapping_add(1); }); + if self.models.status(task).active.as_deref() == Some(id) { + return Ok(true); + } if wait { - self.wait_model_worker(task)?; + self.wait_model_worker(task, id)?; } self.apply_model(task, wait); + if wait { + let status = self.models.status(task); + if status.active.as_deref() != Some(id) { + return Err(status.load_error.unwrap_or_else(|| format!("{id} is not loaded: try again when the model worker is ready"))); + } + } Ok(true) } - fn wait_model_worker(&self, task: Task) -> std::result::Result<(), String> { + fn wait_model_worker(&self, task: Task, id: &str) -> std::result::Result<(), String> { let started = Instant::now(); - while self.models.status(task).busy.is_some() { + loop { + let status = self.models.status(task); + let Some(operation) = status.busy else { return Ok(()) }; + if status.busy_model != Some(id) { + return Err(format!("{id} cannot load while another model operation is running ({operation}): try again when it is ready")); + } if started.elapsed() >= ROTO_MODEL_WAIT { return Err(format!("{} is still loading after {} minutes", self.chosen_model(task), ROTO_MODEL_WAIT.as_secs() / 60)); } std::thread::sleep(Duration::from_millis(5)); } - Ok(()) } /// The Roto Brush model Settings chose, when it is one (not the classic engine). @@ -289,7 +308,9 @@ impl Session { /// Whether the chosen Roto Brush model is still loading or downloading. pub fn roto_model_loading(&self) -> bool { - self.models.status(Task::Mask).busy.is_some() && self.chosen_model(Task::Mask) != seg::CLASSICAL + let status = self.models.status(Task::Mask); + let chosen = self.chosen_model(Task::Mask); + status.active.as_deref() != Some(chosen) && status.busy.is_some() && status.busy_model == Some(chosen) } /// The Roto Brush model Settings chose and its state (`roto.status`). @@ -299,8 +320,8 @@ impl Session { "selected": self.chosen_model(Task::Mask), "installed": self.roto_model_info().is_none_or(|m| self.model_path(m).is_some()), "active": st.active, - "loading": st.busy.is_some(), - "error": st.load_error.or(st.error), + "loading": self.roto_model_loading(), + "error": st.load_error.or(st.error.filter(|_| st.error_model == Some(self.chosen_model(Task::Mask)))), }) } @@ -342,10 +363,13 @@ impl Session { /// refused either way. The classic engine needs nothing. pub fn ensure_roto_model(&mut self, wait: bool) -> std::result::Result<(), String> { let Some(info) = self.roto_model_info() else { return Ok(()) }; + if self.models.status(Task::Mask).active.as_deref() == Some(info.id) { + return Ok(()); + } // A download or load under way (from Settings, or `select` without `wait`) finishes first. // Nothing is in flight in the browser build, so this never sleeps there. if wait { - self.wait_model_worker(Task::Mask)?; + self.wait_model_worker(Task::Mask, info.id)?; } if self.model_path(info).is_none() { if self.models.status(Task::Mask).busy.is_some() { @@ -431,6 +455,7 @@ fn spawn( what: String, wait: bool, generation: Option, + model: &'static str, job: impl FnOnce(&Arc>) -> std::result::Result + Send + 'static, ) { if cfg!(target_arch = "wasm32") { @@ -440,18 +465,22 @@ fn spawn( s.load_error = Some(message.clone()); } s.error = Some(message); + s.error_model = Some(model); }); return; } set(status, |s| { s.busy = Some(what); + s.busy_model = Some(model); s.error = None; + s.error_model = None; }); let st = status.clone(); let run = move || { let r = job(&st); set(&st, |s| { s.busy = None; + s.busy_model = None; if generation.is_some_and(|g| g != s.generation) { // Keep a completion marker so polling starts the newly selected model. s.done = Some((String::new(), false)); @@ -465,6 +494,7 @@ fn spawn( s.load_error = Some(e.clone()); } s.error = Some(e); + s.error_model = Some(model); } } }); @@ -474,11 +504,13 @@ fn spawn( } else if let Err(e) = std::thread::Builder::new().name("ec-model".into()).spawn(run) { set(status, |s| { s.busy = None; + s.busy_model = None; let message = format!("cannot start: {e}"); if generation.is_some() { s.load_error = Some(message.clone()); } s.error = Some(message); + s.error_model = Some(model); }); } } @@ -595,9 +627,7 @@ fn install(s: &mut Session, p: &Value, task: Task) -> Result { }; let dest = install_bytes(&dir, m, &bytes).map_err(|e| bad(&cmd, e))?; let wait = b_p(p, "wait").unwrap_or(false); - if s.retry_installed_model(task, m.id, wait).map_err(|error| bad(&cmd, error))? && wait { - waited(s, task, &cmd)?; - } + s.retry_installed_model(task, m.id, wait).map_err(|error| bad(&cmd, error))?; Ok(json!({"id": m.id, "path": dest.display().to_string()})) } @@ -611,7 +641,7 @@ fn download(s: &mut Session, p: &Value, task: Task) -> Result { } let size = if m.size >= 10_000_000 { format!("{} MB", m.size / 1_000_000) } else { format!("{:.1} MB", m.size as f64 / 1e6) }; let wait = b_p(p, "wait").unwrap_or(false); - spawn(s.models.lane(task), k.label, format!("Downloading {} ({size})…", m.name), wait, None, move |_| { + spawn(s.models.lane(task), k.label, format!("Downloading {} ({size})…", m.name), wait, None, m.id, move |_| { std::fs::create_dir_all(&dir).map_err(|e| format!("cannot create {}: {e}", dir.display()))?; let tmp = dir.join(format!("{}.part", m.file_name)); let out = std::process::Command::new("curl") @@ -723,7 +753,7 @@ mod availability_tests { fn stale_model_worker_completions_do_not_poison_the_current_choice() { for failure in [false, true] { let lane = Arc::new(Mutex::new(Status::default())); - spawn(&lane, "Roto Brush", "Loading old choice".into(), true, Some(0), move |status| { + spawn(&lane, "Roto Brush", "Loading old choice".into(), true, Some(0), "mobilesam", move |status| { set(status, |s| s.generation += 1); if failure { Err("old choice failed".into()) } else { Ok("old choice loaded".into()) } }); @@ -740,7 +770,7 @@ mod availability_tests { fn current_load_and_download_failures_remain_actionable() { for generation in [Some(0), None] { let lane = Arc::new(Mutex::new(Status::default())); - spawn(&lane, "Roto Brush", "Working".into(), true, generation, move |status| { + spawn(&lane, "Roto Brush", "Working".into(), true, generation, "mobilesam", move |status| { if generation.is_none() { set(status, |s| s.generation += 1); } @@ -788,9 +818,48 @@ mod availability_tests { s.error = Some("another model could not be downloaded".into()); }); assert_eq!(session.ensure_roto_model(false), Ok(())); + assert!(session.roto_model()["error"].is_null()); + set(session.models.lane(Task::Mask), |status| { + status.error_model = Some("mobilesam"); + status.error = Some("selected model download failed".into()); + }); + assert_eq!(session.roto_model()["error"], "selected model download failed"); std::fs::remove_dir_all(dir).unwrap(); } + #[test] + fn an_active_model_and_selected_install_do_not_wait_for_other_downloads() { + crate::tests::roto::hold_roto_cache_test_lock(); + let mut session = Session::default(); + session.choose_model(Task::Mask, "mobilesam"); + set(session.models.lane(Task::Mask), |status| { + status.active = Some("mobilesam".into()); + status.busy = Some("Downloading SAM 2".into()); + status.busy_model = Some("sam2.1-hiera-tiny"); + }); + assert_eq!(session.ensure_roto_model(true), Ok(())); + assert_eq!(session.retry_installed_model(Task::Mask, "mobilesam", true), Ok(true)); + let status = session.models.status(Task::Mask); + assert!(status.busy.is_some(), "the unrelated operation must still be running"); + assert!(!session.roto_model_loading()); + assert_eq!(session.roto_model()["loading"], false); + } + + #[test] + fn waiting_on_an_inactive_model_refuses_unrelated_work_without_waiting() { + crate::tests::roto::hold_roto_cache_test_lock(); + let mut session = Session::default(); + session.choose_model(Task::Mask, "mobilesam"); + set(session.models.lane(Task::Mask), |status| { + status.busy = Some("Downloading SAM 2".into()); + status.busy_model = Some("sam2.1-hiera-tiny"); + }); + assert!(session.ensure_roto_model(true).unwrap_err().contains("another model operation")); + assert!(session.retry_installed_model(Task::Mask, "mobilesam", true).unwrap_err().contains("another model operation")); + assert!(session.models.status(Task::Mask).busy.is_some()); + assert!(!session.roto_model_loading()); + } + #[cfg(not(target_arch = "wasm32"))] #[test] fn downloads_do_not_block_an_inactive_model_or_retry_a_failed_load() { @@ -802,7 +871,7 @@ mod availability_tests { let mut session = Session { models_dir: Some(dir.clone()), ..Default::default() }; session.choose_model(Task::Mask, info.id); let lane = session.models.lane(Task::Mask).clone(); - spawn(&lane, "Roto Brush", "Downloading another model".into(), true, None, |_| Err("download failed".into())); + spawn(&lane, "Roto Brush", "Downloading another model".into(), true, None, "sam2.1-hiera-tiny", |_| Err("download failed".into())); assert_eq!(session.models.status(Task::Mask).error.as_deref(), Some("download failed")); assert!(session.models.status(Task::Mask).load_error.is_none()); let error = session.ensure_roto_model(true).unwrap_err(); @@ -810,7 +879,7 @@ mod availability_tests { let failure = session.models.status(Task::Mask).load_error; assert!(failure.is_some(), "the synthetic weights were actually loaded and rejected"); for download_failed in [true, false] { - spawn(&lane, "Roto Brush", "Downloading another model".into(), true, None, move |_| { + spawn(&lane, "Roto Brush", "Downloading another model".into(), true, None, "sam2.1-hiera-tiny", move |_| { if download_failed { Err("download failed".into()) } else { Ok("downloaded".into()) } }); session.poll_models(); @@ -858,7 +927,7 @@ mod availability_tests { session.choose_model(Task::Mask, info.id); let generation = session.models.status(Task::Mask).generation; let lane = session.models.lane(Task::Mask).clone(); - spawn(&lane, "Roto Brush", "Loading the previous file".into(), false, Some(generation), move |status| { + spawn(&lane, "Roto Brush", "Loading the previous file".into(), false, Some(generation), info.id, move |status| { let started = Instant::now(); while status.lock().unwrap().generation == generation { if started.elapsed() > Duration::from_secs(3) { @@ -868,7 +937,7 @@ mod availability_tests { } Err("previous file failed".into()) }); - assert!(session.retry_installed_model(Task::Mask, info.id, true).unwrap()); + assert!(session.retry_installed_model(Task::Mask, info.id, true).unwrap_err().contains("SHA-256")); let status = session.models.status(Task::Mask); assert!(status.busy.is_none()); assert!( From 33f545cf647572d52a323a72269ca41740f35729 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 30 Oct 2026 03:51:09 +0100 Subject: [PATCH 67/77] #1539: keep SAM2 weights install-only until mirrored --- crates/engine/src/analysis/models.rs | 16 +++++++++++++++ crates/engine/src/tests/models.rs | 12 ++++++++---- crates/segment/README.md | 3 ++- crates/segment/src/lib.rs | 28 ++++++++++++++++----------- crates/ui-egui/src/panels/settings.rs | 26 ++++++++++++++++++++++++- docs/agents.md | 5 +++-- docs/roto-brush-sam2.md | 10 +++++----- 7 files changed, 76 insertions(+), 24 deletions(-) diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index a64e880d1..5d5cf436d 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -635,6 +635,9 @@ fn download(s: &mut Session, p: &Value, task: Task) -> Result { let k = kind(task); let cmd = format!("{}.model.download", k.prefix); let m = model_p(p, &cmd, task)?; + if m.url.is_empty() { + return Err(bad(&cmd, format!("{} has no app download available: use Install from File with the SHA-256-verified weights", m.name))); + } let dir = s.models_dir().ok_or_else(|| bad(&cmd, "no settings folder to install into"))?; if s.models.status(task).busy.is_some() { return Err(bad(&cmd, "a model is already downloading or loading")); @@ -748,6 +751,19 @@ pub fn specs() -> Vec { mod availability_tests { use super::*; + /// SAM 2 cannot start a network operation; file installation still verifies the published hash. + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn sam2_is_install_from_file_only() { + let mut session = Session::default(); + for model in [seg::SAM2_TINY, seg::SAM2_SMALL, seg::SAM2_BASE_PLUS, seg::SAM2_LARGE] { + let error = session.execute("roto.model.download", json!({"id": model.id})).unwrap_err(); + assert!(error.to_string().contains("Install from File")); + assert!(session.models.status(Task::Mask).busy.is_none()); + assert!(install_bytes(Path::new("unused"), &model, b"invalid weights").unwrap_err().contains("SHA-256 differs")); + } + } + #[cfg(not(target_arch = "wasm32"))] #[test] fn stale_model_worker_completions_do_not_poison_the_current_choice() { diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index e0347ffe1..c34d02ed7 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -249,19 +249,23 @@ fn reloading_for_a_stroke_does_not_block_or_retry_a_failed_load() { let _ = std::fs::remove_dir_all(d); } -/// With `EFFECTCRAFT_DOWNLOAD_SIZE=sam2.1-hiera-small` (or another registry id): the in-app download (Settings ▸ Roto -/// Brush ▸ Download) fetches the pinned file over the network, verifies it, installs it with its notice and loads it. -/// Skipped without the variable: it needs the network and hundreds of megabytes. +/// With `EFFECTCRAFT_DOWNLOAD_SIZE=mobilesam` (or another downloadable registry id): the in-app +/// download verifies and loads the file. SAM 2 must refuse until a storytold mirror is available. +/// Skipped without the variable: downloadable entries need the network and large files. #[cfg(not(target_arch = "wasm32"))] #[test] fn downloading_a_size_installs_and_loads_it() { let Ok(id) = std::env::var("EFFECTCRAFT_DOWNLOAD_SIZE") else { return }; crate::tests_roto::hold_roto_cache_test_lock(); let (mut s, d) = session("download"); + let info = effectcraft_segment::info(&id).unwrap(); + if info.url.is_empty() { + assert!(s.execute_checked("roto.model.download", json!({"id": id, "wait": true})).unwrap_err().to_string().contains("Install from File")); + return; + } let r = s.execute_checked("roto.model.download", json!({"id": id, "wait": true})).unwrap(); assert_eq!(r["id"], id.as_str(), "{r}"); let dir = s.models_dir().unwrap(); - let info = effectcraft_segment::info(&id).unwrap(); let file = dir.join(info.file_name); assert_eq!(std::fs::metadata(&file).unwrap().len(), info.size, "the whole file arrived"); assert!(dir.join(format!("{}.NOTICE.txt", info.file_name)).exists()); diff --git a/crates/segment/README.md b/crates/segment/README.md index e6e4c6c98..06ffce7ad 100644 --- a/crates/segment/README.md +++ b/crates/segment/README.md @@ -12,7 +12,8 @@ Trained models for Roto Brush and face tracking, each kind behind one swappable depends on one model's layout. The classical face tracker stays built in as the fallback. - **`MODELS`**: the registry. Every entry must be open source under a licence compatible with EffectCraft's MIT OR Apache-2.0 (`ALLOWED_LICENCES`), with its task, authors, homepage, licence - URL, the official weights URL, size and SHA-256. `load` refuses anything else. + URL, an app download URL when available, size and SHA-256. `load` refuses anything else. + SAM 2 is Install from File only until a storytold-controlled mirror is available. - **Weights are never bundled.** They are downloaded on demand (or installed from a file) by the engine (`roto.model.*` / `face.model.*`, Settings ▸ Roto Brush / Face Tracking), verified against the registry's SHA-256 (`sha256`, dependency-free) and kept with a notice naming their diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index 04894cb2b..0d9cc3134 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -141,7 +141,7 @@ pub struct ModelInfo { pub licence_url: &'static str, /// The project's page. pub homepage: &'static str, - /// The official weights file. + /// App download URL; empty when only Install from File is available. pub url: &'static str, pub file_name: &'static str, pub size: u64, @@ -168,8 +168,8 @@ pub const MOBILE_SAM: ModelInfo = ModelInfo { }; /// SAM 2.1 Tiny (Meta AI, Apache-2.0): a video object tracker with streaming memory. The weights are -/// the Hugging Face `transformers` checkpoint, pinned to a commit so the SHA-256 stays valid; they -/// are loaded by `effectcraft-sam2` (native only), not by [`load`]. +/// the Hugging Face `transformers` checkpoint, identified by SHA-256. Until a storytold mirror +/// is available, install it from a file. `effectcraft-sam2` (native only) loads it, not [`load`]. pub const SAM2_TINY: ModelInfo = ModelInfo { id: "sam2.1-hiera-tiny", task: Task::Mask, @@ -179,7 +179,7 @@ pub const SAM2_TINY: ModelInfo = ModelInfo { licence: "Apache-2.0", licence_url: "https://github.com/facebookresearch/sam2/blob/main/LICENSE", homepage: "https://github.com/facebookresearch/sam2", - url: "https://huggingface.co/facebook/sam2.1-hiera-tiny/resolve/de431c4043854a71d8101e17995dfe596bf101a5/model.safetensors", + url: "", file_name: "sam2.1-hiera-tiny.safetensors", size: 155_908_064, sha256: "48c14467e5cf9e51870511feb72c89688e82dd74523142c0538b663e193ac2a7", @@ -195,7 +195,7 @@ pub const SAM2_SMALL: ModelInfo = ModelInfo { licence: "Apache-2.0", licence_url: "https://github.com/facebookresearch/sam2/blob/main/LICENSE", homepage: "https://github.com/facebookresearch/sam2", - url: "https://huggingface.co/facebook/sam2.1-hiera-small/resolve/ee5bba1d82bb8749febdf90f45e84b687142ba03/model.safetensors", + url: "", file_name: "sam2.1-hiera-small.safetensors", size: 184_305_280, sha256: "0a4067b11ce1e23d5229203f11c718a823060d15a4b23fa2372a7d4b77cbbc60", @@ -211,7 +211,7 @@ pub const SAM2_BASE_PLUS: ModelInfo = ModelInfo { licence: "Apache-2.0", licence_url: "https://github.com/facebookresearch/sam2/blob/main/LICENSE", homepage: "https://github.com/facebookresearch/sam2", - url: "https://huggingface.co/facebook/sam2.1-hiera-base-plus/resolve/b7320756a13354e7530a63935656d35b2f91a290/model.safetensors", + url: "", file_name: "sam2.1-hiera-base-plus.safetensors", size: 323_476_296, sha256: "2012733a0de5d03efd1bba550a2847c4551be9ef2e0d497c83074df66189f780", @@ -227,7 +227,7 @@ pub const SAM2_LARGE: ModelInfo = ModelInfo { licence: "Apache-2.0", licence_url: "https://github.com/facebookresearch/sam2/blob/main/LICENSE", homepage: "https://github.com/facebookresearch/sam2", - url: "https://huggingface.co/facebook/sam2.1-hiera-large/resolve/665f8e2ad61cf5f53d65644ff27c8ee525124610/model.safetensors", + url: "", file_name: "sam2.1-hiera-large.safetensors", size: 897_897_416, sha256: "dc407dce21301fd94abb395c5099b4f2c455fdc8a8f261ac3d0ea6d4cd197230", @@ -314,9 +314,10 @@ fn loader(id: &str) -> Option { /// The notice kept next to installed weights: what they are, who made them, their licence. pub fn notice(m: &ModelInfo) -> String { + let source = if m.url.is_empty() { m.homepage } else { m.url }; format!( - "{}\n\nAuthors: {}\nLicence: {} (full text: https://spdx.org/licenses/{}.html; as stated by the authors: {})\nSource: {}\nProject: {}\n\nDownloaded unmodified from the source above. EffectCraft runs it with its own implementation\nand is not affiliated with or endorsed by its authors.\n", - m.name, m.authors, m.licence, m.licence, m.licence_url, m.url, m.homepage + "{}\n\nAuthors: {}\nLicence: {} (full text: https://spdx.org/licenses/{}.html; as stated by the authors: {})\nSource: {}\nProject: {}\n\nInstalled unmodified. EffectCraft runs it with its own implementation\nand is not affiliated with or endorsed by its authors.\n", + m.name, m.authors, m.licence, m.licence, m.licence_url, source, m.homepage ) } @@ -418,14 +419,19 @@ mod tests { for m in MODELS { assert!(ALLOWED_LICENCES.contains(&m.licence), "{}", m.id); assert_eq!(m.sha256.len(), 64); - assert!(m.url.starts_with("https://") && m.licence_url.starts_with("https://") && m.homepage.starts_with("https://")); + assert!(m.url.is_empty() || m.url.starts_with("https://")); + assert!(m.licence_url.starts_with("https://") && m.homepage.starts_with("https://")); assert!(!m.authors.is_empty() && notice(m).contains(m.authors)); } assert_eq!( models(Task::Mask).map(|m| m.id).collect::>(), ["mobilesam", "sam2.1-hiera-tiny", "sam2.1-hiera-small", "sam2.1-hiera-base-plus", "sam2.1-hiera-large"] ); - assert!(SAM2_TINY.url.contains("/resolve/") && SAM2_TINY.file_name.ends_with(".safetensors"), "pinned to a commit, not a branch"); + for model in [SAM2_TINY, SAM2_SMALL, SAM2_BASE_PLUS, SAM2_LARGE] { + assert!(model.url.is_empty(), "SAM 2 has no app download until a storytold mirror is available"); + assert!(model.file_name.ends_with(".safetensors")); + assert!(notice(&model).contains(&format!("Source: {}", model.homepage))); + } assert_eq!(models(Task::Face).map(|m| m.id).collect::>(), ["mediapipe-face"]); assert!(load("mediapipe-face", b"nope").err().is_some_and(|e| e.contains("does not match"))); // A file that isn't the published one is refused before anything is parsed. diff --git a/crates/ui-egui/src/panels/settings.rs b/crates/ui-egui/src/panels/settings.rs index 8b4133fbc..5fdbf1c50 100644 --- a/crates/ui-egui/src/panels/settings.rs +++ b/crates/ui-egui/src/panels/settings.rs @@ -513,7 +513,7 @@ fn models_ui(app: &mut EffectcraftApp, ui: &mut egui::Ui, t: &Tokens, acts: &mut return; } ui.horizontal(|ui| { - if !installed && !web { + if !installed && !web && m["url"].as_str().is_some_and(|url| !url.is_empty()) { let r = ui.add_enabled(busy.is_none(), egui::Button::new(tr("Download"))); reg(app, &format!("settings.{prefix}.download.{id}"), &r, "Download"); if r.clicked() { @@ -924,3 +924,27 @@ pub fn recovery(app: &mut EffectcraftApp, ctx: &egui::Context, t: &Tokens) { app.ui.status = e; } } + +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + + /// Install-only models must not offer a download that the command will reject. + #[test] + fn sam2_settings_offer_file_installation_without_downloads() { + let mut app = EffectcraftApp::new(effectcraft_engine::Session::default()); + let ctx = egui::Context::default(); + let tokens = app.tokens; + let mut actions = vec![]; + let mut output = ctx + .run_ui(egui::RawInput { screen_rect: Some(egui::Rect::from_min_size(egui::Pos2::ZERO, vec2(900.0, 2200.0))), ..Default::default() }, |ui| { + models_ui(&mut app, ui, &tokens, &mut actions, Task::Mask) + }); + output.textures_delta.clear(); + for id in ["sam2.1-hiera-tiny", "sam2.1-hiera-small", "sam2.1-hiera-base-plus", "sam2.1-hiera-large"] { + assert!(app.auto.find(&format!("settings.roto.download.{id}")).is_none()); + assert!(app.auto.find(&format!("settings.roto.install.{id}")).is_some()); + } + assert!(actions.is_empty()); + } +} diff --git a/docs/agents.md b/docs/agents.md index ef88acccb..1a5035f8f 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -504,8 +504,9 @@ reference. Matte settings are ordinary properties (`effects/#1/rotoBrushMatte/se Trained model (Roto Brush 2.0 / 3.0): - `roto.models` lists the built-in engine and every registered model, with its authors, licence, size, URL, SHA-256 and whether it is installed, selected and active. -- `roto.model.download {id}` fetches and verifies the official weights in the background; - `roto.model.install {path, id?}` installs a file you already have. +- `roto.model.download {id}` fetches and verifies weights when the registry provides an app download URL; + `roto.model.install {path, id?}` installs a file you already have. SAM 2 is Install from File only until + a storytold-controlled mirror is available; its size and SHA-256 are verified before installation. - `roto.model.select {id: classical|mobilesam|sam2.1-hiera-tiny|sam2.1-hiera-small|sam2.1-hiera-base-plus|sam2.1-hiera-large}` chooses which one Roto Brush uses (SAM 2 also tracks the object through the clip, see [roto-brush-sam2.md](roto-brush-sam2.md); `roto.status` then reports `source` and per-frame `tracker.frames` with `confidence` and `occluded`). Loading runs in the background; `"wait": true` (also on download and install) returns once it is in use. - `roto.propagate` never falls back to the classic engine silently: a chosen model that is not installed, or failed diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index f2a2cb20d..5b8038552 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -4,10 +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 ▸ **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.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. +1. Settings ▸ Roto Brush ▸ SAM 2.1 Tiny ▸ **Install from File…** with the published `model.safetensors` + (155.9 MB). The file's size and SHA-256 are checked and the licence notice is kept next to it. + Agents: `roto.model.install {id: "sam2.1-hiera-tiny", path: "model.safetensors"}`. + SAM 2 has no in-app download until the weights are mirrored under storytold control. 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 @@ -59,7 +59,7 @@ instances can see different pixels. Objects inside one instance always see the s Four checkpoints are in the registry (Settings ▸ Roto Brush); the size is detected from the file, and each is pinned to a Hugging Face commit with its SHA-256. All run the same tracker; a bigger backbone costs time and memory. -| Model | Download | GPU encode (RX 7900 GRE) | Agreement with PyTorch | +| Model | File size | GPU encode (RX 7900 GRE) | Agreement with PyTorch | |---|---|---|---| | SAM 2.1 Tiny (`sam2.1-hiera-tiny`) | 155.9 MB | 0.24 s | encoder 6.8e-6 | | SAM 2.1 Small (`sam2.1-hiera-small`) | 184.3 MB | 0.20 s | encoder 1.0e-5 | From a21b9b175ec7704e82f93f4b4357437a3e6e9899 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 30 Oct 2026 04:22:21 +0100 Subject: [PATCH 68/77] #1539: retain checkpoint provenance and clarify model waits --- crates/engine/src/analysis/models.rs | 32 +++++++++++++++++++++++++++ crates/segment/README.md | 2 ++ crates/segment/src/lib.rs | 18 +++++++++++---- crates/ui-egui/src/panels/settings.rs | 8 +++++++ docs/agents.md | 13 +++++++---- docs/roto-brush-sam2.md | 5 ++++- 6 files changed, 69 insertions(+), 9 deletions(-) diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index 5d5cf436d..63550545a 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -575,6 +575,8 @@ fn list(s: &mut Session, task: Task) -> Result { "licenceUrl": m.licence_url, "homepage": m.homepage, "url": m.url, + "source": m.source, + "sourceFile": m.source.rsplit('/').next().unwrap_or(m.file_name), "size": m.size, "sha256": m.sha256, "installed": s.model_path(m).is_some(), @@ -600,6 +602,9 @@ fn select(s: &mut Session, p: &Value, task: Task) -> Result { model_p(p, &cmd, task)?; } let wait = b_p(p, "wait").unwrap_or(false); + if wait && id != seg::CLASSICAL && s.models.status(task).active.as_deref() != Some(id) { + s.wait_model_worker(task, id).map_err(|error| bad(&cmd, error))?; + } s.choose_model(task, id); if wait { s.apply_model(task, true); @@ -756,11 +761,38 @@ mod availability_tests { #[test] fn sam2_is_install_from_file_only() { let mut session = Session::default(); + let listed = session.execute("roto.models", json!({})).unwrap(); for model in [seg::SAM2_TINY, seg::SAM2_SMALL, seg::SAM2_BASE_PLUS, seg::SAM2_LARGE] { let error = session.execute("roto.model.download", json!({"id": model.id})).unwrap_err(); assert!(error.to_string().contains("Install from File")); assert!(session.models.status(Task::Mask).busy.is_none()); assert!(install_bytes(Path::new("unused"), &model, b"invalid weights").unwrap_err().contains("SHA-256 differs")); + let row = listed["models"].as_array().unwrap().iter().find(|row| row["id"] == model.id).unwrap(); + assert_eq!(row["source"], model.source); + assert_eq!(row["sourceFile"], "model.safetensors"); + } + } + + /// A waiting selection cannot claim completion while unrelated work prevents its model loading. + #[test] + fn waiting_selection_refuses_an_unrelated_worker_unless_the_model_is_active() { + crate::tests::roto::hold_roto_cache_test_lock(); + for active in [false, true] { + let mut session = Session::default(); + session.choose_model(Task::Mask, if active { "mobilesam" } else { seg::CLASSICAL }); + set(session.models.lane(Task::Mask), |status| { + status.active = active.then(|| "mobilesam".into()); + status.busy = Some("Downloading SAM 2".into()); + status.busy_model = Some("sam2.1-hiera-tiny"); + }); + let result = session.execute("roto.model.select", json!({"id": "mobilesam", "wait": true})); + if active { + assert_eq!(result.unwrap()["active"], "mobilesam"); + } else { + assert!(result.unwrap_err().to_string().contains("another model operation")); + assert_eq!(session.chosen_model(Task::Mask), seg::CLASSICAL, "a refused selection must not change settings"); + } + assert!(session.models.status(Task::Mask).busy.is_some()); } } diff --git a/crates/segment/README.md b/crates/segment/README.md index 06ffce7ad..a95c92f93 100644 --- a/crates/segment/README.md +++ b/crates/segment/README.md @@ -14,6 +14,8 @@ Trained models for Roto Brush and face tracking, each kind behind one swappable EffectCraft's MIT OR Apache-2.0 (`ALLOWED_LICENCES`), with its task, authors, homepage, licence URL, an app download URL when available, size and SHA-256. `load` refuses anything else. SAM 2 is Install from File only until a storytold-controlled mirror is available. + The checkpoint `source` stays separate from the app download URL, so file-install guidance and + licence notices identify the verified artifact even when app downloads are unavailable. - **Weights are never bundled.** They are downloaded on demand (or installed from a file) by the engine (`roto.model.*` / `face.model.*`, Settings ▸ Roto Brush / Face Tracking), verified against the registry's SHA-256 (`sha256`, dependency-free) and kept with a notice naming their diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index 0d9cc3134..7fc0e4ff6 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -143,6 +143,8 @@ pub struct ModelInfo { pub homepage: &'static str, /// App download URL; empty when only Install from File is available. pub url: &'static str, + /// Checkpoint source for manual installation and notices; never used by the app downloader. + pub source: &'static str, pub file_name: &'static str, pub size: u64, /// SHA-256 of the weights file (lower-case hex). @@ -162,6 +164,7 @@ pub const MOBILE_SAM: ModelInfo = ModelInfo { licence_url: "https://github.com/ChaoningZhang/MobileSAM/blob/master/LICENSE", homepage: "https://github.com/ChaoningZhang/MobileSAM", url: "https://raw.githubusercontent.com/ChaoningZhang/MobileSAM/master/weights/mobile_sam.pt", + source: "https://raw.githubusercontent.com/ChaoningZhang/MobileSAM/master/weights/mobile_sam.pt", file_name: "mobile_sam.pt", size: 40_728_226, sha256: "6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f", @@ -180,6 +183,7 @@ pub const SAM2_TINY: ModelInfo = ModelInfo { licence_url: "https://github.com/facebookresearch/sam2/blob/main/LICENSE", homepage: "https://github.com/facebookresearch/sam2", url: "", + source: "https://huggingface.co/facebook/sam2.1-hiera-tiny/blob/de431c4043854a71d8101e17995dfe596bf101a5/model.safetensors", file_name: "sam2.1-hiera-tiny.safetensors", size: 155_908_064, sha256: "48c14467e5cf9e51870511feb72c89688e82dd74523142c0538b663e193ac2a7", @@ -196,6 +200,7 @@ pub const SAM2_SMALL: ModelInfo = ModelInfo { licence_url: "https://github.com/facebookresearch/sam2/blob/main/LICENSE", homepage: "https://github.com/facebookresearch/sam2", url: "", + source: "https://huggingface.co/facebook/sam2.1-hiera-small/blob/ee5bba1d82bb8749febdf90f45e84b687142ba03/model.safetensors", file_name: "sam2.1-hiera-small.safetensors", size: 184_305_280, sha256: "0a4067b11ce1e23d5229203f11c718a823060d15a4b23fa2372a7d4b77cbbc60", @@ -212,6 +217,7 @@ pub const SAM2_BASE_PLUS: ModelInfo = ModelInfo { licence_url: "https://github.com/facebookresearch/sam2/blob/main/LICENSE", homepage: "https://github.com/facebookresearch/sam2", url: "", + source: "https://huggingface.co/facebook/sam2.1-hiera-base-plus/blob/b7320756a13354e7530a63935656d35b2f91a290/model.safetensors", file_name: "sam2.1-hiera-base-plus.safetensors", size: 323_476_296, sha256: "2012733a0de5d03efd1bba550a2847c4551be9ef2e0d497c83074df66189f780", @@ -228,6 +234,7 @@ pub const SAM2_LARGE: ModelInfo = ModelInfo { licence_url: "https://github.com/facebookresearch/sam2/blob/main/LICENSE", homepage: "https://github.com/facebookresearch/sam2", url: "", + source: "https://huggingface.co/facebook/sam2.1-hiera-large/blob/665f8e2ad61cf5f53d65644ff27c8ee525124610/model.safetensors", file_name: "sam2.1-hiera-large.safetensors", size: 897_897_416, sha256: "dc407dce21301fd94abb395c5099b4f2c455fdc8a8f261ac3d0ea6d4cd197230", @@ -243,6 +250,7 @@ pub const FACE_LANDMARKER: ModelInfo = ModelInfo { licence_url: "https://storage.googleapis.com/mediapipe-assets/Model%20Card%20MediaPipe%20Face%20Mesh%20V2.pdf", homepage: "https://developers.google.com/edge/mediapipe/solutions/vision/face_landmarker", url: "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task", + source: "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task", file_name: "face_landmarker.task", size: 3_758_596, sha256: "64184e229b263107bc2b804c6625db1341ff2bb731874b0bcc2fe6544e0bc9ff", @@ -314,10 +322,9 @@ fn loader(id: &str) -> Option { /// The notice kept next to installed weights: what they are, who made them, their licence. pub fn notice(m: &ModelInfo) -> String { - let source = if m.url.is_empty() { m.homepage } else { m.url }; format!( - "{}\n\nAuthors: {}\nLicence: {} (full text: https://spdx.org/licenses/{}.html; as stated by the authors: {})\nSource: {}\nProject: {}\n\nInstalled unmodified. EffectCraft runs it with its own implementation\nand is not affiliated with or endorsed by its authors.\n", - m.name, m.authors, m.licence, m.licence, m.licence_url, source, m.homepage + "{}\n\nAuthors: {}\nLicence: {} (full text: https://spdx.org/licenses/{}.html; as stated by the authors: {})\nSource: {}\nSHA-256: {}\nProject: {}\n\nInstalled unmodified. EffectCraft runs it with its own implementation\nand is not affiliated with or endorsed by its authors.\n", + m.name, m.authors, m.licence, m.licence, m.licence_url, m.source, m.sha256, m.homepage ) } @@ -420,6 +427,7 @@ mod tests { assert!(ALLOWED_LICENCES.contains(&m.licence), "{}", m.id); assert_eq!(m.sha256.len(), 64); assert!(m.url.is_empty() || m.url.starts_with("https://")); + assert!(m.source.starts_with("https://")); assert!(m.licence_url.starts_with("https://") && m.homepage.starts_with("https://")); assert!(!m.authors.is_empty() && notice(m).contains(m.authors)); } @@ -430,7 +438,9 @@ mod tests { for model in [SAM2_TINY, SAM2_SMALL, SAM2_BASE_PLUS, SAM2_LARGE] { assert!(model.url.is_empty(), "SAM 2 has no app download until a storytold mirror is available"); assert!(model.file_name.ends_with(".safetensors")); - assert!(notice(&model).contains(&format!("Source: {}", model.homepage))); + assert!(model.source.contains("/blob/") && model.source.ends_with("/model.safetensors")); + assert!(notice(&model).contains(&format!("Source: {}", model.source))); + assert!(notice(&model).contains(model.sha256)); } assert_eq!(models(Task::Face).map(|m| m.id).collect::>(), ["mediapipe-face"]); assert!(load("mediapipe-face", b"nope").err().is_some_and(|e| e.contains("does not match"))); diff --git a/crates/ui-egui/src/panels/settings.rs b/crates/ui-egui/src/panels/settings.rs index 5fdbf1c50..69f7623dd 100644 --- a/crates/ui-egui/src/panels/settings.rs +++ b/crates/ui-egui/src/panels/settings.rs @@ -508,6 +508,13 @@ fn models_ui(app: &mut EffectcraftApp, ui: &mut egui::Ui, t: &Tokens, acts: &mut } } } + if let (Some(source), Some(file)) = (m["source"].as_str(), m["sourceFile"].as_str()) { + let response = ui.link(file).on_hover_text(crate::i18n::tr("Source")); + reg(app, &format!("settings.{prefix}.source.{id}"), &response, file); + if response.clicked() { + ui.ctx().open_url(egui::OpenUrl::new_tab(source)); + } + } }); if classic { return; @@ -944,6 +951,7 @@ mod tests { for id in ["sam2.1-hiera-tiny", "sam2.1-hiera-small", "sam2.1-hiera-base-plus", "sam2.1-hiera-large"] { assert!(app.auto.find(&format!("settings.roto.download.{id}")).is_none()); assert!(app.auto.find(&format!("settings.roto.install.{id}")).is_some()); + assert_eq!(app.auto.find(&format!("settings.roto.source.{id}")).unwrap().label, "model.safetensors"); } assert!(actions.is_empty()); } diff --git a/docs/agents.md b/docs/agents.md index 1a5035f8f..d43cbd689 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -395,8 +395,9 @@ Face tracking's trained model works like Roto Brush's (below): `face.models` lis tracker and MediaPipe Face Landmarker with its authors and licence; `face.model.download {id}` / `face.model.install {path, id?}` install it (verified); `face.model.select {id: classical|mediapipe-face}` chooses it; `face.model.remove {id}` deletes it. Downloads and loading -run in the background; `"wait": true` on download, install or select returns once the model is in -use (or with its error), which scripts need before tracking. `track.mask` with a +run in the background; `"wait": true` waits for the requested operation (or returns its error). +Download/install do not change the selected model. After installing, use +`face.model.select {id, wait:true}` to select and load it before tracking. `track.mask` with a face method returns `faceModel`, the engine it tries first (the classical one when the model finds no face in the mask). @@ -503,12 +504,16 @@ reference. Matte settings are ordinary properties (`effects/#1/rotoBrushMatte/se Trained model (Roto Brush 2.0 / 3.0): - `roto.models` lists the built-in engine and every registered model, with its authors, - licence, size, URL, SHA-256 and whether it is installed, selected and active. + licence, size, app download URL (empty for install-only models), checkpoint `source` and `sourceFile`, + SHA-256 and whether it is installed, selected and active. - `roto.model.download {id}` fetches and verifies weights when the registry provides an app download URL; `roto.model.install {path, id?}` installs a file you already have. SAM 2 is Install from File only until a storytold-controlled mirror is available; its size and SHA-256 are verified before installation. - `roto.model.select {id: classical|mobilesam|sam2.1-hiera-tiny|sam2.1-hiera-small|sam2.1-hiera-base-plus|sam2.1-hiera-large}` chooses which one Roto Brush uses (SAM 2 also tracks the object through the clip, see [roto-brush-sam2.md](roto-brush-sam2.md); `roto.status` then reports `source` and per-frame `tracker.frames` with `confidence` and `occluded`). Loading runs - in the background; `"wait": true` (also on download and install) returns once it is in use. + in the background. `"wait": true` waits for the requested operation; download/install load the model + only if it is already selected. They do not activate an unselected model: after installation, use + `roto.model.select {id, wait:true}` to select and load it. Unrelated model work can return a retryable + busy error rather than reporting that a waiting selection has finished. - `roto.propagate` never falls back to the classic engine silently: a chosen model that is not installed, or failed to load, is an error, and without `"wait": true` so is a model still loading. `roto.status` reports `model` (`selected`, `installed`, `active`, `loading`, `error`) and a `notice` while the chosen model is not the one in use. diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index 5b8038552..7078fa103 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -8,7 +8,10 @@ Apache-2.0). It is optional: nothing changes until you install the model and cho (155.9 MB). The file's size and SHA-256 are checked and the licence notice is kept next to it. Agents: `roto.model.install {id: "sam2.1-hiera-tiny", path: "model.safetensors"}`. SAM 2 has no in-app download until the weights are mirrored under storytold control. -2. Choose it (`roto.model.select {id: "sam2.1-hiera-tiny"}`). + Each model card links its pinned Hugging Face `model.safetensors` checkpoint, rather than the + project's different `.pt` files. Its checkpoint source and SHA-256 are retained in the notice. +2. Choose it (`roto.model.select {id: "sam2.1-hiera-tiny", wait: true}` for agents). Installing an + unselected model does not activate it, even with `wait: true` on installation. 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 object leaves the frame and returns. Each frame's probability map is then refined by the same graph cut as every other From 05384890536c570a5d5220fa74e2d7151eb88dac Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 30 Oct 2026 04:46:19 +0100 Subject: [PATCH 69/77] #1539: omit point padding when a box prompt is present --- crates/sam2/src/prompt.rs | 48 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 47 insertions(+), 1 deletion(-) diff --git a/crates/sam2/src/prompt.rs b/crates/sam2/src/prompt.rs index 91496bba3..129ce192b 100644 --- a/crates/sam2/src/prompt.rs +++ b/crates/sam2/src/prompt.rs @@ -140,7 +140,6 @@ impl PromptEncoder { let v: Vec = e.iter().zip(self.point_row(row)).map(|(a, b)| a + b).collect(); push(&v); } - push(&self.not_a_point.clone()); } } Ok(Tensor::from_vec(tokens, (1, n, DIM), &Device::Cpu)?) @@ -176,6 +175,53 @@ impl PromptEncoder { mod tests { use super::*; + /// A tiny deterministic checkpoint exercises token labels without downloading trained weights. + fn encoder() -> PromptEncoder { + let mut map = std::collections::HashMap::new(); + for (name, shape) in [ + ("shared_embedding.positional_embedding", vec![2, DIM / 2]), + ("no_mask_embed.weight", vec![1, DIM]), + ("mask_embed.conv1.weight", vec![4, 1, 2, 2]), + ("mask_embed.conv1.bias", vec![4]), + ("mask_embed.conv2.weight", vec![16, 4, 2, 2]), + ("mask_embed.conv2.bias", vec![16]), + ("mask_embed.conv3.weight", vec![DIM, 16, 1, 1]), + ("mask_embed.conv3.bias", vec![DIM]), + ("mask_embed.layer_norm1.weight", vec![4]), + ("mask_embed.layer_norm1.bias", vec![4]), + ("mask_embed.layer_norm2.weight", vec![16]), + ("mask_embed.layer_norm2.bias", vec![16]), + ] { + map.insert(format!("prompt_encoder.{name}"), Tensor::zeros(shape, candle_core::DType::F32, &Device::Cpu).unwrap()); + } + let rows: Vec = (1..=4).flat_map(|row| vec![row as f32 * 10.0; DIM]).collect(); + map.insert("prompt_encoder.point_embed.weight".into(), Tensor::from_vec(rows, (4, DIM), &Device::Cpu).unwrap()); + map.insert("prompt_encoder.not_a_point_embed.weight".into(), Tensor::from_vec(vec![-1000.0f32; DIM], (1, DIM), &Device::Cpu).unwrap()); + PromptEncoder::load(&Weights::from_map(map).unwrap()).unwrap() + } + + /// The reference pads only point-only prompts; box corners are already complete tokens. + #[test] + fn boxes_do_not_add_point_padding_and_click_only_prompts_still_do() { + let encoder = encoder(); + let points = vec![([10.0, 20.0], false), ([30.0, 40.0], true)]; + for points in [vec![], points.clone()] { + let sparse = encoder.sparse(&Prompt { points: points.clone(), bbox: Some([0.0, 0.0, 100.0, 100.0]) }).unwrap(); + assert_eq!(sparse.dims(), &[1, points.len() + 2, DIM]); + let rows = sparse.to_vec3::().unwrap(); + for (token, label) in rows[0].iter().zip(if points.is_empty() { vec![30.0, 40.0] } else { vec![10.0, 20.0, 30.0, 40.0] }) { + assert!(token[..DIM / 2].iter().all(|value| *value == label)); + assert!(token[DIM / 2..].iter().all(|value| *value == label + 1.0)); + } + } + let clicks = encoder.sparse(&Prompt { points, bbox: None }).unwrap(); + assert_eq!(clicks.dims(), &[1, 3, DIM]); + assert!(clicks.to_vec3::().unwrap()[0][2].iter().all(|value| *value == -1000.0)); + let none = encoder.sparse_none().unwrap(); + assert_eq!(none.dims(), &[1, 2, DIM]); + assert!(none.flatten_all().unwrap().to_vec1::().unwrap().iter().all(|value| *value == -1000.0)); + } + #[test] fn hostile_prompts_are_rejected() { let ok = Prompt { points: vec![([10.0, 20.0], true)], bbox: None }; From fa46bbcb8843aebfd21a349919cc1f6f966f4e0f Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 30 Oct 2026 08:06:06 +0100 Subject: [PATCH 70/77] #1539: match reference box padding and use a flat GPU module --- crates/sam2/src/{gpu/mod.rs => gpu.rs} | 2 +- crates/sam2/src/prompt.rs | 11 +++++++---- 2 files changed, 8 insertions(+), 5 deletions(-) rename crates/sam2/src/{gpu/mod.rs => gpu.rs} (99%) diff --git a/crates/sam2/src/gpu/mod.rs b/crates/sam2/src/gpu.rs similarity index 99% rename from crates/sam2/src/gpu/mod.rs rename to crates/sam2/src/gpu.rs index 3193dcf80..bc83b605f 100644 --- a/crates/sam2/src/gpu/mod.rs +++ b/crates/sam2/src/gpu.rs @@ -13,7 +13,7 @@ use wgpu::util::DeviceExt; use crate::{Error, Result}; -const SHADER: &str = concat!(include_str!("kernels.wgsl"), include_str!("kernels_ops.wgsl")); +const SHADER: &str = concat!(include_str!("gpu/kernels.wgsl"), include_str!("gpu/kernels_ops.wgsl")); /// Words in a kernel's parameter block. const PARAM_WORDS: usize = 24; diff --git a/crates/sam2/src/prompt.rs b/crates/sam2/src/prompt.rs index 129ce192b..23d9a7c99 100644 --- a/crates/sam2/src/prompt.rs +++ b/crates/sam2/src/prompt.rs @@ -112,7 +112,7 @@ impl PromptEncoder { self.point_embed.get(i * DIM..(i + 1) * DIM).unwrap_or(&[]) } - /// Sparse tokens `[1, n, 256]` for a prompt (a click-only prompt gets one padding token). + /// Sparse tokens `[1, n, 256]`; the pinned reference pads both click-only and box prompts. pub fn sparse(&self, p: &Prompt) -> Result { p.validate()?; let norm = |c: [f32; 2]| [(c[0] + 0.5) / INPUT_SIZE as f32, (c[1] + 0.5) / INPUT_SIZE as f32]; @@ -140,6 +140,8 @@ impl PromptEncoder { let v: Vec = e.iter().zip(self.point_row(row)).map(|(a, b)| a + b).collect(); push(&v); } + // _embed_boxes adds its own no-point token, separately from _embed_points. + push(&self.not_a_point.clone()); } } Ok(Tensor::from_vec(tokens, (1, n, DIM), &Device::Cpu)?) @@ -200,15 +202,16 @@ mod tests { PromptEncoder::load(&Weights::from_map(map).unwrap()).unwrap() } - /// The reference pads only point-only prompts; box corners are already complete tokens. + /// Transformers 5.19.0 pads click-only prompts in _embed_points and boxes in _embed_boxes. #[test] - fn boxes_do_not_add_point_padding_and_click_only_prompts_still_do() { + fn box_and_click_tokens_match_the_pinned_reference_padding() { let encoder = encoder(); let points = vec![([10.0, 20.0], false), ([30.0, 40.0], true)]; for points in [vec![], points.clone()] { let sparse = encoder.sparse(&Prompt { points: points.clone(), bbox: Some([0.0, 0.0, 100.0, 100.0]) }).unwrap(); - assert_eq!(sparse.dims(), &[1, points.len() + 2, DIM]); + assert_eq!(sparse.dims(), &[1, points.len() + 3, DIM]); let rows = sparse.to_vec3::().unwrap(); + assert!(rows[0].last().unwrap().iter().all(|value| *value == -1000.0)); for (token, label) in rows[0].iter().zip(if points.is_empty() { vec![30.0, 40.0] } else { vec![10.0, 20.0, 30.0, 40.0] }) { assert!(token[..DIM / 2].iter().all(|value| *value == label)); assert!(token[DIM / 2..].iter().all(|value| *value == label + 1.0)); From baa1ca31b3d7a17fc0e8cb09b8209e59dfbb3ebd Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 30 Oct 2026 08:39:43 +0100 Subject: [PATCH 71/77] #1539: preserve stroke corrections and current main APIs --- Cargo.lock | 2 +- apps/effectcraft-cli/tests/it/cli.rs | 41 ----------- crates/effects/src/roto.rs | 2 +- crates/engine/src/tests/models.rs | 69 ++++++++++--------- crates/sam2/src/gpu/encoder.rs | 2 +- crates/sam2/src/gpu/ops.rs | 2 +- crates/sam2/src/gpu/tests.rs | 2 +- crates/track/src/roto.rs | 16 ++--- crates/track/src/roto/tests.rs | 39 ++++++++++- crates/ui-egui/src/panels/roto_tool.rs | 6 +- crates/ui-egui/tests/it/main.rs | 2 + .../ui-egui/tests/{ => it}/ui_roto_objects.rs | 1 + .../tests/{ => it}/ui_roto_span_confidence.rs | 1 + docs/effects.md | 2 +- 14 files changed, 94 insertions(+), 93 deletions(-) rename crates/ui-egui/tests/{ => it}/ui_roto_objects.rs (99%) rename crates/ui-egui/tests/{ => it}/ui_roto_span_confidence.rs (99%) diff --git a/Cargo.lock b/Cargo.lock index b104f63cf..7bd644400 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1838,7 +1838,7 @@ dependencies = [ "effectcraft-segment", "half", "log", - "pollster 0.4.0", + "pollster 1.0.1", "rayon", "wgpu", ] diff --git a/apps/effectcraft-cli/tests/it/cli.rs b/apps/effectcraft-cli/tests/it/cli.rs index 8bdfb3c00..93c899839 100644 --- a/apps/effectcraft-cli/tests/it/cli.rs +++ b/apps/effectcraft-cli/tests/it/cli.rs @@ -870,44 +870,3 @@ fn render_without_progress_writes_no_progress_lines() { let stderr = render_demo_stderr("no-progress.gif", &[]); assert!(!stderr.lines().any(|l| l.starts_with("progress")), "{stderr}"); } - -/// `bench --roto --clip FRAMES --truth MASKS` runs the trackers on the user's own image sequences: here a bright disc -/// drifting over a dark textured frame, with its true masks. Mismatched sequences are refused with a message. -#[test] -fn bench_roto_runs_on_a_users_clip() { - let (frames, masks) = (tmp("clip-frames"), tmp("clip-masks")); - for d in [&frames, &masks] { - let _ = std::fs::remove_dir_all(d); - std::fs::create_dir_all(d).unwrap(); - } - let (w, h) = (96u32, 72u32); - for f in 0..6u32 { - let (cx, cy) = (30.0 + 4.0 * f as f32, 36.0f32); - let (mut img, mut mask) = (image::RgbImage::new(w, h), image::GrayImage::new(w, h)); - for y in 0..h { - for x in 0..w { - let inside = ((x as f32 - cx).powi(2) + (y as f32 - cy).powi(2)).sqrt() < 14.0; - let tex = 40 + ((x * 7 + y * 13) % 30) as u8; - img.put_pixel(x, y, if inside { image::Rgb([230, 120, 40 + (x % 20) as u8]) } else { image::Rgb([tex / 2, tex, tex + 20]) }); - mask.put_pixel(x, y, image::Luma([if inside { 255 } else { 0 }])); - } - } - img.save(frames.join(format!("f{f:03}.png"))).unwrap(); - mask.save(masks.join(format!("f{f:03}.png"))).unwrap(); - } - let (f, m) = (frames.to_str().unwrap(), masks.to_str().unwrap()); - let out = bin().args(["bench", "--roto", "--clip", f, "--truth", m, "--frames", "6"]).output().unwrap(); - let text = String::from_utf8_lossy(&out.stdout).to_string(); - assert!(out.status.success(), "{text}\n{}", String::from_utf8_lossy(&out.stderr)); - let row = text.lines().find(|l| l.contains("classical")).unwrap_or_else(|| panic!("no classical row: {text}")); - let mean: f64 = row.split_whitespace().nth(2).unwrap().parse().unwrap(); - assert!(mean > 0.8, "classical mean IoU on the user's clip: {mean}\n{text}"); - // one mask short: refused - std::fs::remove_file(masks.join("f005.png")).unwrap(); - let out = bin().args(["bench", "--roto", "--clip", f, "--truth", m]).output().unwrap(); - assert!(!out.status.success()); - assert!(String::from_utf8_lossy(&out.stderr).contains("must match"), "{}", String::from_utf8_lossy(&out.stderr)); - // --clip without --truth - let out = bin().args(["bench", "--roto", "--clip", f]).output().unwrap(); - assert!(!out.status.success()); -} diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index e66400416..1a1ce105f 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -696,7 +696,7 @@ impl Chain { let strokes = self.data.strokes_of(f, obj); let p = match frame(f) { Some(img) => rb::tracker_prompt_boxed(&strokes, self.opts.scale, &rb::segment(&img, &strokes, None, &self.opts, self.refine_radius)), - None => rb::tracker_prompt(&strokes, self.opts.scale), + None => return Err(format!("cannot read Roto Brush prompt frame {f}")), }; if p.points.iter().any(|(_, fg)| *fg) { prompts.push((f, p)); diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index c34d02ed7..e1baab2d8 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -101,37 +101,8 @@ 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()); } - // The SAM 2.1 sizes, each from its own file (EFFECTCRAFT_SAM2, _SAM2_SMALL, _SAM2_BASE_PLUS, _SAM2_LARGE = - // path/to/model.safetensors of facebook/sam2.1-hiera-*): installed with its notice, loaded (verified, native - // only) and used by Roto Brush 3.0, then dropped and removed. #[cfg(not(target_arch = "wasm32"))] - for (var, id) in [ - ("EFFECTCRAFT_SAM2", "sam2.1-hiera-tiny"), - ("EFFECTCRAFT_SAM2_SMALL", "sam2.1-hiera-small"), - ("EFFECTCRAFT_SAM2_BASE_PLUS", "sam2.1-hiera-base-plus"), - ("EFFECTCRAFT_SAM2_LARGE", "sam2.1-hiera-large"), - ] { - let Ok(path) = std::env::var(var) else { continue }; - let r = s.execute_checked("roto.model.install", json!({"path": path})).unwrap(); - assert_eq!(r["id"], id, "{var}"); - let dir = s.models_dir().unwrap(); - let notice = std::fs::read_to_string(dir.join(format!("{id}.safetensors.NOTICE.txt"))).unwrap(); - assert!(notice.contains("Meta AI") && notice.contains("Apache-2.0"), "{notice}"); - s.execute_checked("roto.model.select", json!({"id": id})).unwrap(); - wait_loaded(&mut s, Task::Mask); - assert_eq!(s.models.status(Task::Mask).active.as_deref(), Some(id), "{var}"); - 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(id)); - // the list says it is installed, and exactly one size is recommended - let list = s.execute_checked("roto.models", json!({})).unwrap(); - let entry = list["models"].as_array().unwrap().iter().find(|m| m["id"] == id).unwrap(); - assert_eq!(entry["installed"], true); - assert_eq!(list["models"].as_array().unwrap().iter().filter(|m| m["recommended"] == true).count(), 1); - s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); - assert!(crate::effects::roto::model_for(¶ms).is_none()); - assert_eq!(s.execute_checked("roto.model.remove", json!({"id": id})).unwrap()["removed"], true); - assert!(!dir.join(format!("{id}.safetensors")).exists(), "{id} removed"); - } + verify_installed_sam2_sizes(&mut s); // 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") { @@ -150,6 +121,40 @@ fn models_are_listed_chosen_and_verified() { let _ = std::fs::remove_dir_all(d); } +#[cfg(not(target_arch = "wasm32"))] +fn verify_installed_sam2_sizes(s: &mut Session) { + // The SAM 2.1 sizes, each from its own file (EFFECTCRAFT_SAM2, _SAM2_SMALL, _SAM2_BASE_PLUS, _SAM2_LARGE = + // path/to/model.safetensors of facebook/sam2.1-hiera-*): installed with its notice, loaded (verified, native + // only) and used by Roto Brush 3.0, then dropped and removed. + for (var, id) in [ + ("EFFECTCRAFT_SAM2", "sam2.1-hiera-tiny"), + ("EFFECTCRAFT_SAM2_SMALL", "sam2.1-hiera-small"), + ("EFFECTCRAFT_SAM2_BASE_PLUS", "sam2.1-hiera-base-plus"), + ("EFFECTCRAFT_SAM2_LARGE", "sam2.1-hiera-large"), + ] { + let Ok(path) = std::env::var(var) else { continue }; + let r = s.execute_checked("roto.model.install", json!({"path": path})).unwrap(); + assert_eq!(r["id"], id, "{var}"); + let dir = s.models_dir().unwrap(); + let notice = std::fs::read_to_string(dir.join(format!("{id}.safetensors.NOTICE.txt"))).unwrap(); + assert!(notice.contains("Meta AI") && notice.contains("Apache-2.0"), "{notice}"); + s.execute_checked("roto.model.select", json!({"id": id})).unwrap(); + wait_loaded(s, Task::Mask); + assert_eq!(s.models.status(Task::Mask).active.as_deref(), Some(id), "{var}"); + 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(id)); + // the list says it is installed, and exactly one size is recommended + let list = s.execute_checked("roto.models", json!({})).unwrap(); + let entry = list["models"].as_array().unwrap().iter().find(|m| m["id"] == id).unwrap(); + assert_eq!(entry["installed"], true); + assert_eq!(list["models"].as_array().unwrap().iter().filter(|m| m["recommended"] == true).count(), 1); + s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); + assert!(crate::effects::roto::model_for(¶ms).is_none()); + assert_eq!(s.execute_checked("roto.model.remove", json!({"id": id})).unwrap()["removed"], true); + assert!(!dir.join(format!("{id}.safetensors")).exists(), "{id} removed"); + } +} + /// A Roto Brush model that has been idle for ten minutes is unloaded (the choice stays), one in use is not, /// and a model that does not track use is never unloaded. #[test] @@ -167,7 +172,7 @@ fn an_idle_roto_model_is_unloaded() { self.0 } } - crate::tests_roto::hold_roto_cache_test_lock(); + crate::tests::roto::hold_roto_cache_test_lock(); let (mut s, d) = session("idle"); for (idle, unloaded) in [(Some(5.0), false), (None, false), (Some(601.0), true)] { let m: Arc = Arc::new(Fake(idle)); @@ -256,7 +261,7 @@ fn reloading_for_a_stroke_does_not_block_or_retry_a_failed_load() { #[test] fn downloading_a_size_installs_and_loads_it() { let Ok(id) = std::env::var("EFFECTCRAFT_DOWNLOAD_SIZE") else { return }; - crate::tests_roto::hold_roto_cache_test_lock(); + crate::tests::roto::hold_roto_cache_test_lock(); let (mut s, d) = session("download"); let info = effectcraft_segment::info(&id).unwrap(); if info.url.is_empty() { diff --git a/crates/sam2/src/gpu/encoder.rs b/crates/sam2/src/gpu/encoder.rs index 1d390a9d1..a33f9069c 100644 --- a/crates/sam2/src/gpu/encoder.rs +++ b/crates/sam2/src/gpu/encoder.rs @@ -237,7 +237,7 @@ impl GpuEncoder { let px = g.upload(&pixels.contiguous()?.flatten_all()?.to_vec1::()?); let (mut h, mut w) = (ph / 4, pw / 4); let tokens = g.zeros(h * w * hidden)?; - g.patch_embed(&px, &self.patch_w, &self.patch_b, &tokens, ph, pw, hidden)?; + g.patch_embed(&px, &self.patch_w, &self.patch_b, &tokens, (ph, pw, hidden))?; let mut x = g.zeros(h * w * hidden)?; g.add(&tokens, &self.pos_buffer(h, w)?, &x, h * w * hidden)?; let mut stages: Vec<(Buf, usize, usize, usize)> = vec![]; diff --git a/crates/sam2/src/gpu/ops.rs b/crates/sam2/src/gpu/ops.rs index 3d9b8f184..08143724e 100644 --- a/crates/sam2/src/gpu/ops.rs +++ b/crates/sam2/src/gpu/ops.rs @@ -169,7 +169,7 @@ impl Gpu { } /// The patch embedding: `[3, h, w]` pixels, weight `[co, 3, 7, 7]`, bias `[co]`; output `[h/4, w/4, co]`. - pub fn patch_embed(&self, pixels: &Buf, weight: &Buf, bias: &Buf, out: &Buf, h: usize, w: usize, co: usize) -> Result<()> { + pub fn patch_embed(&self, pixels: &Buf, weight: &Buf, bias: &Buf, out: &Buf, (h, w, co): (usize, usize, usize)) -> Result<()> { let (ho, wo) = ((h + 6 - 7) / 4 + 1, (w + 6 - 7) / 4 + 1); need(pixels, 3 * h * w, "patch_embed pixels")?; need(weight, co * 3 * 49, "patch_embed weight")?; diff --git a/crates/sam2/src/gpu/tests.rs b/crates/sam2/src/gpu/tests.rs index 674bd3b8b..e974c34df 100644 --- a/crates/sam2/src/gpu/tests.rs +++ b/crates/sam2/src/gpu/tests.rs @@ -340,7 +340,7 @@ fn patch_embed_rope_and_upsample_match_the_cpu_definitions() { let (gp, gw, gb) = (g.upload(&px), g.upload(&wt), g.upload(&bias)); let (ho, wo) = (h / 4, w / 4); let out = g.zeros(ho * wo * co).unwrap(); - g.patch_embed(&gp, &gw, &gb, &out, h, w, co).unwrap(); + g.patch_embed(&gp, &gw, &gb, &out, (h, w, co)).unwrap(); let t = Tensor::from_vec(px, (1, 3, h, w), &Device::Cpu).unwrap(); let k = Tensor::from_vec(wt, (co, 3, 7, 7), &Device::Cpu).unwrap(); let want = diff --git a/crates/track/src/roto.rs b/crates/track/src/roto.rs index 18c6eaa77..c742654e9 100644 --- a/crates/track/src/roto.rs +++ b/crates/track/src/roto.rs @@ -950,18 +950,18 @@ fn points_along(s: &Stroke, scale: f64, max: usize) -> Vec<[f32; 2]> { /// SAM 2 was trained on a handful of clicks: with sixteen or more points it doubts its own mask /// (measured: confidence 0.25 with 16 points, 0.00 with 32, 0.99 with 1 to 7), so this picks at /// most [`TRACKER_FG_POINTS`] foreground and [`TRACKER_BG_POINTS`] background points, spread -/// along the strokes. -pub fn tracker_prompt(strokes: &[&Stroke], scale: f64) -> Prompt { +/// along the strokes. Like [`stroke_prompt`], later corrections win over earlier paint. +pub fn tracker_prompt(strokes: &[&Stroke], scale: f64, w: usize, h: usize) -> Prompt { + let hard = resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, scale, w, h); + let prompt = stroke_prompt(strokes, &hard, scale, w, h); let mut points = vec![]; - for (kind, fg, max) in [(StrokeKind::Fg, true, TRACKER_FG_POINTS), (StrokeKind::Bg, false, TRACKER_BG_POINTS)] { - let all: Vec<[f32; 2]> = strokes.iter().filter(|s| s.kind == kind).flat_map(|s| points_along(s, scale, 3)).collect(); + for (fg, max) in [(true, TRACKER_FG_POINTS), (false, TRACKER_BG_POINTS)] { + let all: Vec<[f32; 2]> = prompt.points.iter().filter_map(|(p, label)| (*label == fg).then_some(*p)).collect(); let n = all.len().min(max); // evenly spaced picks, keeping the first and the last points.extend((0..n).filter_map(|i| all.get(if n > 1 { i * (all.len() - 1) / (n - 1) } else { all.len() / 2 }).map(|p| (*p, fg)))); } - // a box the user drew on the frame prompts the tracker too - let bbox = strokes.iter().filter(|s| s.kind == StrokeKind::Box).find_map(|s| box_of(s, scale)); - Prompt { points, bbox, ..Default::default() } + Prompt { points, ..prompt } } /// [`tracker_prompt`] plus a box around the classic engine's matte of the same strokes. Clicks alone @@ -969,7 +969,7 @@ pub fn tracker_prompt(strokes: &[&Stroke], scale: f64) -> Prompt { /// multi-part figure, against 1.00 with the box); the box says how far the object reaches. A box /// stroke the user drew wins over the matte's. pub fn tracker_prompt_boxed(strokes: &[&Stroke], scale: f64, classic: &FrameSeg) -> Prompt { - let mut p = tracker_prompt(strokes, scale); + let mut p = tracker_prompt(strokes, scale, classic.w, classic.h); if p.bbox.is_none() { let (w, h) = (classic.w, classic.h); let (mut x0, mut y0, mut x1, mut y1) = (usize::MAX, usize::MAX, 0usize, 0usize); diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index 50ccaeaa4..bfc00df74 100644 --- a/crates/track/src/roto/tests.rs +++ b/crates/track/src/roto/tests.rs @@ -591,13 +591,46 @@ fn a_correction_keeps_the_model_prompted_with_what_is_left() { fn tracker_prompts_carry_a_handful_of_points() { let st = base_strokes(0); let refs: Vec<&Stroke> = st.iter().collect(); - let p = tracker_prompt(&refs, 1.0); + let p = tracker_prompt(&refs, 1.0, W as usize, H as usize); let (fg, bg) = (p.points.iter().filter(|(_, f)| *f).count(), p.points.iter().filter(|(_, f)| !*f).count()); assert!((1..=TRACKER_FG_POINTS).contains(&fg) && (1..=TRACKER_BG_POINTS).contains(&bg), "{fg} fg, {bg} bg"); // a long, finely sampled stroke still gives at most the limit, and an empty one gives none let long = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 2.0, points: (0..400).map(|i| [i as f64 * 0.4, 50.0]).collect(), object: 0 }; - assert!(tracker_prompt(&[&long], 1.0).points.len() <= TRACKER_FG_POINTS); - assert!(tracker_prompt(&[], 1.0).points.is_empty()); + assert!(tracker_prompt(&[&long], 1.0, W as usize, H as usize).points.len() <= TRACKER_FG_POINTS); + assert!(tracker_prompt(&[], 1.0, W as usize, H as usize).points.is_empty()); +} + +/// Tracker clicks follow the same correction order as the classic matte and model prompts. +#[test] +fn tracker_prompts_follow_latest_corrections_and_keep_remaining_foreground() { + let (w, h) = (100, 80); + let green = Stroke { kind: StrokeKind::Fg, frame: 0, radius: 15.0, points: vec![[50.0, 40.0]], object: 0 }; + let red = Stroke { kind: StrokeKind::Bg, radius: 16.0, ..green.clone() }; + for strokes in [&[&green, &red][..], &[&red, &green][..]] { + let classic = FrameSeg::empty(w, h); + for prompt in [tracker_prompt(strokes, 1.0, w, h), tracker_prompt_boxed(strokes, 1.0, &classic)] { + let expected = std::ptr::eq(strokes[1], &green); + assert_eq!(prompt.points.iter().any(|(_, fg)| *fg), expected, "{:?}", prompt.points); + let hard = resolve(strokes, StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h); + for (point, fg) in &prompt.points { + assert_eq!(owner(&hard, w, h, *point), Some(if *fg { HARD_FG } else { HARD_BG })); + } + assert!(!prompt.points.is_empty()); + } + } + // A small correction removes the sampled centre, while most of the foreground dot remains. + let tap = Stroke { radius: 2.0, ..red }; + let strokes = [&green, &tap]; + let hard = resolve(&strokes, StrokeKind::Fg, StrokeKind::Bg, 1.0, w, h); + for prompt in [tracker_prompt(&strokes, 1.0, w, h), tracker_prompt_boxed(&strokes, 1.0, &FrameSeg::empty(w, h))] { + assert!(prompt.points.iter().any(|(_, fg)| *fg)); + assert!(prompt.points.iter().any(|(_, fg)| !*fg)); + for (point, fg) in &prompt.points { + assert_eq!(owner(&hard, w, h, *point), Some(if *fg { HARD_FG } else { HARD_BG })); + } + assert!(prompt.points.iter().filter(|(_, fg)| *fg).count() <= TRACKER_FG_POINTS); + assert!(prompt.points.iter().filter(|(_, fg)| !*fg).count() <= TRACKER_BG_POINTS); + } } #[test] diff --git a/crates/ui-egui/src/panels/roto_tool.rs b/crates/ui-egui/src/panels/roto_tool.rs index cfed82e9b..194bba26a 100644 --- a/crates/ui-egui/src/panels/roto_tool.rs +++ b/crates/ui-egui/src/panels/roto_tool.rs @@ -452,7 +452,7 @@ pub fn bar_buttons(app: &mut EffectcraftApp, ui: &mut egui::Ui, p: &egui::Painte let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(104.0, 20.0)); let strokes = fx::data(¶ms).strokes.iter().filter(|st| st.object == current && st.kind.is_segmentation()).count(); let label = if strokes == 0 { tr_args("Object {}", &[&(current + 1)]) } else { tr_args("Object {} ({})", &[&(current + 1), &strokes]) }; - if widgets::text_button(ui, r, &label, false, &t, egui::Id::new(("roto-object", uid))).clicked() { + if widgets::text_button(ui, r, &label, false, &t, egui::Id::unique(("roto-object", uid))).clicked() { // cycle through the objects that have strokes (and the empty one being drawn) let mut all = objects.clone(); if !known { @@ -466,7 +466,7 @@ pub fn bar_buttons(app: &mut EffectcraftApp, ui: &mut egui::Ui, p: &egui::Painte app.auto.add("layerPanel.roto.object", r, &label); x += 108.0; let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(26.0, 20.0)); - if widgets::text_button(ui, r, "+", false, &t, egui::Id::new(("roto-object-new", uid))).clicked() { + if widgets::text_button(ui, r, "+", false, &t, egui::Id::unique(("roto-object-new", uid))).clicked() { let next = objects.iter().copied().chain([current]).max().map_or(0, |m| m + 1); let _ = app.session.execute("roto.options", json!({"object": next})); } @@ -475,7 +475,7 @@ pub fn bar_buttons(app: &mut EffectcraftApp, ui: &mut egui::Ui, p: &egui::Painte // remove the current object's strokes (the object stays selected; with none left it is simply empty) if strokes > 0 { let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(26.0, 20.0)); - let c = widgets::text_button(ui, r, "×", false, &t, egui::Id::new(("roto-object-clear", uid))).on_hover_text(tr("Remove this object's strokes")); + let c = widgets::text_button(ui, r, "×", false, &t, egui::Id::unique(("roto-object-clear", uid))).on_hover_text(tr("Remove this object's strokes")); if c.clicked() && let Err(e) = app.session.execute("roto.clearStrokes", json!({"layer": layer.id.0, "effect": uid, "object": current})) { diff --git a/crates/ui-egui/tests/it/main.rs b/crates/ui-egui/tests/it/main.rs index 9ef4c4b54..e3a249122 100644 --- a/crates/ui-egui/tests/it/main.rs +++ b/crates/ui-egui/tests/it/main.rs @@ -72,6 +72,8 @@ mod ui_quick_offset; mod ui_ram_preview; mod ui_render_queue; mod ui_render_queue_keyboard; +mod ui_roto_objects; +mod ui_roto_span_confidence; mod ui_script_console; mod ui_scriptui; mod ui_scrolling; diff --git a/crates/ui-egui/tests/ui_roto_objects.rs b/crates/ui-egui/tests/it/ui_roto_objects.rs similarity index 99% rename from crates/ui-egui/tests/ui_roto_objects.rs rename to crates/ui-egui/tests/it/ui_roto_objects.rs index 252ecfb3d..cbf52d07d 100644 --- a/crates/ui-egui/tests/ui_roto_objects.rs +++ b/crates/ui-egui/tests/it/ui_roto_objects.rs @@ -59,6 +59,7 @@ fn count(image: &image::RgbaImage, want: [u8; 3]) -> usize { #[test] fn the_boundary_view_outlines_each_object_and_show_object_picks_one() { + let _gpu = crate::gpu_test::lock(); let mut s = Session::default(); s.execute("comp.new", json!({"name": "View", "width": 640, "height": 360, "duration": 4})).unwrap(); s.execute("layer.newSolid", json!({"name": "Box", "color": "#e04020", "width": 80, "height": 80})).unwrap(); diff --git a/crates/ui-egui/tests/ui_roto_span_confidence.rs b/crates/ui-egui/tests/it/ui_roto_span_confidence.rs similarity index 99% rename from crates/ui-egui/tests/ui_roto_span_confidence.rs rename to crates/ui-egui/tests/it/ui_roto_span_confidence.rs index 04379f1ae..d9730b778 100644 --- a/crates/ui-egui/tests/ui_roto_span_confidence.rs +++ b/crates/ui-egui/tests/it/ui_roto_span_confidence.rs @@ -64,6 +64,7 @@ fn close(got: [u8; 4], want: [u8; 3]) -> bool { #[test] fn the_span_bar_shows_the_trackers_confidence() { + let _gpu = crate::gpu_test::lock(); let mut s = Session::default(); s.execute("comp.new", json!({"name": "View", "width": 640, "height": 360, "duration": 4})).unwrap(); s.execute("layer.newSolid", json!({"name": "Box", "color": "#e04020", "width": 80, "height": 80})).unwrap(); diff --git a/docs/effects.md b/docs/effects.md index e233dd47f..3490af5a2 100644 --- a/docs/effects.md +++ b/docs/effects.md @@ -3,7 +3,7 @@ -EffectCraft ships 308 effects with 3034 parameters in total, grouped into the same categories as the Effects & Presets panel. Every effect is our own implementation, written from public behaviour descriptions and standard image-processing literature. Parameter names, order, popup options, units, defaults and ranges follow the reference application so that projects, expressions and muscle memory carry over. +EffectCraft ships 308 effects with 2987 parameters in total, grouped into the same categories as the Effects & Presets panel. Every effect is our own implementation, written from public behaviour descriptions and standard image-processing literature. Parameter names, order, popup options, units, defaults and ranges follow the reference application so that projects, expressions and muscle memory carry over. - **GPU**: 280 effects also run on the GPU compositor with identical results. - **32**: 308 effects process 32-bit float (HDR, overbright) pixels without clamping. From 76f82d7d552b7fe2fae2aaee6cf2eaf1ece1623f Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 30 Oct 2026 08:53:17 +0100 Subject: [PATCH 72/77] #1539: clarify fallback visibility and benchmark correction rounds --- apps/effectcraft-cli/src/main.rs | 2 +- crates/track/src/roto/bench.rs | 8 ++++---- docs/roto-brush-sam2.md | 13 ++++++++----- 3 files changed, 13 insertions(+), 10 deletions(-) diff --git a/apps/effectcraft-cli/src/main.rs b/apps/effectcraft-cli/src/main.rs index dc47671e2..69a56c696 100644 --- a/apps/effectcraft-cli/src/main.rs +++ b/apps/effectcraft-cli/src/main.rs @@ -1056,7 +1056,7 @@ fn bench_roto(args: &Args) -> Result<(), Failure> { if sam2.is_none() { note!("bench --roto: no SAM 2.1 size is installed; no SAM 2 column"); } - // `--fixes`: also count the correction strokes each needs before every frame reaches IoU 0.9 (up to 6). + // `--fixes`: count correction rounds before every frame reaches IoU 0.9 (up to 6 rounds, each up to two strokes). let fixes = args.flag("--fixes"); let measure = |sc: &bench::Scenario, run: &dyn Fn(&bench::Scenario) -> Option| if fixes { bench::with_fixes(sc, 0.9, 6, run) } else { run(sc) }; let mut runs = vec![]; diff --git a/crates/track/src/roto/bench.rs b/crates/track/src/roto/bench.rs index c086bfd57..98ada3903 100644 --- a/crates/track/src/roto/bench.rs +++ b/crates/track/src/roto/bench.rs @@ -46,7 +46,7 @@ pub struct Run { /// Mean IoU over the frames after the object's last absence (re-appearance), if it had one. pub after_return: Option, pub per_frame: Vec, - /// Correction strokes needed before every frame reaches the target IoU (`>N`: not within N); set by + /// Correction rounds needed before every frame reaches the target IoU (`>N`: not within N); set by /// [`with_fixes`], `None` when not measured. pub fixes: Option, /// The most memory (MB) the tracker held for its memory banks and cached frames (trackers that have none: `None`). @@ -301,7 +301,7 @@ pub fn run_tracker(sc: &Scenario, tracker: &str, model: &dyn MaskModel) -> Optio } /// The correction a user would paint on frame `f` where `pred` is wrong: a short foreground stroke through -/// the true object's centre, and a background stroke on the biggest false-positive blob's centre. +/// the true object's centre, and a background stroke at the centroid of all false-positive pixels. fn correction(sc: &Scenario, f: usize, pred: &[u8]) -> Vec { let truth = sc.truth.get(f).map(Vec::as_slice).unwrap_or(&[]); let centre = |keep: &dyn Fn(usize) -> bool| -> Option<[f64; 2]> { @@ -324,8 +324,8 @@ fn correction(sc: &Scenario, f: usize, pred: &[u8]) -> Vec { } /// Repeat `run` on `sc`, each time painting the correction a user would on the first frame under `target` IoU, -/// until every frame reaches it or `max` corrections were needed. Sets `fixes` on the first run's result: the -/// number of corrections, or `>max`. +/// until every frame reaches it or `max` correction rounds were needed. Sets `fixes` on the first run's result: +/// the number of rounds (each can add foreground and background strokes), or `>max`. pub fn with_fixes(sc: &Scenario, target: f64, max: usize, run: &dyn Fn(&Scenario) -> Option) -> Option { let mut cur = sc.clone(); let mut first: Option = None; diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index 7078fa103..b165a29bc 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -126,8 +126,10 @@ the CPU. Pending per-object probability maps have a separate 256 MiB budget, including their allocated capacity and frame slots. This bounds the full-resolution results waiting to be united when objects first prompted on different frames -report in separate passes. Reaching the budget cancels tracking and reports the reason through the existing -`trackerNote`; propagation falls back to the classic engine. A single object's results need no pending buffer. +report in separate passes. Reaching the budget cancels tracking and falls back to the classic engine. The reason is +available as `trackerNote` in live job progress only. Completion removes that progress and shows a success toast, so +`wait: true` callers can miss the fallback reason; the completion toast alone does not confirm the tracker was used. +A single object's results need no pending buffer. ## How it works `crates/sam2` is a Rust port (on [candle](https://github.com/huggingface/candle), native only) of the Hugging Face @@ -189,9 +191,10 @@ loses it when it is hidden; SAM 2's streaming memory is what lets it re-acquire textured disc), so the "move" and "light" rows mostly show all three working; the last two rows are the difference. ### Corrections needed (`--fixes`) -`bench --roto --fixes` also counts the correction strokes each tracker needs before every frame reaches IoU 0.9: after each -run it paints what a user would on the first frame under the target (a short foreground stroke through the true object's -centre, and a background stroke on the biggest false-positive blob), up to 6. 16 frames, debug build: +`bench --roto --fixes` also counts the correction rounds each tracker needs before every frame reaches IoU 0.9: after each +run it paints on the first frame under the target, up to 6 rounds. A round can add a short foreground stroke through the +true object's centre and a background stroke at the centroid of all false-positive pixels. With separate false-positive +blobs, that centroid can fall between them. The counts below are rounds, not individual strokes. 16 frames, debug build: | Clip | Classical | MobileSAM | SAM 2.1 Tiny | |---|---|---|---| From ee6830fefb125f1fba4c1bd85632e9c0783ebd25 Mon Sep 17 00:00:00 2001 From: bflatastic Date: Fri, 30 Oct 2026 15:01:40 +0700 Subject: [PATCH 73/77] sam2: integration tests in one binary (tests/it/main.rs) Co-Authored-By: Claude Opus 5.5 --- crates/sam2/tests/{ => it}/figure.rs | 0 crates/sam2/tests/it/main.rs | 4 ++++ crates/sam2/tests/{ => it}/reference.rs | 0 docs/roto-brush-sam2.md | 2 +- 4 files changed, 5 insertions(+), 1 deletion(-) rename crates/sam2/tests/{ => it}/figure.rs (100%) create mode 100644 crates/sam2/tests/it/main.rs rename crates/sam2/tests/{ => it}/reference.rs (100%) diff --git a/crates/sam2/tests/figure.rs b/crates/sam2/tests/it/figure.rs similarity index 100% rename from crates/sam2/tests/figure.rs rename to crates/sam2/tests/it/figure.rs diff --git a/crates/sam2/tests/it/main.rs b/crates/sam2/tests/it/main.rs new file mode 100644 index 000000000..b8232c63f --- /dev/null +++ b/crates/sam2/tests/it/main.rs @@ -0,0 +1,4 @@ +//! The crate's integration tests, one binary (each module is a former `tests/.rs`). + +mod figure; +mod reference; diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/it/reference.rs similarity index 100% rename from crates/sam2/tests/reference.rs rename to crates/sam2/tests/it/reference.rs diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index b165a29bc..f81e854c3 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -134,7 +134,7 @@ A single object's results need no pending buffer. ## How it works `crates/sam2` is a Rust port (on [candle](https://github.com/huggingface/candle), native only) of the Hugging Face `transformers` SAM 2 implementation, Apache-2.0 (see its `NOTICE`). Every stage is checked against PyTorch by -`crates/sam2/tests/reference.rs` (they skip without the weights; reference tensors come from +`crates/sam2/tests/it/reference.rs` (they skip without the weights; reference tensors come from `tools/sam2_reference.py` and `tools/sam2_video_reference.py`): | Stage | Max abs error vs PyTorch | From efb53c31fc2be976507bf946c075f196e1cc4896 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 30 Oct 2026 09:34:23 +0100 Subject: [PATCH 74/77] #1539: clarify saved model identity and runtime fallback warnings --- docs/roto-brush-sam2.md | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index b165a29bc..d5e5b26f0 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -22,16 +22,18 @@ Apache-2.0). It is optional: nothing changes until you install the model and cho chose and whether it is installed, loaded and loading; `notice` says when it is not the one in use for the layer. `tracker.saved` names the model the saved frames came from. -Propagate never falls back to the classic engine without saying so. With SAM 2 chosen but not installed, or failing to -load, it refuses with the reason. While it is still loading, `wait: true` (agents) loads it first; without `wait` it +With SAM 2 chosen but not installed, or failing to load, Propagate refuses with the reason. A tracker failure during +propagation can fall back to classic; its warning is live progress only, as described under [Memory and limits](#memory-and-limits). +While it is still loading, `wait: true` (agents) loads it first; without `wait` it refuses, and the load goes on in the background. Deselecting SAM 2 (choosing Classic) runs Propagate exactly as before, and Roto Brush 1.0 never uses a model. Tracked frames are saved in the project with the instance, as Freeze's are, so a saved project renders the same when it reopens, in the app or with `effectcraft-cli`, without the model being loaded or chosen there. They apply while the strokes and span they came from are unchanged. Changing a stroke or the span makes them stale until the next Propagate. -Choosing another model, or the classic engine, leaves them in place and `roto.status` says so; a classic Propagate -replaces them with classic mattes. Object Matte saves nothing. +Choosing another model, or the classic engine, leaves them in place and `roto.status` says so. They retain the original +model tag and per-object cache keys when rendered or re-saved; the next Propagate replaces them using the selected +model (or classic engine). Object Matte saves nothing. ## Several objects in one instance A Roto Brush instance can hold several objects. The **Object N** button in the Layer panel's bar picks the object new From a4f5c516ed8c08e82f75238eb1c744033796233d Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 30 Oct 2026 10:00:33 +0100 Subject: [PATCH 75/77] #1539: validate benchmark inputs and preserve shared refinement strokes --- apps/effectcraft-cli/src/main.rs | 22 +++++++----- apps/effectcraft-cli/tests/it/cli.rs | 54 ++++++++++++++++++++++++---- crates/effects/src/roto.rs | 27 ++++++++++++++ crates/engine/src/cmd/roto_cmds.rs | 15 ++++++-- crates/engine/src/tests/roto.rs | 34 ++++++++++++++++++ crates/sam2/src/session.rs | 5 +-- 6 files changed, 139 insertions(+), 18 deletions(-) diff --git a/apps/effectcraft-cli/src/main.rs b/apps/effectcraft-cli/src/main.rs index 69a56c696..b1b53b488 100644 --- a/apps/effectcraft-cli/src/main.rs +++ b/apps/effectcraft-cli/src/main.rs @@ -1035,7 +1035,7 @@ fn bench_weights( fn bench_roto(args: &Args) -> Result<(), Failure> { use effectcraft_segment::{Loaded, MOBILE_SAM, MaskModel}; use effectcraft_track::roto::bench; - let frames = args.num("--frames")?.unwrap_or(24.0).clamp(12.0, 240.0) as usize; + let frames = args.num("--frames")?.unwrap_or(24.0).clamp(1.0, 240.0) as usize; let session = effectcraft_host::session(); // `--mobilesam FILE` uses the weights at FILE (a file that cannot be used is an error); // otherwise the installed ones. @@ -1060,13 +1060,14 @@ fn bench_roto(args: &Args) -> Result<(), Failure> { let fixes = args.flag("--fixes"); let measure = |sc: &bench::Scenario, run: &dyn Fn(&bench::Scenario) -> Option| if fixes { bench::with_fixes(sc, 0.9, 6, run) } else { run(sc) }; let mut runs = vec![]; - let mut scenarios = bench::scenarios(frames); - // `--clip FRAMES_DIR --truth MASKS_DIR`: the user's own clip (image sequences, same file order) with its true mattes. - if let Some(dir) = args.opt("--clip") { + // User clips pair unique filename stems, allowing JPEG frames with PNG truth masks. + let scenarios = if let Some(dir) = args.opt("--clip") { let truth = args.opt("--truth").ok_or_else(|| Failure::Error("--clip needs --truth MASKS_DIR".into()))?; let name = std::path::Path::new(dir).file_name().map_or_else(|| "clip".to_string(), |n| n.to_string_lossy().chars().take(10).collect()); - scenarios = vec![load_bench_clip(&name, dir, truth, frames).map_err(Failure::Error)?]; - } + vec![load_bench_clip(&name, dir, truth, frames).map_err(Failure::Error)?] + } else { + bench::scenarios(frames) + }; for sc in scenarios { runs.extend(measure(&sc, &|s| Some(bench::run_flow(s, "classical", None)))); if let Some(m) = &mobilesam { @@ -1085,14 +1086,14 @@ fn bench_roto(args: &Args) -> Result<(), Failure> { Ok(()) } -/// The image files of `dir` (png, jpg, jpeg), in file-name order. +/// The image files of `dir` (png, jpg, jpeg), ordered by stem then full file name. fn image_files(dir: &str) -> Result, String> { let mut files: Vec = std::fs::read_dir(dir) .map_err(|e| format!("cannot read {dir}: {e}"))? .filter_map(|e| e.ok().map(|e| e.path())) .filter(|p| p.extension().and_then(|e| e.to_str()).is_some_and(|e| ["png", "jpg", "jpeg"].contains(&e.to_ascii_lowercase().as_str()))) .collect(); - files.sort(); + files.sort_by(|a, b| a.file_stem().cmp(&b.file_stem()).then_with(|| a.file_name().cmp(&b.file_name()))); Ok(files) } @@ -1103,6 +1104,11 @@ fn load_bench_clip(name: &str, frames_dir: &str, truth_dir: &str, limit: usize) if frames.is_empty() || frames.len() != masks.len() { return Err(format!("{} frames in {frames_dir} but {} masks in {truth_dir}: they must match one to one", frames.len(), masks.len())); } + if frames.iter().map(|p| p.file_stem()).ne(masks.iter().map(|p| p.file_stem())) + || frames.windows(2).chain(masks.windows(2)).any(|pair| pair.first().and_then(|p| p.file_stem()) == pair.get(1).and_then(|p| p.file_stem())) + { + return Err("frame and truth file names must match one to one, with unique stems (extensions may differ)".into()); + } let n = frames.len().min(limit.clamp(1, 240)); let (mut imgs, mut truth) = (vec![], vec![]); for (f, m) in frames.iter().zip(&masks).take(n) { diff --git a/apps/effectcraft-cli/tests/it/cli.rs b/apps/effectcraft-cli/tests/it/cli.rs index 93c899839..b3279260a 100644 --- a/apps/effectcraft-cli/tests/it/cli.rs +++ b/apps/effectcraft-cli/tests/it/cli.rs @@ -689,7 +689,7 @@ fn bench_roto_runs_on_a_users_clip() { std::fs::create_dir_all(d).unwrap(); } let (w, h) = (96u32, 72u32); - for f in 0..6u32 { + for f in 0..14u32 { let (cx, cy) = (30.0 + 4.0 * f as f32, 36.0f32); let (mut img, mut mask) = (image::RgbImage::new(w, h), image::GrayImage::new(w, h)); for y in 0..h { @@ -700,15 +700,23 @@ fn bench_roto_runs_on_a_users_clip() { mask.put_pixel(x, y, image::Luma([if inside { 255 } else { 0 }])); } } - img.save(frames.join(format!("f{f:03}.png"))).unwrap(); - mask.save(masks.join(format!("f{f:03}.png"))).unwrap(); + let extension = if f == 0 { "jpg" } else { "png" }; + let stem = match f { + 0 => "a".into(), + 1 => "a.jpg".into(), + _ => format!("f{f:03}"), + }; + img.save(frames.join(format!("{stem}.{extension}"))).unwrap(); + mask.save(masks.join(format!("{stem}.png"))).unwrap(); } let (f, m) = (frames.to_str().unwrap(), masks.to_str().unwrap()); - let out = bin().env("EFFECTCRAFT_CONFIG_DIR", &config).args(["bench", "--roto", "--clip", f, "--truth", m, "--frames", "6"]).output().unwrap(); + let out = bin().env("EFFECTCRAFT_CONFIG_DIR", &config).args(["bench", "--roto", "--clip", f, "--truth", m, "--frames", "6", "--json"]).output().unwrap(); let text = String::from_utf8_lossy(&out.stdout).to_string(); assert!(out.status.success(), "{text}\n{}", String::from_utf8_lossy(&out.stderr)); - let row = text.lines().find(|l| l.contains("classical")).unwrap_or_else(|| panic!("no classical row: {text}")); - let mean: f64 = row.split_whitespace().nth(2).unwrap().parse().unwrap(); + let rows: Value = serde_json::from_str(&text).unwrap(); + let row = rows.as_array().unwrap().iter().find(|row| row["tracker"] == "classical").unwrap(); + assert_eq!(row["per_frame"].as_array().unwrap().len(), 6, "a user clip respects the requested frame limit"); + let mean = row["mean_iou"].as_f64().unwrap(); assert!(mean > 0.8, "classical mean IoU on the user's clip: {mean}\n{text}"); // one mask short: refused std::fs::remove_file(masks.join("f005.png")).unwrap(); @@ -720,6 +728,40 @@ fn bench_roto_runs_on_a_users_clip() { assert!(!out.status.success()); } +#[test] +fn bench_roto_rejects_mismatched_names_before_loading_images() { + let root = tmp("bench-clip-names"); + let (frames, masks) = (root.join("frames"), root.join("masks")); + let _ = std::fs::remove_dir_all(&root); + for dir in [&frames, &masks] { + std::fs::create_dir_all(dir).unwrap(); + std::fs::write(dir.join("f000.png"), []).unwrap(); + } + std::fs::write(frames.join("f001.png"), []).unwrap(); + std::fs::write(masks.join("f002.png"), []).unwrap(); + let out = bin() + .env("EFFECTCRAFT_CONFIG_DIR", root.join("config")) + .args(["bench", "--roto", "--frames", "1", "--clip"]) + .arg(&frames) + .arg("--truth") + .arg(&masks) + .output() + .unwrap(); + let error = String::from_utf8_lossy(&out.stderr); + assert!(!out.status.success()); + assert!(error.contains("names must match"), "validate the entire name set before decoding: {error}"); + std::fs::remove_file(frames.join("f001.png")).unwrap(); + std::fs::remove_file(masks.join("f002.png")).unwrap(); + for dir in [&frames, &masks] { + std::fs::write(dir.join("f000.jpg"), []).unwrap(); + } + let out = + bin().env("EFFECTCRAFT_CONFIG_DIR", root.join("config")).args(["bench", "--roto", "--clip"]).arg(&frames).arg("--truth").arg(&masks).output().unwrap(); + assert!(!out.status.success()); + assert!(String::from_utf8_lossy(&out.stderr).contains("unique stems")); + std::fs::remove_dir_all(root).unwrap(); +} + /// `bench --roto --sam2 FILE` / `--mobilesam FILE` say why a file named on the command line cannot be /// used (a missing file, one that is not the published weights), instead of "not installed". #[test] diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 1a1ce105f..379c772c5 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -694,6 +694,9 @@ impl Chain { let mut prompts = vec![]; for f in self.data.stroke_frames().into_iter().filter(|f| self.data.in_span(*f)) { let strokes = self.data.strokes_of(f, obj); + if !strokes.iter().any(|s| s.kind == StrokeKind::Fg) { + continue; + } let p = match frame(f) { Some(img) => rb::tracker_prompt_boxed(&strokes, self.opts.scale, &rb::segment(&img, &strokes, None, &self.opts, self.refine_radius)), None => return Err(format!("cannot read Roto Brush prompt frame {f}")), @@ -1429,6 +1432,30 @@ mod tests { } } + #[test] + fn prompt_segmentation_reads_only_frames_with_that_objects_foreground() { + let mut data = RotoData::default(); + for (kind, frame, object) in [(StrokeKind::Fg, 0, 0), (StrokeKind::Refine, 1, 0), (StrokeKind::Fg, 2, 1)] { + data.add(Stroke { kind, frame, radius: 3.0, points: vec![[35.0, 40.0]], object }, 2, [0, 2]); + } + data.span = [0, 2]; + let mut params = params_with(&data); + params.values.insert(INPUT_KEY.into(), Value::Str("foreground-only-prompts".into())); + let mut chain = Chain::new(¶ms, [100.0, 80.0], 25.0, 1.0); + let fake = Arc::new(TwoObjects(Default::default())); + chain.model = Some(fake.clone()); + let mut prompt_reads = vec![]; + let mut get = |f: i64| { + if fake.0.lock().unwrap().is_empty() { + prompt_reads.push(f); + } + Some(Arc::new(frame(f as f64 / 25.0))) + }; + assert_eq!(chain.track(&mut get, &|| false, &mut |_| {}).unwrap().unwrap().frames, 3); + assert_eq!(prompt_reads, [0, 2], "shared Refine Edge and other objects' frames need no prompt graph cut"); + assert_eq!(fake.0.lock().unwrap()[0], [1, 1]); + } + #[test] fn separate_object_passes_cancel_at_the_pending_mask_budget() { use std::sync::atomic::Ordering::Relaxed; diff --git a/crates/engine/src/cmd/roto_cmds.rs b/crates/engine/src/cmd/roto_cmds.rs index 5124d413d..792c60019 100644 --- a/crates/engine/src/cmd/roto_cmds.rs +++ b/crates/engine/src/cmd/roto_cmds.rs @@ -378,9 +378,20 @@ fn clear_strokes(s: &mut Session, p: &Value) -> Result { Some(k) => Some(StrokeKind::from_name(k).ok_or_else(|| bad(cmd, format!("unknown kind `{k}`")))?), None => None, }; - let object = f_p(p, "object").map(|v| v.round() as u32); + let object = match p.get("object") { + None => None, + Some(value) => { + let value = value + .as_f64() + .filter(|v| v.is_finite() && (0.0..=rt::MAX_OBJECT_ID as f64).contains(v)) + .ok_or_else(|| bad(cmd, format!("`object` is 0–{}", rt::MAX_OBJECT_ID)))?; + Some(value.round() as u32) + } + }; let n0 = d.strokes.len(); - d.strokes.retain(|st| !(frame.is_none_or(|f| st.frame == f) && kind.is_none_or(|k| st.kind == k) && object.is_none_or(|o| st.object == o))); + d.strokes.retain(|st| { + !(frame.is_none_or(|f| st.frame == f) && kind.is_none_or(|k| st.kind == k) && object.is_none_or(|o| st.kind.is_segmentation() && st.object == o)) + }); if d.strokes.iter().all(|s| !s.kind.is_segmentation()) { d.base = None; } else if d.base.is_some_and(|b| !d.strokes.iter().any(|s| s.frame == b && s.kind.is_segmentation())) { diff --git a/crates/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index e686962e1..99c155265 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -1125,3 +1125,37 @@ fn strokes_belong_to_the_current_object() { let json = s.project.to_json(); assert!(json.contains("\\\"object\\\":31"), "the project file carries the object"); } + +#[test] +fn clearing_one_object_keeps_shared_refine_edge_strokes() { + let mut s = Session::default(); + s.execute("comp.new", json!({"width": 64, "height": 64, "duration": 1})).unwrap(); + let layer = s.execute("layer.newSolid", json!({})).unwrap()["layer"].as_u64().unwrap(); + for (object, kind) in [(0, "fg"), (0, "bg"), (0, "refine"), (0, "refineErase"), (1, "fg")] { + s.execute("roto.stroke", json!({"layer": layer, "object": object, "kind": kind, "points": [[20, 20]], "radius": 4})).unwrap(); + } + let before = s.project.to_json(); + let result = s.execute("roto.clearStrokes", json!({"layer": layer, "object": 0})).unwrap(); + assert_eq!(result["removed"], 2, "only the object's foreground/background strokes are owned by it"); + assert_eq!(result["strokes"], 3); + assert!(s.undo()); + assert_eq!(s.project.to_json(), before); + for kind in ["refine", "refineErase"] { + let result = s.execute("roto.clearStrokes", json!({"layer": layer, "kind": kind})).unwrap(); + assert_eq!(result["removed"], 1, "explicitly clearing shared {kind} still works"); + } +} + +#[test] +fn clearing_strokes_rejects_invalid_object_filters_without_editing() { + let mut s = Session::default(); + s.execute("comp.new", json!({"width": 64, "height": 64, "duration": 1})).unwrap(); + let layer = s.execute("layer.newSolid", json!({})).unwrap()["layer"].as_u64().unwrap(); + s.execute("roto.stroke", json!({"layer": layer, "points": [[20, 20]], "radius": 4})).unwrap(); + let before = s.project.to_json(); + for object in [json!(-1), json!(-0.1), json!(32), json!(1e30), json!(null), json!("0")] { + let result = s.execute("roto.clearStrokes", json!({"layer": layer, "object": object})); + assert!(result.is_err(), "invalid object {object} must not delete object zero: {result:?}"); + assert_eq!(s.project.to_json(), before); + } +} diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 48ec72a31..0524ee89e 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -96,8 +96,8 @@ pub enum Backend { Auto, } -/// The GPU path: kernels for the backbone and the memory attention. Any failure switches it off for good -/// and the CPU path takes over. +/// The GPU path: kernels for the backbone and memory attention. Device failures disable it; +/// memory-attention budget limits fall back to the CPU for that call only. struct GpuParts { encoder: GpuEncoder, attention: GpuMemoryAttention, @@ -229,6 +229,7 @@ impl VideoModel { if let Some(g) = self.gpu.as_ref().filter(|g| g.ok.load(Ordering::Relaxed)) { match g.attention.forward(current, current_pos, memory, memory_pos, pointers) { Ok(t) => return Ok(t), + Err(Error::Budget(reason)) => log::info!("SAM 2 memory attention on the CPU for this call: {reason}"), Err(e) => self.gpu_failed("memory attention", &e), } } From 4f642085877aad76913d471de94f8e69d2edda2c Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 30 Oct 2026 10:17:19 +0100 Subject: [PATCH 76/77] #1539: keep per-object mattes when only one object is followed --- crates/effects/src/roto.rs | 43 +++++++++++++++++++++++++++++++++++--- 1 file changed, 40 insertions(+), 3 deletions(-) diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 379c772c5..8b021eb87 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -689,6 +689,7 @@ impl Chain { // Each object has its own prompts: a few clicks from its strokes and a box around what the classic engine // makes of them on that frame (clicks alone pick a part of the object, not all of it). let objects = self.data.objects(); + let several = objects.len() > 1; let mut by_object: Vec> = vec![]; for &obj in &objects { let mut prompts = vec![]; @@ -751,7 +752,6 @@ impl Chain { // a frame is occluded when every object is absent; otherwise the doubtful object sets the confidence let (occluded, confidence) = (!present, if present { confidence } else { worst }); // with several objects each keeps its own matte too (Show Object, the per-object overlay) - let several = objects.len() > 1; let mut parts = Vec::with_capacity(segs.len()); for (obj, seg) in segs { if several { @@ -803,8 +803,8 @@ impl Chain { } } -/// Each object's matte of layer frame `f` where the tracker followed several objects (full resolution; computes -/// nothing). Show Object limits it to the one chosen. Empty with fewer than two objects. +/// Each tracked object's matte of layer frame `f` in a multi-object instance (full resolution; computes +/// nothing). Show Object limits it to the one chosen. Empty with fewer than two objects in the instance. pub fn object_mattes(params: &Params, layer_size: [f64; 2], fps: f64, f: i64) -> Vec<(u32, Arc)> { let mut chain = Chain::new(params, layer_size, fps, 1.0); let objects = chain.data.objects(); @@ -1432,6 +1432,43 @@ mod tests { } } + #[test] + fn a_single_followed_object_keeps_its_matte_in_a_multi_object_instance() { + for outside_span in [false, true] { + let mut data = RotoData::default(); + data.add(Stroke { kind: StrokeKind::Fg, frame: 0, radius: 3.0, points: vec![[35.0, 40.0]], object: 0 }, 2, [0, 5]); + data.add(Stroke { kind: StrokeKind::Fg, frame: if outside_span { 5 } else { 0 }, radius: 3.0, points: vec![[85.0, 70.0]], object: 1 }, 2, [0, 5]); + if !outside_span { + data.add(Stroke { kind: StrokeKind::Bg, frame: 0, radius: 8.0, points: vec![[85.0, 70.0]], object: 1 }, 2, [0, 5]); + } + data.span = [0, 2]; + assert_eq!(data.objects(), [0, 1]); + let mut params = params_with(&data); + params.values.insert(INPUT_KEY.into(), Value::Str(format!("one-followed-object-{outside_span}"))); + let mut chain = Chain::new(¶ms, [100.0, 80.0], 25.0, 1.0); + chain.model = Some(Arc::new(TwoObjects(Default::default()))); + chain.keys = data.chain_keys(seed_with(¶ms, [100.0, 80.0], 25.0, effectcraft_segment::SAM2_TINY.id)); + let mut get = |f: i64| Some(Arc::new(frame(f as f64 / 25.0))); + assert_eq!(chain.track(&mut get, &|| false, &mut |_| {}).unwrap().unwrap().frames, 3); + let matte = chain.object_seg(1, 0).expect("the followed object keeps its own matte"); + assert!(matte.area() > 0); + assert!(chain.object_seg(1, 1).is_none(), "the unprompted object has no invented matte"); + chain.show_object = 1; + assert_eq!(chain.ensure(1, &mut get, &|| false).unwrap().matte, matte.matte); + let saved = tracks_of(&chain).unwrap(); + assert_eq!(saved.objects.get(&0).unwrap().len(), 3); + assert!(!saved.objects.contains_key(&1)); + params.values.insert(TRACKED.into(), Value::Str(saved.to_json())); + params.values.insert("rotoBrushMatte/object".into(), Value::Scalar(1.0)); + let reopened = Chain::new(¶ms, [100.0, 80.0], 25.0, 1.0); + assert_eq!(reopened.object_seg(1, 0).unwrap().matte, matte.matte); + assert_eq!(reopened.ensure(1, &mut |_| None, &|| false).unwrap().matte, matte.matte); + let overlays = object_mattes(¶ms, [100.0, 80.0], 25.0, 1); + assert_eq!(overlays.len(), 1); + assert_eq!((overlays[0].0, overlays[0].1.area()), (0, matte.area())); + } + } + #[test] fn prompt_segmentation_reads_only_frames_with_that_objects_foreground() { let mut data = RotoData::default(); From 5eaed8b2182cedf5534b9436bea47de046e964ac Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Fri, 30 Oct 2026 10:20:55 +0100 Subject: [PATCH 77/77] #1539: reopen saved object mattes with their cache evicted --- crates/effects/src/roto.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/effects/src/roto.rs b/crates/effects/src/roto.rs index 8b021eb87..d6f27b9bf 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -1460,6 +1460,7 @@ mod tests { assert!(!saved.objects.contains_key(&1)); params.values.insert(TRACKED.into(), Value::Str(saved.to_json())); params.values.insert("rotoBrushMatte/object".into(), Value::Scalar(1.0)); + forget(chain.keys.values().map(|key| object_key(*key, 0))); let reopened = Chain::new(¶ms, [100.0, 80.0], 25.0, 1.0); assert_eq!(reopened.object_seg(1, 0).unwrap().matte, matte.matte); assert_eq!(reopened.ensure(1, &mut |_| None, &|| false).unwrap().matte, matte.matte);