pauljd/effectcraft
1
0
Fork 0
mirror of https://github.com/storytold/effectcraft.git synced 2026-11-11 23:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/main' into HEAD

Browse Source 
# Conflicts:
#	crates/engine/Cargo.toml
This commit is contained in:
bflatastic
2026-10-30 17:22:12 +07:00
parent 2a72f03872 3d285116d8
commit 334bf0f785
68 changed files with 11695 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+346 -9
View File
@@ -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",
]
Cargo.toml
+1
View File
@@ -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" }
apps/effectcraft-cli/Cargo.toml
+3
View File
@@ -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"
apps/effectcraft-cli/src/main.rs
+126
View File
@@ -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 <info|commands|exec|run|props|get|se
bench --ops [--small] [--layers N] [--comps N] [--footage N]
everyday operations on a large generated project: startup, open,
save, auto-save, edits + undo/redo, timeline, Project panel
bench --roto [--frames N] [--sam2 FILE] [--mobilesam FILE] [--fixes] [--clip FRAMES_DIR --truth MASKS_DIR] [--json]
Roto Brush trackers on synthetic clips with known mattes: movement,
lighting change, occlusion, leaving and re-entering the frame
script FILE.jsx [F.ecproj] | --eval CODE run JavaScript with the After Effects-style object model
(app.project, comps, layers, properties…); prints writeLn
output and the result; errors exit 1 with file:line:col
@@ -137,6 +142,11 @@ options: --project F.ecproj | --demo | --empty --save | --save-as F --bridge
/// Options that take a value.
const VALUED: &[&str] = &[
"--n",
"--clip",
"--truth",
"--frames",
"--sam2",
"--mobilesam",
"--play",
"--layers",
"--comps",
@@ -187,6 +197,7 @@ const VALUED: &[&str] = &[
/// Options without a value. Any other `--option` is a usage error.
const FLAGS: &[&str] = &[
"--json",
"--fixes",
"--save",
"--demo",
"--empty",
@@ -199,6 +210,7 @@ const FLAGS: &[&str] = &[
"--work-area",
"--queue",
"--ops",
"--roto",
"--small",
"--adv3d",
"--autosave",
@@ -975,6 +987,9 @@ fn bench_cmd(args: &Args) -> 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<T>(
flag: &str,
named: Option<&str>,
installed: Option<std::path::PathBuf>,
build: impl Fn(&[u8]) -> Result<T, String>,
) -> Result<Option<T>, 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<std::sync::Arc<dyn MaskModel>> = 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<std::sync::Arc<dyn MaskModel>> =
bench_weights("--sam2", args.opt("--sam2"), installed_sam2, |bytes| effectcraft_sam2::mask::load(bytes).map(|m| m as std::sync::Arc<dyn MaskModel>))?;
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<bench::Run>| 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<Vec<std::path::PathBuf>, String> {
let mut files: Vec<std::path::PathBuf> = 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<effectcraft_track::roto::bench::Scenario, String> {
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
apps/effectcraft-cli/tests/it/cli.rs
+107
View File
@@ -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() {
crates/effects/src/catalog.rs
+2 -1
View File
@@ -128,7 +128,8 @@ pub fn effects_markdown() -> String {
s.push_str("# Effects\n\n");
s.push_str("<!-- Generated from the effect registry (crates/effects/src/catalog.rs); do not edit by hand.\n");
s.push_str(" Regenerate: UPDATE_DOCS=1 cargo test -p effectcraft-effects --lib effects_doc_is_current -->\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();
crates/effects/src/migrate.rs
+7 -1
View File
@@ -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);
crates/effects/src/object_matte.rs
+5 -5
View File
@@ -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<i64> {
pub fn chain(params: &Params, layer_size: [f64; 2], fps: f64, f: i64, scale: f64, model: Option<Arc<dyn MaskModel>>) -> Option<Chain> {
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
crates/effects/src/roto.rs
+1016 -13
View File
File diff suppressed because it is too large Load Diff
crates/effects/src/roto/pending.rs
+138
View File
@@ -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<i64, Vec<Option<TrackedFrame>>>,
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::<Option<TrackedFrame>>()).saturating_add(256)
}
fn mask_bytes(mask: &TrackedFrame) -> usize {
mask.prob.capacity().saturating_mul(size_of::<f32>())
}
/// 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<Option<Vec<Option<TrackedFrame>>>, 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<i64, Vec<Option<TrackedFrame>>> {
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());
}
}
crates/engine/Cargo.toml
+2 -1
View File
@@ -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 }
crates/engine/src/analysis/models.rs
+564 -17
View File
@@ -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<String>,
/// Model whose load/download owns this lane, for selection-specific waits.
pub busy_model: Option<&'static str>,
/// The model in use.
pub active: Option<String>,
/// The last failure, for the Settings page.
pub error: Option<String>,
/// 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<String>,
/// 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<Loaded>) {
/// 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<Loaded>) {
activate(self.lane(task), &self.face_model, task, model);
}
}
@@ -90,6 +107,30 @@ fn activate(lane: &Arc<Mutex<Status>>, 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<Mutex<Status>>, 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<Mutex<Status>>, 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<PathBuf> {
@@ -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<String> {
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<bool, String> {
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<String> {
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<String> {
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<Loaded, String> {
let bytes = std::fs::read(path).map_err(|e| format!("cannot read {}: {e}", path.display()))?;
// SAM 2 runs on candle and is native only; every other model loads through the registry.
#[cfg(not(target_arch = "wasm32"))]
if 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<u64>,
model: &'static str,
job: impl FnOnce(&Arc<Mutex<Status>>) -> std::result::Result<String, String> + 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<Item = &'static ModelInfo> {
seg::models(task).filter(move |m| !wasm || !m.id.starts_with("sam2.1-hiera-"))
}
fn target_models(task: Task) -> impl Iterator<Item = &'static ModelInfo> {
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<Value> {
@@ -282,8 +559,14 @@ fn list(s: &mut Session, task: Task) -> Result<Value> {
"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::<String>);
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<Value> {
"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<Value> {
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<Value> {
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<Value> {
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<CommandSpec> {
"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<CommandSpec> {
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::<Vec<_>>(), ["mediapipe-face"]);
}
}
}
crates/engine/src/analysis/roto.rs
+124 -6
View File
@@ -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<String>,
/// 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<fx::TrackSummary>,
/// Why the tracker was not used, when a tracking model was chosen but the classic engine ran.
pub tracker_note: Option<String>,
}
impl RotoProgress {
@@ -131,6 +142,8 @@ pub struct RotoShared {
pub cancel: AtomicBool,
pub state: Mutex<RotoProgress>,
frozen: Mutex<Option<FrozenCache>>,
/// The tracker's frames, as the saved tracking (JSON) to store with the instance.
saved: Mutex<Option<String>>,
}
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(&params, 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(&params_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<usize, String> {
// 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<i64, fx::FrameTrust> {
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(&params_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<i64, bool> {
let mut out = BTreeMap::new();
crates/engine/src/cmd/roto_cmds.rs
+59 -7
View File
@@ -27,7 +27,7 @@ pub fn specs() -> Vec<CommandSpec> {
"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<CommandSpec> {
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<Vec<[f64; 2]>> {
fn stroke(s: &mut Session, p: &Value) -> Result<Value> {
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<Value> {
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<Value> {
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<Value> {
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<Value> {
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<Value> {
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::<Vec<_>>());
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<Value> {
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<Value> {
}
}
}
// 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(&params_static(&g), size, comp.frame_rate.as_f64(), 1.0);
let tracked: serde_json::Map<String, Value> = 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::<Vec<_>>());
Ok(out)
crates/engine/src/tests/models.rs
+195 -1
View File
@@ -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(&params).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(&params).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(&params).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<f64>);
impl MaskModel for Fake {
fn info(&self) -> &'static ModelInfo {
&effectcraft_segment::SAM2_TINY
}
fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &Prompt) -> effectcraft_segment::Result<Vec<f32>> {
Err("no weights".into())
}
fn idle_secs(&self) -> Option<f64> {
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<dyn MaskModel> = 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<Vec<f32>> {
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<dyn MaskModel> = 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);
}
crates/engine/src/tests/roto.rs
+491
View File
@@ -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<Vec<f32>> {
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<std::sync::Arc<effectcraft_segment::ClipFrame>>,
_: &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::<Vec<f32>>();
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(&gt(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<Vec<f32>> {
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<Vec<f32>> {
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<Arc<ClipFrame>>,
_: &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(&params, 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<Vec<_>> = (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<Vec<_>> = (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<Image>, 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<Vec<(i64, usize, Vec<i64>)>>,
}
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<Vec<f32>> {
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<std::sync::Arc<effectcraft_segment::ClipFrame>>,
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);
}
}
crates/sam2/Cargo.toml
+25
View File
@@ -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"
crates/sam2/NOTICE
+11
View File
@@ -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).
crates/sam2/examples/gpu_matmul.rs
+164
View File
@@ -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<storage, read> a: array<f32>;
@group(0) @binding(1) var<storage, read> b: array<f32>;
@group(0) @binding(2) var<storage, read_write> c: array<f32>;
@group(0) @binding(3) var<uniform> p: P;
var<workgroup> sa: array<f32, 1024>;
var<workgroup> sb: array<f32, 1024>;
@compute @workgroup_size(16, 16)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
let row0 = wg.y * 64u;
let col0 = wg.x * 64u;
let tid = li.y * 16u + li.x;
var acc: array<f32, 16>;
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<f32, 4>;
var br: array<f32, 4>;
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<dyn std::error::Error>> {
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::<f32>()?;
let bv = b.flatten_all()?.to_vec1::<f32>()?;
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(&params),
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(())
}
crates/sam2/examples/op_bench.rs
+68
View File
@@ -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<T>(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<dyn std::error::Error>> {
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(())
}
crates/sam2/examples/sam2_bench.rs
+126
View File
@@ -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<dyn std::error::Error>> {
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<Tensor> {
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::<f64>() / 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(())
}
crates/sam2/src/decoder.rs
+260
View File
@@ -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<Attention> {
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<Tensor> {
let split = |t: Tensor| -> Result<Tensor> { 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<Linear>,
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<FeedForward> {
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::<Result<_>>()?,
proj_out: linear(w, &format!("{prefix}.proj_out"), hidden, output)?,
sigmoid_output,
})
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
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<Block>,
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<FeedForward>,
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<Tensor> {
Ok(x.broadcast_add(&b.reshape((1, b.dim(0)?, 1, 1))?)?)
}
impl MaskDecoder {
pub fn load(w: &Weights) -> Result<MaskDecoder> {
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::<Result<Vec<_>>>()?;
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::<Result<_>>()?,
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<DecoderOutput> {
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<Tensor> { 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<Tensor> = self.hyper.iter().enumerate().map(|(i, m)| m.forward(&mask_tokens_out.narrow(1, i, 1)?.squeeze(1)?)).collect::<Result<_>>()?;
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::<f32>()?;
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::<f32>()?;
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)?))
}
crates/sam2/src/gpu.rs
+471
View File
@@ -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<HashMap<&'static str, Arc<wgpu::ComputePipeline>>>,
encoder: Mutex<Option<wgpu::CommandEncoder>>,
turn: Mutex<()>,
error: Arc<Mutex<Option<String>>>,
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> {
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<std::result::Result<Arc<Gpu>, 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<Arc<Gpu>> {
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<bool> {
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<Gpu> {
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<Mutex<Option<String>>> = 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<u32> = 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<Buf> {
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<wgpu::ComputePipeline> {
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<wgpu::BindGroupEntry<'_>> =
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<Vec<f32>> {
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::<u8, f32>(&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)],
&params,
(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<Vec<f32>> {
Ok(t.flatten_all()?.to_vec1::<f32>()?)
}
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<GLinear> {
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<GLinear> {
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<GLinear> {
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<GNorm> {
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;
crates/sam2/src/gpu/encoder.rs
+289
View File
@@ -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<GLinear>,
}
/// 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<Gpu>,
cfg: HieraConfig,
patch_w: Buf,
patch_b: Buf,
blocks: Vec<GBlock>,
stage_ends: Vec<usize>,
neck: Vec<GLinear>,
top_down: Vec<usize>,
skip0: GLinear,
skip1: GLinear,
/// The position embedding of each grid size, uploaded on first use.
pos: Mutex<HashMap<(usize, usize), Buf>>,
/// The CPU backbone, which owns the position embedding tables.
cpu: Arc<Hiera>,
eps: f32,
}
impl GpuEncoder {
/// Upload the backbone, the neck and the skip convolutions. `cpu` supplies the position embedding.
pub fn load(gpu: Arc<Gpu>, w: &Weights, cpu: Arc<Hiera>) -> Result<GpuEncoder> {
let cfg = cpu.config().clone();
let flat = |t: Tensor| -> Result<Vec<f32>> { Ok(t.flatten_all()?.to_vec1::<f32>()?) };
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<usize> = 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::<Result<Vec<_>>>()?;
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<Buf> {
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<Buf> {
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::<f32>()?;
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<GpuFeatures> {
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::<f32>()?);
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<Tensor> {
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)?) })
}
}
crates/sam2/src/gpu/kernels.wgsl
+165
View File
@@ -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<vec4<u32>, 6> }
@group(0) @binding(0) var<storage, read> in0: array<f32>;
@group(0) @binding(1) var<storage, read> in1: array<f32>;
@group(0) @binding(2) var<storage, read> in2: array<f32>;
@group(0) @binding(3) var<storage, read_write> out0: array<f32>;
@group(0) @binding(4) var<storage, read> in3: array<f32>;
@group(0) @binding(5) var<uniform> p: P;
fn u(i: u32) -> u32 { return p.v[i / 4u][i % 4u]; }
fn f(i: u32) -> f32 { return bitcast<f32>(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<workgroup> sa: array<f32, 1024>;
var<workgroup> sb: array<f32, 1024>;
@compute @workgroup_size(16, 16)
fn matmul(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
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<f32, 16>;
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<f32, 4>;
var br: array<f32, 4>;
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<storage, read> w16: array<u32>;
@compute @workgroup_size(16, 16)
fn matmul_w16(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
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<f32, 16>;
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<f32, 4>;
var br: array<f32, 4>;
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;
}
}
}
}
crates/sam2/src/gpu/kernels_ops.wgsl
+206
View File
@@ -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<u32>, nw: vec3<u32>, li: vec3<u32>) -> u32 { return (wg.y * nw.x + wg.x) * 256u + li.x; }
var<workgroup> red: array<f32, 64>;
// 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<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
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<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
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<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
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<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
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<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
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<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
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<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
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<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
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<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
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<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
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];
}
crates/sam2/src/gpu/memattn.rs
+184
View File
@@ -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<Gpu>,
layers: Vec<GLayer>,
final_ln: GNorm,
cos: Buf,
sin: Buf,
}
fn attn(g: &Gpu, w: &Weights, prefix: &str, kv: usize) -> Result<GAttn> {
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<Gpu>, w: &Weights) -> Result<GpuMemoryAttention> {
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::<Result<Vec<_>>>()?;
let (cos, sin) = rope_tables(10000.0)?;
let flat = |t: Tensor| -> Result<Vec<f32>> { Ok(t.flatten_all()?.to_vec1::<f32>()?) };
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<Buf> {
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<Buf> {
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<Buf> {
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<Tensor> {
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<Vec<f32>> { Ok(t.contiguous()?.flatten_all()?.to_vec1::<f32>()?) };
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)?)
}
}
crates/sam2/src/gpu/memenc.rs
+127
View File
@@ -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<Gpu>,
down: Vec<(GConv, GNorm)>,
final_conv: GLinear,
feature_projection: GLinear,
fuser: Vec<GFuser>,
projection: GLinear,
}
fn conv(g: &Gpu, w: &Weights, name: &str, ci: usize, co: usize, k: usize, groups: usize) -> Result<GConv> {
let flat = |t: Tensor| -> Result<Vec<f32>> { Ok(t.flatten_all()?.to_vec1::<f32>()?) };
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<Gpu>, w: &Weights) -> Result<GpuMemoryEncoder> {
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::<f32>()?;
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::<Result<Vec<_>>>()?;
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<Tensor> {
let g = &self.gpu;
let _turn = g.exclusive();
let flat = |t: &Tensor| -> Result<Vec<f32>> { Ok(t.contiguous()?.flatten_all()?.to_vec1::<f32>()?) };
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)?)
}
}
crates/sam2/src/gpu/ops.rs
+219
View File
@@ -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> {
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(())
}
}
crates/sam2/src/gpu/tests.rs
+454
View File
@@ -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<Gpu>);
impl std::ops::Deref for Shared {
type Target = Gpu;
fn deref(&self) -> &Gpu {
&self.1
}
}
pub(crate) fn gpu() -> Option<Shared> {
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<f32> {
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::<f32>());
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<f32> {
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<f32> = 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<f32> = (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<f32> = 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<f32> {
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<f32> = x
.chunks(cols)
.flat_map(|r| {
let mean = r.iter().sum::<f32>() / cols as f32;
let var = r.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / cols as f32;
r.iter().enumerate().map(|(i, v)| (v - mean) / (var + 1e-6).sqrt() * w[i] + b[i]).collect::<Vec<_>>()
})
.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::<Vec<_>>();
let gx = g.upload(&x);
g.softmax(&gx, rows, cols).unwrap();
let want: Vec<f32> = x
.chunks(cols)
.flat_map(|r| {
let m = r.iter().cloned().fold(f32::MIN, f32::max);
let e: Vec<f32> = r.iter().map(|v| (v - m).exp()).collect();
let s: f32 = e.iter().sum();
e.into_iter().map(|v| v / s).collect::<Vec<_>>()
})
.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<f32> = 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::<f32>().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<f32> = (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::<f32>().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<f32> = x
.chunks(cols)
.flat_map(|r| {
let mean = r.iter().sum::<f32>() / cols as f32;
let var = r.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / cols as f32;
r.iter().enumerate().map(|(i, v)| gelu((v - mean) / (var + 1e-6).sqrt() * wt[i] + b[i])).collect::<Vec<_>>()
})
.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");
}
crates/sam2/src/hiera.rs
+449
View File
@@ -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<usize>,
pub blocks: Vec<usize>,
pub heads: Vec<usize>,
pub windows: Vec<usize>,
pub global_blocks: Vec<usize>,
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<T: Copy>(v: &[T], i: usize, what: &str) -> Result<T> {
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<Tensor> {
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<Tensor> {
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<usize>,
}
impl Attention {
fn forward(&self, x: &Tensor) -> Result<Tensor> {
let (b, h, w, _) = x.dims4()?;
let qkv = apply(&self.qkv, x)?.reshape((b, h * w, 3, self.heads, ()))?;
let part = |i: usize| -> Result<Tensor> { 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<usize>,
}
impl BlockPlan {
pub(crate) fn new(cfg: &HieraConfig, stage: usize, block_in_stage: usize, total: usize) -> Result<BlockPlan> {
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<Vec<BlockPlan>> {
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<Linear>,
window: usize,
query_stride: Option<usize>,
}
impl Block {
fn load(w: &Weights, cfg: &HieraConfig, stage: usize, block_in_stage: usize, total: usize) -> Result<Block> {
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<Tensor> {
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<f32> {
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<f32>,
pos_embed_window: Vec<f32>,
window0: usize,
hidden: usize,
pub blocks: Vec<Block>,
stage_ends: Vec<usize>,
}
impl Hiera {
pub fn load(w: &Weights, cfg: &HieraConfig) -> Result<Hiera> {
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<Vec<f32>> { Ok(t.flatten_all()?.to_vec1::<f32>()?) };
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<Tensor> {
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<Tensor> {
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<Vec<Tensor>> {
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<Tensor>)> {
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<f32> = (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::<f32>().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::<f32>().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());
}
}
crates/sam2/src/lib.rs
+145
View File
@@ -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<candle_core::Error> for Error {
fn from(e: candle_core::Error) -> Error {
Error::Compute(e.to_string())
}
}
pub type Result<T> = std::result::Result<T, Error>;
/// 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<Size> {
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::<usize>();
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");
}
}
crates/sam2/src/mask.rs
+227
View File
@@ -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<Tensor> {
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<usize> {
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<VideoModel>,
info: &'static ModelInfo,
last: Mutex<Option<(u64, Arc<Encoded>)>>,
used: Mutex<std::time::Instant>,
pub(crate) peak: std::sync::atomic::AtomicU64,
}
impl Sam2Mask {
pub fn new(model: Arc<VideoModel>) -> Sam2Mask {
Sam2Mask::with_info(model, &SAM2_TINY)
}
/// A model of a given registry entry (its size).
pub fn with_info(model: Arc<VideoModel>, 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<VideoModel> {
&self.model
}
fn encoded(&self, rgb: &[[f32; 3]], w: usize, h: usize) -> Result<Arc<Encoded>> {
let key = fingerprint(rgb, w, h);
if let Some((k, e)) = self.last.lock().unwrap_or_else(PoisonError::into_inner).as_ref()
&& *k == key
{
return Ok(e.clone());
}
let enc = Arc::new(self.model.encode(&preprocess(rgb, w, h)?)?);
*self.last.lock().unwrap_or_else(PoisonError::into_inner) = Some((key, enc.clone()));
Ok(enc)
}
fn probability(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &effectcraft_segment::Prompt) -> Result<Vec<f32>> {
let enc = self.encoded(rgb, w, h)?;
let (sx, sy) = (INPUT_SIZE as f32 / w as f32, INPUT_SIZE as f32 / h as f32);
let p = Prompt {
points: prompt.points.iter().map(|(c, fg)| ([c[0] * sx, c[1] * sy], *fg)).collect(),
bbox: prompt.bbox.map(|b| [b[0] * sx, b[1] * sy, b[2] * sx, b[3] * sy]),
};
if p.points.is_empty() && p.bbox.is_none() {
return Err(Error::Input("SAM 2 needs at least one point or a box".into()));
}
let multimask = p.points.len() + 2 * usize::from(p.bbox.is_some()) <= 1;
let out = self.model.image.segment(&enc, &p, multimask)?;
let ious = out.iou.flatten_all()?.to_vec1::<f32>()?;
let best = ious.iter().enumerate().fold((0usize, f32::NEG_INFINITY), |b, (i, v)| if *v > b.1 { (i, *v) } else { b }).0;
let logits = out.masks.narrow(1, best, 1)?.upsample_bilinear2d(h, w, false)?;
Ok(candle_nn::ops::sigmoid(&logits)?.flatten_all()?.to_vec1::<f32>()?)
}
}
impl MaskModel for Sam2Mask {
fn info(&self) -> &'static ModelInfo {
self.info
}
fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &effectcraft_segment::Prompt) -> effectcraft_segment::Result<Vec<f32>> {
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<String> {
Some(self.model.device_name())
}
fn peak_state_bytes(&self) -> Option<u64> {
Some(self.peak.load(std::sync::atomic::Ordering::Relaxed))
}
fn idle_secs(&self) -> Option<f64> {
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<Arc<Sam2Mask>, 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<String> {
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<bool>) -> Option<String> {
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::<f32>().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}");
}
}
crates/sam2/src/memattn.rs
+173
View File
@@ -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<f32> = (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<Tensor> {
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<Tensor> {
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<RopeAttention> {
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<Tensor> {
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<Tensor> { 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<Layer>,
final_ln: LayerNorm,
tables: (Tensor, Tensor),
}
impl MemoryAttention {
pub fn load(w: &Weights) -> Result<MemoryAttention> {
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::<Result<Vec<_>>>()?;
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<Tensor> {
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::<f32>().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::<f32>().unwrap();
let s = sin.narrow(0, 0, 1).unwrap().to_vec2::<f32>().unwrap();
assert!(c[0].iter().all(|v| *v == 1.0) && s[0].iter().all(|v| *v == 0.0));
}
}
crates/sam2/src/memory.rs
+138
View File
@@ -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<Conv> {
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<Tensor> {
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<Tensor> {
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<Tensor> {
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<FuserBlock>,
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<MemoryEncoder> {
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::<Result<Vec<_>>>()?;
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<Memory> {
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<Memory> {
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()? })
}
}
crates/sam2/src/model.rs
+87
View File
@@ -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<ImageModel> {
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> {
ImageModel::load(&Weights::open(path)?)
}
/// Encode a normalised `[1, 3, 1024, 1024]` frame.
pub fn encode(&self, pixels: &Tensor) -> Result<Encoded> {
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<DecoderOutput> {
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<DecoderOutput> {
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
}
}
crates/sam2/src/neck.rs
+221
View File
@@ -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<Tensor> {
static CACHE: OnceLock<PositionCache> = OnceLock::new();
CACHE.get_or_init(|| PositionCache::new(1024 * 1024 / size_of::<f32>())).get(num_pos_features, h, w)
}
const MAX_POSITION_ELEMENTS: usize = 128 * 1024 * 1024 / size_of::<f32>();
struct PositionCache {
entries: Mutex<HashMap<(usize, usize, usize), Tensor>>,
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<Tensor> {
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<Tensor> {
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<f32> = (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<f32> {
// 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<Tensor> {
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<Conv1x1>,
top_down: Vec<usize>,
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<Tensor>,
pub positions: Vec<Tensor>,
}
impl Neck {
/// Position tables shared by the CPU and GPU paths, released with this neck.
pub(crate) fn position(&self, h: usize, w: usize) -> Result<Tensor> {
self.positions.get(self.pos_features, h, w)
}
pub fn load(w: &Weights, channels: &[usize], fpn: usize, top_down: &[usize]) -> Result<Neck> {
let convs = channels
.iter()
.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::<Result<Vec<_>>>()?;
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<NeckOutput> {
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<Tensor> = 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<Hiera>,
pub neck: Neck,
levels: usize,
}
/// The multi-scale image features the decoder and memory modules use, highest resolution first
/// (256², 128², 64² for a 1024 input), `[B, C, H, W]`.
pub struct ImageFeatures {
pub features: Vec<Tensor>,
pub positions: Vec<Tensor>,
}
impl VisionEncoder {
/// The encoder of whichever size the checkpoint is.
pub fn load(w: &Weights) -> Result<VisionEncoder> {
let cfg = HieraConfig::for_size(crate::Size::detect(w)?);
let channels: Vec<usize> = 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<ImageFeatures> {
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<Tensor>| -> Vec<Tensor> {
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::<f32>().unwrap(), neck.position(2, 4).unwrap().flatten_all().unwrap().to_vec1::<f32>().unwrap());
neck.position(3, 4).unwrap();
let entries = cache.entries.lock().unwrap();
assert_eq!(entries.len(), 1, "older tables are evicted to fit the budget");
assert!(entries.values().map(Tensor::elem_count).sum::<usize>() <= 64);
drop(entries);
let large = neck.position(8, 8).unwrap();
assert_eq!(large.dims(), &[1, 4, 8, 8]);
assert_eq!(cache.entries.lock().unwrap().len(), 1, "a table over budget is returned without retention");
}
}
crates/sam2/src/nn.rs
+171
View File
@@ -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<Linear> {
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<Tensor> {
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<LayerNorm> {
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<Mlp> {
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<Tensor> {
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<Tensor> {
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<Tensor> {
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<usize> = 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<Tensor> {
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<Tensor> {
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<Tensor> {
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<f32> = (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::<f32>().unwrap(), b.flatten_all().unwrap().to_vec1::<f32>().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);
}
}
crates/sam2/src/prompt.rs
+237
View File
@@ -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<f32>,
point_embed: Vec<f32>,
not_a_point: Vec<f32>,
no_mask: Tensor,
mask: MaskEmbed,
}
fn rows(t: Tensor) -> Result<Vec<f32>> {
Ok(t.flatten_all()?.to_vec1::<f32>()?)
}
impl PromptEncoder {
pub fn load(w: &Weights) -> Result<PromptEncoder> {
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<f32> {
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<Tensor> {
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<Tensor> {
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<f32> = 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<f32> = 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<f32> = 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<Tensor> {
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<Tensor> {
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<Tensor> {
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<Tensor> {
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<Tensor> { 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<f32> = (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::<f32>().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::<f32>().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::<f32>().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());
}
}
crates/sam2/src/resize.rs
+124
View File
@@ -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<f32>)> {
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<f32> = (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<Vec<f32>> {
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<Tensor> {
validate_frame_size(w, h, mask.len())?;
let side = INPUT_SIZE;
let f: Vec<f32> = 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::<f32>().unwrap();
assert!(v.iter().all(|x| *x == 0.0 || *x == 1.0));
assert!(v.iter().sum::<f32>() > 0.0);
}
}
crates/sam2/src/session.rs
+733
View File
@@ -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<Tensor>;
}
/// Cooperative cancellation shared with the caller.
#[derive(Clone, Default)]
pub struct Cancel(Arc<AtomicBool>);
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<f32>,
/// 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<GpuParts>,
pub image: ImageModel,
memory_encoder: MemoryEncoder,
memory_attention: MemoryAttention,
temporal_pos: Vec<f32>,
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> {
VideoModel::load_with(w, Backend::Cpu)
}
/// On the GPU when there is one (see [`Backend::Auto`]).
pub fn load_auto(w: &Weights) -> Result<VideoModel> {
VideoModel::load_with(w, Backend::Auto)
}
pub fn load_with(w: &Weights, backend: Backend) -> Result<VideoModel> {
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::<f32>()?,
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<GpuParts> {
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<Encoded> {
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<Memory> {
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<Tensor> {
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> {
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<FrameMask>,
}
#[derive(Default)]
struct ObjectState {
cond: BTreeMap<usize, FrameOut>,
non_cond: BTreeMap<usize, FrameOut>,
}
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<f64>,
}
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<VideoModel>,
frames: Box<dyn FrameSource + 'a>,
objects: BTreeMap<ObjectId, ObjectState>,
cache: VecDeque<(usize, Arc<FrameFeatures>)>,
budget: Budget,
peak: usize,
}
fn sine_embedding_1d(positions: &[f32], dim: usize) -> Result<Tensor> {
let half = dim / 2;
let dim_t: Vec<f32> = (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<f32> {
Ok(t.flatten_all()?.to_vec1::<f32>()?.first().copied().unwrap_or(0.0))
}
impl<'a> VideoTracker<'a> {
pub fn new(model: Arc<VideoModel>, frames: Box<dyn FrameSource + 'a>) -> Result<VideoTracker<'a>> {
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<Arc<FrameFeatures>> {
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<Tensor> {
// 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<f32>, Vec<Tensor>) = (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<f32> = 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<FrameMask> {
// 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<Tensor> { 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::<f32>()?;
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::<f32>()?, 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<usize, FrameOut>| f.abs_diff(frame) <= MAX_POINTERS || cond.keys().any(|c| f.abs_diff(*c) <= MEMORY_FRAMES);
let doomed: Vec<usize> = 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<FrameMask> {
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<FrameMask> {
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::<f32>()?;
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<FrameMask> {
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::<f32>()?, (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::<f32>()?;
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<usize>, 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<usize>,
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<ObjectId> = 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<FrameMask> {
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<Item = (ObjectId, usize)> + '_ {
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::<usize, usize>::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");
}
}
crates/sam2/src/track.rs
+181
View File
@@ -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<Arc<ClipFrame>>;
/// 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<Frames<'a>>,
sizes: Rc<RefCell<HashMap<usize, (usize, usize)>>>,
}
impl FrameSource for Source<'_> {
fn num_frames(&self) -> usize {
self.count
}
fn frame(&self, index: usize) -> Result<Tensor> {
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<TrackedFrame> {
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::<f32>()?;
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<usize, usize> = 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<u8> = 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<Error> = 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<Arc<ClipFrame>>,
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(_))));
}
}
crates/sam2/src/weights.rs
+112
View File
@@ -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<String, Tensor>,
}
impl Weights {
/// Read a `.safetensors` file; rejects non-finite values and unsupported dtypes.
pub fn open(path: &Path) -> Result<Weights> {
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<Weights> {
let map = candle_core::safetensors::load_buffer(bytes, &Device::Cpu).map_err(|e| Error::Weights(e.to_string()))?;
Weights::from_map(map)
}
/// Validate an already-loaded set of tensors.
pub fn from_map(map: HashMap<String, Tensor>) -> Result<Weights> {
let mut out = HashMap::with_capacity(map.len());
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::<f32>()?;
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<Vec<usize>> {
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<Tensor> {
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<String, Tensor> {
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();
}
}
crates/sam2/tests/it/figure.rs
+77
View File
@@ -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<bool>) {
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::<f32>().unwrap();
(iou(&m, &gt), out.iou.flatten_all().unwrap().to_vec1::<f32>().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}");
}
crates/sam2/tests/it/main.rs
+4
View File
@@ -0,0 +1,4 @@
//! The crate's integration tests, one binary (each module is a former `tests/<name>.rs`).
mod figure;
mod reference;
crates/sam2/tests/it/reference.rs
+805
View File
@@ -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 `<repo>/sam2-tiny`, or the folder in
//! `EFFECTCRAFT_SAM2_DIR`) and the fixture tensors (`<repo>/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<String, Tensor>)> {
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::<f32>().unwrap()
}
/// The largest magnitude in a tensor.
fn peak(t: &Tensor) -> f32 {
t.abs().unwrap().flatten_all().unwrap().max(0).unwrap().to_scalar::<f32>().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::<usize>());
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<usize> = 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::<f32>());
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<String, Tensor>)> {
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::<f32>().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<Tensor> {
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::<f32>().unwrap();
let want_score = r.get(&format!("video.{tag}.{}.score", m.frame)).unwrap().flatten_all().unwrap().to_vec1::<f32>().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<_>>(), 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<_>>(), 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<Tensor> {
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<std::sync::Arc<ClipFrame>> = (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::<Vec<_>>(), (100..100 + n as i64).collect::<Vec<_>>());
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<std::sync::Arc<ClipFrame>> = (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::<Vec<_>>());
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::<f32>().unwrap();
let want_score = r.get(&format!("maskprompt.{tag}.{}.score", m.frame)).unwrap().flatten_all().unwrap().to_vec1::<f32>().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<std::sync::Arc<ClipFrame>> = (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<f32> = (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::<f32>().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::<f32>().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}");
}
crates/segment/README.md
+8 -1
View File
@@ -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:
crates/segment/src/lib.rs
+199 -9
View File
@@ -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<Vec<f32>>;
/// 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<u64> {
None
}
/// What runs the model: the GPU's name or "CPU" (`None`: it does not say).
fn device_name(&self) -> Option<String> {
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<f64> {
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<Vec<(i64, Prompt)>>,
}
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<f32>,
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<Arc<ClipFrame>>,
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<Loaded> {
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<Loader> {
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::<Vec<_>>(),
["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::<Vec<_>>(), ["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);
}
}
}
crates/track/src/roto.rs
+106
View File
@@ -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<u32> {
let mut v: Vec<u32> = 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>) -> FrameSeg {
let Some(first) = parts.first() else { return FrameSeg::default() };
let (w, h) = (first.w, first.h);
let mut parts: Vec<FrameSeg> = 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<f32>, 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
crates/track/src/roto/bench.rs
+447
View File
@@ -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<Image>,
/// Visible-object matte per frame (`w × h`, 0/1); empty when the object is out of view.
pub truth: Vec<Vec<u8>>,
/// Frame the strokes are painted on.
pub base: usize,
pub strokes: Vec<Stroke>,
}
/// 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<f64>,
pub per_frame: Vec<f64>,
/// 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<String>,
/// The most memory (MB) the tracker held for its memory banks and cached frames (trackers that have none: `None`).
pub peak_mb: Option<f64>,
/// The matte of every frame (not serialised).
#[serde(skip)]
pub mattes: Vec<Vec<u8>>,
}
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<u8>) {
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<Stroke> {
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<Image>, truth: Vec<Vec<u8>>) -> std::result::Result<Scenario, String> {
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<Scenario> {
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<u8>], secs: f64) -> Run {
let per_frame: Vec<f64> = 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<u8>]| {
let pairs: Vec<f64> = v.windows(2).filter_map(|w| Some(iou(w.first()?, w.get(1)?))).collect();
pairs.iter().sum::<f64>() / 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::<f64>() / tail.len() as f64)
});
Run {
scenario: sc.name.clone(),
tracker: tracker.to_string(),
mean_iou: per_frame.iter().sum::<f64>() / 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<Option<FrameSeg>> = 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<Vec<u8>> = 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<Run> {
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<Arc<ClipFrame>> = 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<u8>> = 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<Stroke> {
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<Run>) -> Option<Run> {
let mut cur = sc.clone();
let mut first: Option<Run> = 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));
}
}
crates/track/src/roto/tests.rs
+151 -8
View File
@@ -65,7 +65,13 @@ pub(crate) fn frame_at(c: [f64; 2]) -> (Image, Vec<u8>) {
}
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::<Vec<_>>(), 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::<Vec<_>>(), &[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<f32> = (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<f32> = gt.iter().map(|v| f32::from(*v)).collect();
let seg = segment_tracked(&img, &refs, soft, false, &opts, 0.0);
assert!(iou(&seg.matte, &gt) > 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, &gt) > 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, &gt) > 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]);
}
crates/ui-egui/Cargo.toml
+1
View File
@@ -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 }
crates/ui-egui/src/app/viewer_frames.rs
+11 -1
View File
@@ -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 {
crates/ui-egui/src/panels/roto_tool.rs
+172 -4
View File
@@ -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<bool> {
let width = inp.img.width as usize;
if width == 0 || !inp.scale.is_finite() || inp.scale <= 0.0 || inp.offset.iter().any(|v| !v.is_finite()) {
return vec![false; inp.img.data.len()];
}
inp.img
.data
.iter()
.enumerate()
.map(|(i, _)| {
let x = ((i % width) as f64 + 0.5 - inp.offset[0]) / inp.scale;
let y = ((i / width) as f64 + 0.5 - inp.offset[1]) / inp.scale;
if x < 0.0 || y < 0.0 || x >= seg.w as f64 || y >= seg.h as f64 {
return false;
}
(y.floor() as usize)
.checked_mul(seg.w)
.and_then(|row| row.checked_add(x.floor() as usize))
.and_then(|index| seg.matte.get(index))
.is_some_and(|m| *m != 0)
})
.collect()
}
/// A stroke being drawn.
#[derive(Clone, Debug)]
struct Drag {
@@ -94,6 +130,19 @@ pub fn draw_overlay(
let alpha: Vec<f32> = (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<bool>)> = 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(&params, 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<bool>, x: i64, y: i64| {
x >= 0
&& y >= 0
&& x < w as i64
&& y < h as i64
&& (y as usize)
.checked_mul(w)
.and_then(|row| row.checked_add(x as usize))
.and_then(|index| m.get(index))
.copied()
.unwrap_or(false)
};
per_object
.iter()
.rev()
.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(&params);
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(&params).objects();
let current = app.session.state.roto.object;
let known = objects.contains(&current);
let r = Rect::from_min_size(pos2(x, bar.center().y - 10.0), vec2(104.0, 20.0));
let strokes = fx::data(&params).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<bool> = (0..24).map(|i| matches!((i % 6, i / 6), (3 | 4, 1 | 2))).collect();
assert_eq!(object_mask_in_input(&seg, &inp), expected);
}
#[test]
fn confidence_colours_run_from_green_to_red_and_absent_is_slate() {
let (g, a, r, s) = (trust_color(0.99, false), trust_color(0.6, false), trust_color(0.1, false), trust_color(0.99, true));
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);
}
}
crates/ui-egui/src/panels/settings.rs
+36 -1
View File
@@ -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());
}
}
crates/ui-egui/tests/it/main.rs
+3
View File
@@ -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;
crates/ui-egui/tests/it/roto_model_guard.rs
+26
View File
@@ -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<dyn MaskModel>) -> 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);
}
}
crates/ui-egui/tests/it/ui_roto_objects.rs
+100
View File
@@ -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<Vec<f32>> {
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<Arc<ClipFrame>>,
_: &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");
}
crates/ui-egui/tests/it/ui_roto_span_confidence.rs
+119
View File
@@ -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<Vec<f32>> {
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<Arc<ClipFrame>>,
_: &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:?}");
}
}
docs/adr/0002-sam2-roto-tracker.md
+19
View File
@@ -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.
docs/agents.md
+19 -7
View File
@@ -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).
docs/architecture.md
+2 -1
View File
@@ -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)) |
docs/effects.md
+2 -2
View File
@@ -3,7 +3,7 @@
<!-- Generated from the effect registry (crates/effects/src/catalog.rs); do not edit by hand.
Regenerate: UPDATE_DOCS=1 cargo test -p effectcraft-effects --lib effects_doc_is_current -->
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.<category>.<name>`) 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
docs/roto-brush-sam2.md
+209
View File
@@ -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.
tools/requirements-sam2-reference.txt
+4
View File
@@ -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
tools/sam2_reference.py
+85
View File
@@ -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])
tools/sam2_video_reference.py
+136
View File
@@ -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])
xtask/src/main.rs
+1
View File
@@ -50,6 +50,7 @@ const LAYERS: &[(&str, u8)] = &[
("expr", 3),
("export", 3),
("gpu", 3),
("sam2", 3),
("lottie", 3),
("format", 3),
("plugin", 3),