diff --git a/Cargo.lock b/Cargo.lock index 091a6ebf8..896d08dc2 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" @@ -1394,7 +1455,10 @@ dependencies = [ "effectcraft-gpu", "effectcraft-host", "effectcraft-render", + "effectcraft-sam2", + "effectcraft-segment", "effectcraft-time", + "effectcraft-track", "effectcraft-ui-egui", "image", "serde_json", @@ -1445,6 +1509,7 @@ dependencies = [ "effectcraft-psd", "effectcraft-raster", "effectcraft-render", + "effectcraft-sam2", "effectcraft-segment", "effectcraft-svg", "effectcraft-text", @@ -1764,6 +1829,21 @@ dependencies = [ "web-time", ] +[[package]] +name = "effectcraft-sam2" +version = "0.18.0" +dependencies = [ + "bytemuck", + "candle-core", + "candle-nn", + "effectcraft-segment", + "half", + "log", + "pollster 1.0.1", + "rayon", + "wgpu", +] + [[package]] name = "effectcraft-script" version = "0.18.0" @@ -1852,6 +1932,7 @@ dependencies = [ "effectcraft-path", "effectcraft-project", "effectcraft-render", + "effectcraft-segment", "effectcraft-text", "effectcraft-time", "eframe", @@ -2060,6 +2141,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" @@ -2585,6 +2678,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" @@ -2727,6 +2832,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" @@ -2780,7 +3004,7 @@ dependencies = [ "js-sys", "libc", "r-efi 6.0.0", - "rand_core", + "rand_core 0.10.1", "wasm-bindgen", ] @@ -2928,9 +3152,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", ] @@ -2970,6 +3197,8 @@ dependencies = [ "allocator-api2", "equivalent", "foldhash 0.2.0", + "serde", + "serde_core", ] [[package]] @@ -3605,6 +3834,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" dependencies = [ "libc", + "stable_deref_trait", ] [[package]] @@ -3856,6 +4086,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" @@ -4295,7 +4535,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]] @@ -4564,6 +4804,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" @@ -4606,7 +4855,7 @@ dependencies = [ "chacha20", "core_detect", "num-traits", - "rand", + "rand 0.10.3", "rand_xorshift", "regex-syntax", "rusty-fork", @@ -4694,6 +4943,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" @@ -4702,7 +4961,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]] @@ -4711,13 +4989,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]] @@ -5039,6 +5327,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" @@ -5114,6 +5413,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" @@ -5641,6 +5946,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" @@ -5932,6 +6251,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" @@ -7470,6 +7795,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" @@ -7482,7 +7819,7 @@ dependencies = [ "indexmap", "memchr", "nix", - "typed-path", + "typed-path 0.13.0", "windows-sys 0.61.2", "zopfli", ] diff --git a/Cargo.toml b/Cargo.toml index e105557ed..269638bad 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/apps/effectcraft-cli/Cargo.toml b/apps/effectcraft-cli/Cargo.toml index b5784f4eb..abfc014eb 100644 --- a/apps/effectcraft-cli/Cargo.toml +++ b/apps/effectcraft-cli/Cargo.toml @@ -20,6 +20,9 @@ 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 } +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 816186bed..6ba352247 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] [--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 (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 //! @@ -124,6 +126,9 @@ 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})), @@ -997,6 +1012,117 @@ 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(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. + 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 (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> = + 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"); + } + // `--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![]; + // 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()); + 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 { + runs.extend(measure(&sc, &|s| Some(bench::run_flow(s, "mobilesam", Some(m.as_ref()))))); + } + if let Some(m) = &sam2 { + 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") { + say!("{}", serde_json::to_string_pretty(&runs).map_err(|e| Failure::Error(e.to_string()))?); + } else { + say!("{}", bench::table(&runs)); + } + Ok(()) +} + +/// 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_by(|a, b| a.file_stem().cmp(&b.file_stem()).then_with(|| a.file_name().cmp(&b.file_name()))); + 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())); + } + 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) { + 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 c9d601ee9..e9a96dd7e 100644 --- a/apps/effectcraft-cli/tests/it/cli.rs +++ b/apps/effectcraft-cli/tests/it/cli.rs @@ -739,6 +739,113 @@ 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")); + let config = frames.join("config"); + 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..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 { + 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 }])); + } + } + 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", "--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 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(); + 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().env("EFFECTCRAFT_CONFIG_DIR", &config).args(["bench", "--roto", "--clip", f]).output().unwrap(); + 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] +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().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().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}"); + } + 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/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 6804ea2c7..8a9750282 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 => {} } @@ -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, 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 2467a719d..d6f27b9bf 100644 --- a/crates/effects/src/roto.rs +++ b/crates/effects/src/roto.rs @@ -29,9 +29,9 @@ 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 effectcraft_track::roto::{self as rb, Gmm, RotoData, SegOpts, rle}; use rayon::prelude::*; use serde::{Deserialize, Serialize}; @@ -40,12 +40,16 @@ 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. 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; @@ -78,6 +82,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)), @@ -95,6 +100,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, @@ -191,9 +197,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() @@ -208,6 +220,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 [ @@ -221,7 +238,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 } @@ -298,9 +315,44 @@ 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 +} + +/// 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 { + /// 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 +365,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 +386,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, @@ -341,13 +404,113 @@ 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, + /// Show Object: the matte of this object only (its id plus one; 0 shows all of them together). + pub show_object: u32, } 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); + 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. + 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))) + } + + /// 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; + // 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); + } + let full = match cached(ok, 1.0) { + Some(s) => s, + None => { + let (t, _) = saved?; + 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); 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 Some(seg) = self.object_seg(f, o) + { + return Some(seg); + } + 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. + 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 @@ -356,12 +519,24 @@ 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) { + return Some(s); + } let k = *self.keys.get(&g)?; + 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 +582,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); @@ -422,10 +598,222 @@ 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| 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())) + || 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()) + } + + /// 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; + } + } + + /// 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 + /// 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> { + 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), + 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()) }; + // 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| 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 { + return Ok(Some(TrackSummary::default())); + }; + // 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![]; + 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}")), + }; + if p.points.iter().any(|(_, fg)| *fg) { + prompts.push((f, p)); + } + } + by_object.push(prompts); + } + // 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: 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 = 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); + 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((*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; + 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 }); + // with several objects each keeps its own matte too (Show Object, the per-object overlay) + 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); + summary.min_confidence = if summary.frames == 1 { confidence } else { summary.min_confidence.min(confidence) }; + progress(summary.frames); + }; + 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() || pending_failed.get(), + &mut |o, g, 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; + } + 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 pending.into_incomplete() { + finish(g, parts, &mut summary, progress); + } + Ok(Some(summary)) + } +} + +/// 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(); + 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. @@ -434,7 +822,7 @@ pub fn downscale(s: &FrameSeg, scale: f64) -> FrameSeg { 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(); } @@ -513,6 +901,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}") } @@ -544,6 +935,183 @@ 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, + /// 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 +/// 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)?; + 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() }) +} + +/// 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)?; + 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() }) +} + +/// 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; + // 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 { + 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(); + 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) }); + } 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()); + } + } + } + } + let cache = TrackCache { key: tracks_key(&ch.keys), model: model.to_string(), size: size?, frames, objects }; + (!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(); @@ -731,8 +1299,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| { @@ -776,6 +1344,441 @@ 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(()) + } + } + + /// 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 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)); + 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); + 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(); + 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; + 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(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"); + } + + /// 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"); + 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] + 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]); + 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(); + 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(); + 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(); + 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"); + + // 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 /// blur / decontamination their nested ones) load with their values in the new places. #[test] 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/crates/engine/Cargo.toml b/crates/engine/Cargo.toml index 8e0f23711..826346987 100644 --- a/crates/engine/Cargo.toml +++ b/crates/engine/Cargo.toml @@ -40,7 +40,8 @@ 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 } tempfile = { workspace = true } +effectcraft-sam2 = { workspace = true } diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index b4ddab00d..63550545a 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}; @@ -27,13 +29,21 @@ 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 /// [`Session::poll_models`]). pub done: Option<(String, bool)>, + /// Counts selection and selected-weight changes: an older load 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. @@ -61,7 +71,14 @@ impl Models { self.face_model.lock().ok().and_then(|m| m.clone()) } - fn set_active(&self, task: Task, model: Option) { + /// 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) + } + + pub(crate) fn set_active(&self, task: Task, model: Option) { activate(self.lane(task), &self.face_model, task, model); } } @@ -90,6 +107,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); @@ -126,6 +167,13 @@ 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); + impl Session { /// The folder installed models live in: the host's, else `models` next to the settings. pub fn models_dir(&self) -> Option { @@ -151,7 +199,20 @@ impl Session { Task::Mask => &mut self.prefs.roto.model, Task::Face => &mut self.prefs.face.model, }; - *slot = id.to_string(); + let changed = slot != id; + if changed { + *slot = id.to_string(); + } + 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; + s.error_model = None; + s.load_error = None; + }); } /// Load (or drop) the models Settings chooses; called whenever settings change. @@ -173,7 +234,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.load_error.is_some() { return; } let Some(path) = self.model_path(info) else { @@ -181,19 +242,195 @@ impl Session { return; }; let (face_slot, label) = (self.models.face_model.clone(), kind(task).label); - spawn(&lane, label, format!("Loading {}…", info.name), wait, move |status| { + let generation = st.generation; + spawn(&lane, label, format!("Loading {}…", info.name), wait, Some(generation), info.id, 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)) }); } + /// The Roto Brush model in use (`None`: the classic engine). + pub fn models_roto_active(&self) -> Option { + 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) -> std::result::Result { + if self.chosen_model(task) != id { + return Ok(false); + } + 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, 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, id: &str) -> std::result::Result<(), String> { + let started = Instant::now(); + 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)); + } + } + + /// 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 { + 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`). + 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": self.roto_model_loading(), + "error": st.load_error.or(st.error.filter(|_| st.error_model == Some(self.chosen_model(Task::Mask)))), + }) + } + + /// 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.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)) + } + } + + /// 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 + /// 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(()) }; + 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, info.id)?; + } + 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.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.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)) + } + + /// 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, + }); + } + } + + /// 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(); + let st = self.models.status(Task::Mask); + if chosen != seg::CLASSICAL && st.active.is_none() && st.busy.is_none() && st.load_error.is_none() { + self.apply_model(Task::Mask, false); + } + } + /// 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 { - 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); } } @@ -202,6 +439,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 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) } @@ -212,26 +454,47 @@ fn spawn( label: &'static str, what: String, wait: bool, + generation: Option, + model: &'static str, 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); + 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)); + return; + } match r { 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); + s.error_model = Some(model); } } }); @@ -241,7 +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.error = Some(format!("cannot start: {e}")); + 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); }); } } @@ -264,9 +533,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 { @@ -282,8 +559,14 @@ 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")] + let (recommended, hint) = (false, None::); v.push(json!({ + "recommended": recommended, + "hint": hint, "id": m.id, "name": m.name, "description": m.description, @@ -292,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(), @@ -317,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); @@ -337,15 +625,14 @@ 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))?; let wait = b_p(p, "wait").unwrap_or(false); - s.apply_model(task, wait); - if 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()})) } @@ -353,13 +640,16 @@ 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")); } 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, 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") @@ -407,7 +697,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|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) ), @@ -461,3 +751,260 @@ 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::*; + + /// 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(); + 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()); + } + } + + #[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), "mobilesam", 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!(status.load_error.is_none()); + 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, "mobilesam", 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.load_error.as_deref(), generation.map(|_| "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()); + 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)); + } + } + + #[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(())); + 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() { + 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, "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(); + 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, "sam2.1-hiera-tiny", 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).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).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), info.id, 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_err().contains("SHA-256")); + 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(); + 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/analysis/roto.rs b/crates/engine/src/analysis/roto.rs index 813561969..67b0afc68 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}; @@ -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. @@ -112,6 +119,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 { @@ -131,6 +142,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 { @@ -249,6 +262,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()); } } } @@ -279,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 } @@ -432,10 +446,49 @@ 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() }; + // 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_frames(Some(&w.frames), &mut get, &cancel, &mut report); + if cancel() { + let mut s = lock(&shared.state); + s.cancelled = true; + 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(), + }); + if result.is_err() { + chain.forget_tracking(); + } + 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() { let mut s = lock(&shared.state); @@ -592,9 +645,37 @@ 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 { + 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. + // 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() { if !wait { return Err("a Roto Brush job is already running".into()); @@ -697,6 +778,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 }; @@ -723,6 +815,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(); } @@ -757,17 +853,39 @@ 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 }), } } } 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/cmd/roto_cmds.rs b/crates/engine/src/cmd/roto_cmds.rs index a6d389ca1..792c60019 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 ), @@ -196,6 +204,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)?, @@ -230,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; @@ -261,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(), })) } @@ -367,8 +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 = 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))); + 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())) { @@ -410,6 +433,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)) } @@ -428,9 +454,13 @@ 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))); + // The model Settings chose and whether this layer segments with it (propagate refuses or waits otherwise). + out["model"] = s.roto_model(); out["limits"] = json!(rt::frame_limits(&comp, &layer)); let frame = frame_p(s, p, cid, &layer); out["frame"] = json!(frame); @@ -441,7 +471,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); @@ -463,6 +493,28 @@ 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; 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); + let tracked: serde_json::Map = chain + .keys + .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(), "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()) { + "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::>()); Ok(out) diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index 21d316e75..e1baab2d8 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 { @@ -36,7 +62,13 @@ 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 in four sizes (video trackers; see crates/sam2). + 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] + }; + 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()); @@ -69,6 +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()); } + #[cfg(not(target_arch = "wasm32"))] + 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") { @@ -86,3 +120,163 @@ 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] +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); +} + +/// 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] +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); + 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); + let _ = std::fs::remove_dir_all(d); +} + +/// 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 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/engine/src/tests/roto.rs b/crates/engine/src/tests/roto.rs index 64ee30d5e..99c155265 100644 --- a/crates/engine/src/tests/roto.rs +++ b/crates/engine/src/tests/roto.rs @@ -668,3 +668,494 @@ 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())); + 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(); + 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"); + // 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}"); + 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(); + assert_eq!(st["source"], "classical", "{st}"); + assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 0); + 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 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(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}")) }; + 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(0, 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}"); + // 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. + 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. + // 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}"); + 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 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() { + 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) { + 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 tracker that emits a completed frame before failing must leave only classic results behind. +#[test] +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; + 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, + req: &TrackRequest, + frame: &mut dyn FnMut(i64) -> Option>, + _: &dyn Fn() -> bool, + out: &mut dyn FnMut(usize, i64, TrackedFrame), + ) -> effectcraft_segment::Result<()> { + 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, 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 = st.unwrap(); + assert_ne!(st["source"], "tracker", "{st}"); + assert_eq!(st["tracker"]["frames"].as_object().unwrap().len(), 0); + 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"); + } +} + +/// 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); +} + +/// 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, 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})); + 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}"); + // 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)); +} + +#[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"); +} + +#[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/Cargo.toml b/crates/sam2/Cargo.toml new file mode 100644 index 000000000..e45cc8def --- /dev/null +++ b/crates/sam2/Cargo.toml @@ -0,0 +1,25 @@ +[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] +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 optional GPU path (wgpu compute, any vendor); the CPU path stays the reference. +wgpu = { workspace = true } +pollster = { workspace = true } +bytemuck = { workspace = true } +half = "2.7" 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/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/examples/op_bench.rs b/crates/sam2/examples/op_bench.rs new file mode 100644 index 000000000..4e2b5b4a3 --- /dev/null +++ b/crates/sam2/examples/op_bench.rs @@ -0,0 +1,68 @@ +//! 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()); + // 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 new file mode 100644 index 000000000..0b51dd668 --- /dev/null +++ b/crates/sam2/examples/sam2_bench.rs @@ -0,0 +1,126 @@ +//! 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()); + } + // 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); + // 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); + } + // 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/decoder.rs b/crates/sam2/src/decoder.rs new file mode 100644 index 000000000..d6b25a7e4 --- /dev/null +++ b/crates/sam2/src/decoder.rs @@ -0,0 +1,260 @@ +// 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::{Linear, ops::sigmoid, ops::softmax}; + +use crate::nn::{LayerNorm, apply, gelu, 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(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)?); + apply(&self.o, &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 = apply(&self.proj_in, x)?.relu()?; + for l in &self.layers { + x = apply(l, &x)?.relu()?; + } + let x = apply(&self.proj_out, &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()?; + 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)?; + 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/gpu.rs b/crates/sam2/src/gpu.rs new file mode 100644 index 000000000..bc83b605f --- /dev/null +++ b/crates/sam2/src/gpu.rs @@ -0,0 +1,471 @@ +// 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!("gpu/kernels.wgsl"), include_str!("gpu/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, + /// 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. +#[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>, + turn: 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"))) +} + +/// 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")) +} + +/// 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> { + 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()); + 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), 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 { + &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 { + // `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, + 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(), 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. + 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, f16: false }) + } + + 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(()); + } + 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.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 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) { + 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( + 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")?), + ); + Ok(()) + } +} + +pub mod encoder; +pub mod memattn; +pub mod memenc; +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_weights(&flat(w.get(&format!("{name}.weight"), &[out, inp])?)?), + b: g.upload(&flat(w.get(&format!("{name}.bias"), &[out])?)?), + inp, + out, + }) + } + + /// 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_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_weights(&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/encoder.rs b/crates/sam2/src/gpu/encoder.rs new file mode 100644 index 000000000..a33f9069c --- /dev/null +++ b/crates/sam2/src/gpu/encoder.rs @@ -0,0 +1,289 @@ +// 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, false)?; + 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 _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 { + 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); + // 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)); + } + } + // 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..faea20e76 --- /dev/null +++ b/crates/sam2/src/gpu/kernels.wgsl @@ -0,0 +1,165 @@ +// 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; + } + } + } +} + +// 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/kernels_ops.wgsl b/crates/sam2/src/gpu/kernels_ops.wgsl new file mode 100644 index 000000000..419476600 --- /dev/null +++ b/crates/sam2/src/gpu/kernels_ops.wgsl @@ -0,0 +1,206 @@ + +// ---------------------------------------------------------------- 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), 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; + 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) { + 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 +// 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/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/ops.rs b/crates/sam2/src/gpu/ops.rs new file mode 100644 index 000000000..08143724e --- /dev/null +++ b/crates/sam2/src/gpu/ops.rs @@ -0,0 +1,219 @@ +//! 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. + #[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(), u32::from(gelu)], + 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(()) + } + + /// 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")?; + 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, 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")?; + 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..e974c34df --- /dev/null +++ b/crates/sam2/src/gpu/tests.rs @@ -0,0 +1,454 @@ +//! 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, ()>, Arc); + +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(()); + let guard = TURN.lock().unwrap_or_else(PoisonError::into_inner); + match Gpu::shared() { + Ok(g) => Some(Shared(guard, g)), + Err(e) => { + eprintln!("skipped: {e}"); + None + } + } +} + +/// 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_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 }; + 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, false).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"); +} + +#[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"); +} + +#[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"); +} diff --git a/crates/sam2/src/hiera.rs b/crates/sam2/src/hiera.rs new file mode 100644 index 000000000..82862c829 --- /dev/null +++ b/crates/sam2/src/hiera.rs @@ -0,0 +1,449 @@ +// 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::Tensor; +use candle_nn::{Linear, ops::softmax_last_dim}; + +use crate::nn::{LayerNorm, Mlp, apply, 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, + } + } + + /// `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 { + 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), 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 (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. +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 = 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); + 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()?; + // 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, ()))?; + apply(&self.proj, &out) + } +} + +/// 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, + 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 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 { + 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: 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: plan.window, + query_stride: plan.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(&apply(p, &h)?, s)?, + (Some(p), None) => apply(p, &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(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))?)?; + Ok(x.permute((0, 2, 3, 1))?.contiguous()?) + } + + /// Background embedding resized to the grid plus the tiled window embedding, `[1, H, W, C]`. + 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); + 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 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(); + assert!(window_partition(&x, 0).is_err()); + } +} diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs new file mode 100644 index 000000000..f2aba0390 --- /dev/null +++ b/crates/sam2/src/lib.rs @@ -0,0 +1,145 @@ +//! 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; + +pub mod decoder; +pub mod gpu; +pub mod hiera; +pub mod mask; +pub mod memattn; +pub mod memory; +pub mod model; +pub mod neck; +pub mod nn; +pub mod prompt; +pub mod resize; +pub mod session; +pub mod track; +pub mod weights; + +/// 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), + /// A tensor operation failed (shape mismatch, out of memory, device error). + Compute(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}"), + 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. +#[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", + } + } + + /// 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)] +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 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")); + assert_eq!(Error::Cancelled.to_string(), "SAM 2: cancelled"); + } +} diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs new file mode 100644 index 000000000..382f9b514 --- /dev/null +++ b/crates/sam2/src/mask.rs @@ -0,0 +1,227 @@ +// 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 = 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]; + 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)?) +} + +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 { + 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, + info: &'static ModelInfo, + last: Mutex)>>, + used: Mutex, + pub(crate) peak: std::sync::atomic::AtomicU64, +} + +impl Sam2Mask { + pub fn new(model: Arc) -> Sam2Mask { + 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`]). + 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. + 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.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 { + self.info + } + + 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 device_name(&self) -> Option { + 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()) + } +} + +/// 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); + 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::with_info(Arc::new(model), info))) +} + +/// 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 { + &SAM2_TINY +} + +/// 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 { + 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" => 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(), + ) +} + +/// 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::*; + + #[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())); + } + + /// 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(); + assert!(e.contains("does not match"), "{e}"); + } +} diff --git a/crates/sam2/src/memattn.rs b/crates/sam2/src/memattn.rs new file mode 100644 index 000000000..d455f1ed1 --- /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]`. +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; + 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/src/memory.rs b/crates/sam2/src/memory.rs new file mode 100644 index 000000000..7a73fd5a3 --- /dev/null +++ b/crates/sam2/src/memory.rs @@ -0,0 +1,138 @@ +// 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, apply_gelu, gelu, 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_gelu(&self.pw1, &h)?)?; + 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 = gelu(&norm_channels_first(ln, &conv.forward(&m)?)?)?; + } + 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, _) = 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/model.rs b/crates/sam2/src/model.rs new file mode 100644 index 000000000..36e30e310 --- /dev/null +++ b/crates/sam2/src/model.rs @@ -0,0 +1,87 @@ +// 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(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) + } + + /// 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/neck.rs b/crates/sam2/src/neck.rs new file mode 100644 index 000000000..50733fe79 --- /dev/null +++ b/crates/sam2/src/neck.rs @@ -0,0 +1,221 @@ +// 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 std::collections::HashMap; +use std::sync::{Mutex, OnceLock, PoisonError}; + +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`). +/// 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(); + 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 } + } + + 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) + } +} + +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(); + 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; count]; + 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, + positions: PositionCache, +} + +/// 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 { + /// 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() + .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, positions: PositionCache::new(MAX_POSITION_ELEMENTS) }) + } + + /// `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(self.position(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: std::sync::Arc, + 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 { + /// 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 { + 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) }) + } +} + +#[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 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(), 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 = 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/nn.rs b/crates/sam2/src/nn.rs new file mode 100644 index 000000000..256a60da9 --- /dev/null +++ b/crates/sam2/src/nn.rs @@ -0,0 +1,171 @@ +// 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, CustomOp2, Layout, Shape, Tensor}; +use candle_nn::Linear; + +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 { 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`. +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 { + apply(&self.proj_out, &apply_gelu(&self.proj_in, x)?) + } +} + +/// `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 = 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; + +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 new file mode 100644 index 000000000..23d9a7c99 --- /dev/null +++ b/crates/sam2/src/prompt.rs @@ -0,0 +1,237 @@ +// 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 crate::nn::{LayerNorm, gelu, 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]`; 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]; + 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); + } + // _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)?) + } + + /// 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()?) + } + + /// 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 = gelu(&cl(&step(mask, c1)?, &self.mask.ln1)?)?; + let x = gelu(&cl(&step(&x, c2)?, &self.mask.ln2)?)?; + step(&x, c3) + } +} + +#[cfg(test)] +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() + } + + /// Transformers 5.19.0 pads click-only prompts in _embed_points and boxes in _embed_boxes. + #[test] + 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() + 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)); + } + } + 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 }; + 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/src/resize.rs b/crates/sam2/src/resize.rs new file mode 100644 index 000000000..fdbcbab48 --- /dev/null +++ b/crates/sam2/src/resize.rs @@ -0,0 +1,124 @@ +// 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::Result; +use crate::mask::validate_frame_size; +use crate::prompt::INPUT_SIZE; + +/// 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> { + 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; 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() { + 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; 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(); + 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 { + 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() }; + Ok(Tensor::from_vec(data, (1, 1, side, side), &Device::Cpu)?) +} + +#[cfg(test)] +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(); + 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 new file mode 100644 index 000000000..0524ee89e --- /dev/null +++ b/crates/sam2/src/session.rs @@ -0,0 +1,733 @@ +// 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::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::{Encoded, ImageModel}; +use crate::neck::ImageFeatures; +use crate::nn::{apply, linear}; +use crate::prompt::{INPUT_SIZE, Prompt}; +use crate::resize::antialias_bilinear; +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 { + 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, +} + +/// 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 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, + 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, + temporal_pos: Vec, + 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 { + /// 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 { + 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::()?, + 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])?), + }) + } + + 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 = 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 }); + } + 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(Error::Budget(reason)) => log::info!("SAM 2 memory attention on the CPU for this call: {reason}"), + 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)?) + } +} + +/// 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), +} + +/// 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, + peak: usize, +} + +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<'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}"))); + } + Ok(VideoTracker { model, frames, objects: BTreeMap::new(), cache: VecDeque::new(), budget: Budget::default(), peak: 0 }) + } + + 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 + } + + /// 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))); + } + Ok(()) + } + + /// 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.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 { + // 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 { + 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.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 { + 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)?; + + mark("decoder"); + 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)?; + 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 }; + 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); + } + self.check_budget()?; + 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) + } + + /// 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.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(); + 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); + } + self.check_budget()?; + 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<()> { + 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(|(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())); + } + 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 }; + for id in &ids { + 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(()) + } + + /// 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 new file mode 100644 index 000000000..7e73a1c74 --- /dev/null +++ b/crates/sam2/src/track.rs @@ -0,0 +1,181 @@ +// 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, validate_frame_size}; +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}; + +/// 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 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) + } +} + +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 }) +} + +/// 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(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))); + } + 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"))); + } + 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 (oi, prompts) in req.objects.iter().enumerate() { + for (frame, p) in prompts { + 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)); + } + } + 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 first_of: HashMap = HashMap::new(); + for (oi, 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()))?; + // 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(), + bbox: p.bbox.map(|b| [b[0] * sx, b[1] * sy, b[2] * sx, b[3] * sy]), + }; + tracker.prompt(idx, object_id(oi), &prompt)?; + let e = first_of.entry(oi).or_insert(idx); + *e = (*e).min(idx); + } + 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| { + 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(oi, req.first + m.frame as i64, t), + Err(e) => { + failure.get_or_insert(e); + flag.cancel(); + } + } + }; + // 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); + // 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, + } + } +} + +impl ObjectTracker for Sam2Mask { + fn track( + &self, + req: &TrackRequest, + frames: &mut dyn FnMut(i64) -> Option>, + cancel: &dyn Fn() -> bool, + out: &mut dyn FnMut(usize, i64, TrackedFrame), + ) -> effectcraft_segment::Result<()> { + self.touch(); + let r = self.run_track(req, frames, cancel, out).map_err(|e| e.to_string()); + self.touch(); + 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/crates/sam2/src/weights.rs b/crates/sam2/src/weights.rs new file mode 100644 index 000000000..ca246b143 --- /dev/null +++ b/crates/sam2/src/weights.rs @@ -0,0 +1,112 @@ +// 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) + } + + /// 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()); + 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 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() + } + + 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/it/figure.rs b/crates/sam2/tests/it/figure.rs new file mode 100644 index 000000000..2ca4d0b13 --- /dev/null +++ b/crates/sam2/tests/it/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/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/it/reference.rs b/crates/sam2/tests/it/reference.rs new file mode 100644 index 000000000..01648e329 --- /dev/null +++ b/crates/sam2/tests/it/reference.rs @@ -0,0 +1,805 @@ +//! 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(); + note_size(&w); + 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() +} + +/// 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; + +#[test] +fn hiera_matches_pytorch_block_by_block() { + let Some((w, r)) = load() else { return }; + 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(), 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}"); + 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::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(); + 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()); +} + +/// 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(&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}"); + 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()); + 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 < tol(NECK_TOL, r.get("vision_encoder.neck.3").unwrap())); +} + +/// 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 < 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}"); + 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 < loose(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; + } + 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. +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}"); +} + +#[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}"); +} + +/// 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()?) + } +} + +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; + +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(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); + 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 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; + 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(_)))); +} + +#[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()); +} + +#[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::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 |o, f, t| { + assert_eq!(o, 0); + 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::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()); +} + +/// 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()); +} + +/// 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:?}"); +} + +/// 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::shared() { + Ok(g) => 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 < 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 < 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 < 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). +/// 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()); +} + +/// 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/README.md b/crates/segment/README.md index 79b1d996b..a95c92f93 100644 --- a/crates/segment/README.md +++ b/crates/segment/README.md @@ -12,7 +12,10 @@ 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. + 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 @@ -36,6 +39,10 @@ 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 | +| `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 adeb44881..7fc0e4ff6 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -45,6 +45,77 @@ 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 + } + /// 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 + } + /// 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`). +#[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, 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, + /// 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. +#[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 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(usize, i64, TrackedFrame), + ) -> Result<()>; } /// What a model is for. @@ -70,8 +141,10 @@ 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, + /// 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). @@ -91,11 +164,82 @@ 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", }; +/// SAM 2.1 Tiny (Meta AI, Apache-2.0): a video object tracker with streaming memory. The weights are +/// 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, + 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: "", + 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", +}; + +/// 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: "", + 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", +}; + +/// 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: "", + 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", +}; + +/// 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: "", + 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", +}; + pub const FACE_LANDMARKER: ModelInfo = ModelInfo { id: "mediapipe-face", task: Task::Face, @@ -106,13 +250,14 @@ 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", }; /// 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, 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. @@ -144,21 +289,42 @@ 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)?))), - _ => 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, } } /// The notice kept next to installed weights: what they are, who made them, their licence. pub fn notice(m: &ModelInfo) -> String { 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: {}\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 ) } @@ -260,13 +426,37 @@ 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.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)); } + 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"] + ); + 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!(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"))); // A file that isn't the published one is refused before anything is parsed. 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/track/src/roto.rs b/crates/track/src/roto.rs index 103d695dd..c742654e9 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; @@ -87,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 @@ -129,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] } @@ -223,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(); @@ -902,6 +944,70 @@ 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. 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 (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)))); + } + Prompt { points, ..prompt } +} + +/// [`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, 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); + 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; + +/// 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/bench.rs b/crates/track/src/roto/bench.rs new file mode 100644 index 000000000..98ada3903 --- /dev/null +++ b/crates/track/src/roto/bench.rs @@ -0,0 +1,447 @@ +//! 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 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, 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: 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. + 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: String, + 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, + /// 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`). + pub peak_mb: Option, + /// The matte of every frame (not serialised). + #[serde(skip)] + pub mattes: 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(), object: 0 } +} + +/// 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: 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). +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.clone(), + 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, + fixes: None, + peak_mb: None, + mattes: mattes.to_vec(), + } +} + +/// 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().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); + 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 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); + } + } + let secs = start.elapsed().as_secs_f64(); + 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) +} + +/// 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 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(); + // 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; 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() == 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) { + *slot = segment_tracked(img, &strokes, tf.prob, tf.occluded, &opts, 0.0).matte; + } + }); + r.ok()?; + 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 +/// 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]> { + 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 % 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]) + }; + 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` 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; + 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} {:>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} {:>7}\n", + r.scenario, + r.tracker, + r.mean_iou, + r.min_iou, + r.flicker, + back, + r.secs_per_frame, + r.fixes.as_deref().unwrap_or("-"), + r.peak_mb.map_or_else(|| "-".to_string(), |v| format!("{v:.0}")) + )); + } + 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 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 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)); + } +} diff --git a/crates/track/src/roto/tests.rs b/crates/track/src/roto/tests.rs index 8a5cab98c..bfc00df74 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()); @@ -579,3 +586,139 @@ 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, 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, 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] +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); +} + +#[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]], 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/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/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 { diff --git a/crates/ui-egui/src/panels/roto_tool.rs b/crates/ui-egui/src/panels/roto_tool.rs index ca3c56ac5..194bba26a 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}; @@ -21,6 +21,42 @@ 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) +} + +/// 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 { @@ -94,6 +130,19 @@ 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)| (object_color(obj, o.boundary_color), object_mask_in_input(&seg, &inp))) + .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 +166,30 @@ 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 + && (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() + .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 { @@ -243,6 +315,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 +350,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() { @@ -352,6 +445,45 @@ 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(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::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 { + 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})); + } + 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 += 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::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})); + } + app.auto.add("layerPanel.roto.newObject", r, "New object"); + 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::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})) + { + 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() { @@ -363,3 +495,39 @@ pub fn bar_buttons(app: &mut EffectcraftApp, ui: &mut egui::Ui, p: &egui::Painte } } } + +#[cfg(test)] +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)); + 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); + } +} diff --git a/crates/ui-egui/src/panels/settings.rs b/crates/ui-egui/src/panels/settings.rs index 816b0ca01..69f7623dd 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)); } @@ -505,12 +508,19 @@ 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; } 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() { @@ -921,3 +931,28 @@ 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_eq!(app.auto.find(&format!("settings.roto.source.{id}")).unwrap().label, "model.safetensors"); + } + assert!(actions.is_empty()); + } +} diff --git a/crates/ui-egui/tests/it/main.rs b/crates/ui-egui/tests/it/main.rs index 1d066a503..36a5aa485 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; @@ -72,6 +73,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/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/it/ui_roto_objects.rs b/crates/ui-egui/tests/it/ui_roto_objects.rs new file mode 100644 index 000000000..cbf52d07d --- /dev/null +++ b/crates/ui-egui/tests/it/ui_roto_objects.rs @@ -0,0 +1,100 @@ +//! 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 _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(); + 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(); + } + 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}"); + + // 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(); + } + // 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"); + + // 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"); +} diff --git a/crates/ui-egui/tests/it/ui_roto_span_confidence.rs b/crates/ui-egui/tests/it/ui_roto_span_confidence.rs new file mode 100644 index 000000000..d9730b778 --- /dev/null +++ b/crates/ui-egui/tests/it/ui_roto_span_confidence.rs @@ -0,0 +1,119 @@ +//! 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 _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(); + 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(); + 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(); + + 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); + + // 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(); + 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:?}"); + } +} 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/docs/agents.md b/docs/agents.md index f5197378f..34162aacd 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -417,8 +417,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). @@ -525,11 +526,22 @@ 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.select {id: classical|mobilesam}` chooses which one Roto Brush uses. Loading runs - in the background; `"wait": true` (also on download and install) returns once it is in use. + 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` 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. + 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/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)) | diff --git a/docs/effects.md b/docs/effects.md index b3adb46ec..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. @@ -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 diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md new file mode 100644 index 000000000..0204058d5 --- /dev/null +++ b/docs/roto-brush-sam2.md @@ -0,0 +1,209 @@ +# 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 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. + 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 + Roto Brush frame: your strokes stay hard constraints and edges snap to colour. +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. + `tracker.saved` names the model the saved frames came from. + +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. 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 +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). + +**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. + +## 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 | 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 | +| 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 | + +### 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 +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 +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 +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). + +## 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. + +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 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 +`transformers` SAM 2 implementation, Apache-2.0 (see its `NOTICE`). Every stage is checked against PyTorch by +`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 | +|---|---| +| 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: 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`). + +| 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 (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: + +| Clip | Classical | MobileSAM | SAM 2.1 Tiny | +|---|---|---|---| +| 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 last two rows are the difference. + +### Corrections needed (`--fixes`) +`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 | +|---|---|---|---| +| 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. 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 diff --git a/tools/sam2_reference.py b/tools/sam2_reference.py new file mode 100644 index 000000000..2cfd03fde --- /dev/null +++ b/tools/sam2_reference.py @@ -0,0 +1,85 @@ +#!/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 +import torch.nn.functional as F +from transformers import Sam2Model + + +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() + image = synthetic_image() + # 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): + 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. + 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)) + 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.clone().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]) diff --git a/tools/sam2_video_reference.py b/tools/sam2_video_reference.py new file mode 100644 index 000000000..a05223be9 --- /dev/null +++ b/tools/sam2_video_reference.py @@ -0,0 +1,136 @@ +#!/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) + + # 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(): + 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]) 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),