From 6468bfc06bf96ad3a62578e549da927512f987a1 Mon Sep 17 00:00:00 2001 From: Guinther Kovalski Date: Tue, 6 Oct 2026 20:15:04 -0300 Subject: [PATCH] M12: SAM 3 Object and Describe masks (pure Rust, Metal) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AI masks from Meta's SAM 3, run in LightCraft itself: a new `lightcraft-segment` crate ports the Hugging Face transformers implementation (Apache-2.0) to candle, on the GPU through Metal on macOS (CPU elsewhere), with no C/C++ and no `unsafe`. - Object: click to include, ⌥-click to leave a part out (SAM 3's interactive prompt encoder + two-way mask decoder). - Describe: a text prompt selects every instance ("sky", "the red car"; "car, road" merges several) via the CLIP text encoder and DETR detector. - Both as new masks and as Add / Subtract / Intersect components; + / − buttons under the selected mask combine selections. - The segmentation is stored with the mask (288² quantized logits, ~10-30 KB), so renders, exports and the web build never need the model. - A zoomed-in detail pass (background, ~5 s on an M4 Pro) re-segments the area around a selection from a sharper source: 5-10x finer edges on small objects. - Per-selection Edge slider (hard ↔ soft); hovering a mask in the list shows it in red (the always-on overlay is now off by default, O toggles). - The engine feature `sam` (on in the desktop app) gates the model; the image is encoded once per photo in the background, clicks take ~35 ms. - Weights are not bundled (SAM License): tools/install-sam3.sh / .ps1 download and verify facebook/sam3; docs/ai-masks.md. - Clearer errors (missing file, missing model) and a stderr logger in the app (LIGHTCRAFT_LOG=info). candle is pinned at 0.9.2: later releases make candle-core depend on tokenizers with the Oniguruma C library. Every stage is checked against the Python reference (tests/reference.rs, tools/sam3_reference.py): on Metal the backbone matches to 7e-5 relative error, masks to ~5e-5. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01C8BkS7JDqa1EMcjW6spDGA --- Cargo.lock | 712 ++++++++++++++++++++++++++ Cargo.toml | 5 + NOTICE | 6 + apps/lightcraft/Cargo.toml | 2 +- apps/lightcraft/src/main.rs | 35 +- crates/develop/Cargo.toml | 1 + crates/develop/src/lib.rs | 2 + crates/develop/src/segmask.rs | 131 +++++ crates/develop/src/settings.rs | 25 + crates/engine/Cargo.toml | 5 + crates/engine/src/cmd/masks.rs | 141 ++++- crates/engine/src/lib.rs | 4 + crates/engine/src/segment.rs | 588 +++++++++++++++++++++ crates/engine/src/tests.rs | 34 ++ crates/pipeline/src/masks.rs | 141 ++++- crates/segment/Cargo.toml | 24 + crates/segment/examples/sam3_bench.rs | 25 + crates/segment/src/clip.rs | 77 +++ crates/segment/src/detector.rs | 304 +++++++++++ crates/segment/src/lib.rs | 252 +++++++++ crates/segment/src/mask.rs | 126 +++++ crates/segment/src/neck.rs | 73 +++ crates/segment/src/nn.rs | 238 +++++++++ crates/segment/src/tokenizer.rs | 173 +++++++ crates/segment/src/tracker.rs | 249 +++++++++ crates/segment/src/vit.rs | 166 ++++++ crates/segment/src/weights.rs | 108 ++++ crates/segment/tests/reference.rs | 97 ++++ crates/ui-egui/src/lib.rs | 30 ++ crates/ui-egui/src/panels/detail.rs | 45 ++ crates/ui-egui/src/panels/masking.rs | 158 +++++- crates/ui-egui/src/state.rs | 16 +- crates/ui-egui/src/tests_masking.rs | 12 +- docs/ai-masks.md | 103 ++++ docs/parity.md | 7 +- tools/install-sam3.ps1 | 40 ++ tools/install-sam3.sh | 110 ++++ tools/requirements-sam3-reference.txt | 8 + tools/sam3_reference.py | 71 +++ xtask/src/layers.rs | 1 + 40 files changed, 4328 insertions(+), 17 deletions(-) create mode 100644 crates/develop/src/segmask.rs create mode 100644 crates/engine/src/segment.rs create mode 100644 crates/segment/Cargo.toml create mode 100644 crates/segment/examples/sam3_bench.rs create mode 100644 crates/segment/src/clip.rs create mode 100644 crates/segment/src/detector.rs create mode 100644 crates/segment/src/lib.rs create mode 100644 crates/segment/src/mask.rs create mode 100644 crates/segment/src/neck.rs create mode 100644 crates/segment/src/nn.rs create mode 100644 crates/segment/src/tokenizer.rs create mode 100644 crates/segment/src/tracker.rs create mode 100644 crates/segment/src/vit.rs create mode 100644 crates/segment/src/weights.rs create mode 100644 crates/segment/tests/reference.rs create mode 100644 docs/ai-masks.md create mode 100644 tools/install-sam3.ps1 create mode 100755 tools/install-sam3.sh create mode 100644 tools/requirements-sam3-reference.txt create mode 100644 tools/sam3_reference.py diff --git a/Cargo.lock b/Cargo.lock index 1952f5c4..66873930 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -594,6 +594,12 @@ dependencies = [ "no_std_io2", ] +[[package]] +name = "block" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d8c1fef690941d3e7788d328517591fecc684c084084702d6ff1641e993699a" + [[package]] name = "block2" version = "0.5.1" @@ -678,6 +684,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" @@ -741,6 +753,76 @@ 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", + "candle-metal-kernels", + "candle-ug", + "float8", + "gemm 0.19.0", + "half", + "libm", + "memmap2", + "num-traits", + "num_cpus", + "objc2-foundation 0.3.2", + "objc2-metal 0.3.2", + "rand", + "rand_distr", + "rayon", + "safetensors 0.7.0", + "thiserror 2.0.21", + "yoke 0.8.3", + "zip", +] + +[[package]] +name = "candle-metal-kernels" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fdfe9d06de16ce49961e49084e5b79a75a9bdf157246e7c7b6328e87a7aa25d" +dependencies = [ + "half", + "objc2 0.6.4", + "objc2-foundation 0.3.2", + "objc2-metal 0.3.2", + "once_cell", + "thiserror 2.0.21", + "tracing", +] + +[[package]] +name = "candle-nn" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3045fa9e7aef8567d209a27d56b692f60b96f4d0569f4c3011f8ca6715c65e03" +dependencies = [ + "candle-core", + "candle-metal-kernels", + "half", + "libc", + "num-traits", + "objc2-metal 0.3.2", + "rayon", + "safetensors 0.7.0", + "serde", + "thiserror 2.0.21", +] + +[[package]] +name = "candle-ug" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c22d62be69068bf58987a45f690612739d8d2ea1bf508c1b87dc6815a019575d" +dependencies = [ + "ug", + "ug-metal", +] + [[package]] name = "cc" version = "1.5.1" @@ -987,6 +1069,22 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8b14ccef22fc6f5a8f4d7d768562a182c04ce9a3b3157b91390b52ddfdf1a76" +[[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 = "ecolor" version = "0.36.2" @@ -1129,6 +1227,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" @@ -1296,6 +1406,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", + "rand_distr", +] + [[package]] name = "fnv" version = "1.0.7" @@ -1399,6 +1521,244 @@ dependencies = [ "slab", ] +[[package]] +name = "gemm" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab96b703d31950f1aeddded248bc95543c9efc7ac9c4a21fda8703a83ee35451" +dependencies = [ + "dyn-stack", + "gemm-c32 0.18.2", + "gemm-c64 0.18.2", + "gemm-common 0.18.2", + "gemm-f16 0.18.2", + "gemm-f32 0.18.2", + "gemm-f64 0.18.2", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa0673db364b12263d103b68337a68fbecc541d6f6b61ba72fe438654709eacb" +dependencies = [ + "dyn-stack", + "gemm-c32 0.19.0", + "gemm-c64 0.19.0", + "gemm-common 0.19.0", + "gemm-f16 0.19.0", + "gemm-f32 0.19.0", + "gemm-f64 0.19.0", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-c32" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6db9fd9f40421d00eea9dd0770045a5603b8d684654816637732463f4073847" +dependencies = [ + "dyn-stack", + "gemm-common 0.18.2", + "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 0.19.0", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-c64" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfcad8a3d35a43758330b635d02edad980c1e143dc2f21e6fd25f9e4eada8edf" +dependencies = [ + "dyn-stack", + "gemm-common 0.18.2", + "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 0.19.0", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-common" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a352d4a69cbe938b9e2a9cb7a3a63b7e72f9349174a2752a558a8a563510d0f3" +dependencies = [ + "bytemuck", + "dyn-stack", + "half", + "libm", + "num-complex", + "num-traits", + "once_cell", + "paste", + "pulp 0.21.5", + "raw-cpuid", + "rayon", + "seq-macro", + "sysctl", +] + +[[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 0.22.3", + "raw-cpuid", + "rayon", + "seq-macro", + "sysctl", +] + +[[package]] +name = "gemm-f16" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cff95ae3259432f3c3410eaa919033cd03791d81cebd18018393dc147952e109" +dependencies = [ + "dyn-stack", + "gemm-common 0.18.2", + "gemm-f32 0.18.2", + "half", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "rayon", + "seq-macro", +] + +[[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 0.19.0", + "gemm-f32 0.19.0", + "half", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "rayon", + "seq-macro", +] + +[[package]] +name = "gemm-f32" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc8d3d4385393304f407392f754cd2dc4b315d05063f62cf09f47b58de276864" +dependencies = [ + "dyn-stack", + "gemm-common 0.18.2", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "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 0.19.0", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-f64" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35b2a4f76ce4b8b16eadc11ccf2e083252d8237c1b589558a49b0183545015bd" +dependencies = [ + "dyn-stack", + "gemm-common 0.18.2", + "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 0.19.0", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + [[package]] name = "gethostname" version = "1.1.0" @@ -1582,9 +1942,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", + "rand_distr", "zerocopy", ] @@ -1609,6 +1972,8 @@ dependencies = [ "allocator-api2", "equivalent", "foldhash", + "serde", + "serde_core", ] [[package]] @@ -1620,6 +1985,12 @@ dependencies = [ "foldhash", ] +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + [[package]] name = "hermit-abi" version = "0.5.3" @@ -2158,6 +2529,7 @@ version = "0.2.1" dependencies = [ "lightcraft-color", "lightcraft-geom", + "miniz_oxide 0.8.9", "proptest", "serde", "serde_json", @@ -2182,6 +2554,7 @@ dependencies = [ "lightcraft-raster", "lightcraft-raw", "lightcraft-scenes", + "lightcraft-segment", "log", "miniz_oxide 0.8.9", "rayon", @@ -2324,6 +2697,18 @@ dependencies = [ "serde", ] +[[package]] +name = "lightcraft-segment" +version = "0.2.1" +dependencies = [ + "candle-core", + "candle-nn", + "log", + "serde", + "serde_json", + "thiserror 2.0.21", +] + [[package]] name = "lightcraft-sysmem" version = "0.2.1" @@ -2426,6 +2811,15 @@ dependencies = [ "imgref", ] +[[package]] +name = "malloc_buf" +version = "0.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62bb907fe88d54d8d9ce32a3cceab4218ed2f6b7d35617cafe9adf84e43919cb" +dependencies = [ + "libc", +] + [[package]] name = "maybe-rayon" version = "0.1.1" @@ -2449,6 +2843,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" dependencies = [ "libc", + "stable_deref_trait", ] [[package]] @@ -2460,6 +2855,21 @@ dependencies = [ "autocfg", ] +[[package]] +name = "metal" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ecfd3296f8c56b7c1f6fbac3c71cefa9d78ce009850c45000015f206dc7fa21" +dependencies = [ + "bitflags 2.13.2", + "block", + "core-graphics-types", + "foreign-types", + "log", + "objc", + "paste", +] + [[package]] name = "miniz_oxide" version = "0.8.9" @@ -2629,6 +3039,20 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0676bb32a98c1a483ce53e500a81ad9c3d5b3f7c920c28c24e9cb0980d0b5bc8" +[[package]] +name = "num" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" +dependencies = [ + "num-bigint", + "num-complex", + "num-integer", + "num-iter", + "num-rational", + "num-traits", +] + [[package]] name = "num-bigint" version = "0.4.8" @@ -2639,6 +3063,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "bytemuck", + "num-traits", +] + [[package]] name = "num-derive" version = "0.4.2" @@ -2659,6 +3093,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-iter" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c92800bd69a1eac91786bcfe9da64a897eb72911b8dc3095decbd07429e8048b" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-rational" version = "0.4.2" @@ -2680,6 +3124,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" @@ -2702,6 +3156,15 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "objc" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "915b1b472bc21c53464d6c8461c9d3af805ba1ef837e1cac254428f4a77177b1" +dependencies = [ + "malloc_buf", +] + [[package]] name = "objc-sys" version = "0.3.5" @@ -2868,6 +3331,7 @@ checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272" dependencies = [ "bitflags 2.13.2", "block2 0.6.2", + "libc", "objc2 0.6.4", "objc2-core-foundation", ] @@ -2915,7 +3379,9 @@ checksum = "a0125f776a10d00af4152d74616409f0d4a2053a6f57fa5b7d6aa2854ac04794" dependencies = [ "bitflags 2.13.2", "block2 0.6.2", + "dispatch2", "objc2 0.6.4", + "objc2-core-foundation", "objc2-foundation 0.3.2", ] @@ -3352,6 +3818,43 @@ dependencies = [ "unarray", ] +[[package]] +name = "pulp" +version = "0.21.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96b86df24f0a7ddd5e4b95c94fc9ed8a98f1ca94d3b01bdce2824097e7835907" +dependencies = [ + "bytemuck", + "cfg-if", + "libm", + "num-complex", + "reborrow", + "version_check", +] + +[[package]] +name = "pulp" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "046aa45b989642ec2e4717c8e72d677b13edd831a4d3b6cf37d9a3e54912496a" +dependencies = [ + "bytemuck", + "cfg-if", + "libm", + "num-complex", + "paste", + "pulp-wasm-simd-flag", + "raw-cpuid", + "reborrow", + "version_check", +] + +[[package]] +name = "pulp-wasm-simd-flag" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d8f70e07b9c3962945a74e59ca1c511bba65b6419468acc217c457d93f3c740" + [[package]] name = "pxfm" version = "0.1.30" @@ -3438,6 +3941,16 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_distr" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" +dependencies = [ + "num-traits", + "rand", +] + [[package]] name = "rand_xorshift" version = "0.4.0" @@ -3503,6 +4016,15 @@ dependencies = [ "rgb", ] +[[package]] +name = "raw-cpuid" +version = "11.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "498cd0dc59d73224351ee52a95fee0f1a617a2eae0e7d9d720cc622c73a54186" +dependencies = [ + "bitflags 2.13.2", +] + [[package]] name = "raw-window-handle" version = "0.6.2" @@ -3552,6 +4074,12 @@ dependencies = [ "once_cell", ] +[[package]] +name = "reborrow" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03251193000f4bd3b042892be858ee50e8b3719f2b08e5833ac4353724632430" + [[package]] name = "redox_syscall" version = "0.4.1" @@ -3695,6 +4223,27 @@ dependencies = [ "wait-timeout", ] +[[package]] +name = "safetensors" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44560c11236a6130a46ce36c836a62936dc81ebf8c36a37947423571be0e55b6" +dependencies = [ + "serde", + "serde_json", +] + +[[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" @@ -3728,6 +4277,12 @@ version = "1.0.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" +[[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" @@ -3990,6 +4545,42 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "synstructure" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "901704edd0dfe137f1987838ee4f259e4e063c31371bdb423f7ae38ec6f77f02" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.6", +] + +[[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 = "tempfile" version = "3.27.0" @@ -4179,6 +4770,12 @@ dependencies = [ "rustc-hash 2.1.3", ] +[[package]] +name = "typed-path" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e28f89b80c87b8fb0cf04ab448d5dd0dd0ade2f8891bae878de66a75a28600e" + [[package]] name = "uds_windows" version = "1.2.1" @@ -4190,6 +4787,41 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "ug" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76b761acf8af3494640d826a8609e2265e19778fb43306c7f15379c78c9b05b0" +dependencies = [ + "gemm 0.18.2", + "half", + "libloading", + "memmap2", + "num", + "num-traits", + "num_cpus", + "rayon", + "safetensors 0.4.5", + "serde", + "thiserror 1.0.69", + "tracing", + "yoke 0.7.5", +] + +[[package]] +name = "ug-metal" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7adf545a99a086d362efc739e7cf4317c18cbeda22706000fd434d70ea3d95" +dependencies = [ + "half", + "metal", + "objc", + "serde", + "thiserror 1.0.69", + "ug", +] + [[package]] name = "unarray" version = "0.1.4" @@ -5141,6 +5773,53 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7a5a4b21e1a62b67a2970e6831bc091d7b87e119e7f9791aef9702e3bef04448" +[[package]] +name = "yoke" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "120e6aef9aa629e3d4f52dc8cc43a015c7724194c97dfaf45180d2daf2b77f40" +dependencies = [ + "serde", + "stable_deref_trait", + "yoke-derive 0.7.5", + "zerofrom", +] + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive 0.8.4", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "synstructure 0.13.2", +] + +[[package]] +name = "yoke-derive" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec8ebde2db3681e8c9980cc27822030e68752690ddfa9473e739aeb4dbde6d71" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.6", + "synstructure 0.14.0", +] + [[package]] name = "zbus" version = "5.19.0" @@ -5267,6 +5946,39 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f75b4683f6c7f45248d4d64056a24298c6281e0993356d7d1b4a1a962ef10d4a" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.6", + "synstructure 0.14.0", +] + +[[package]] +name = "zip" +version = "7.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c42e33efc22a0650c311c2ef19115ce232583abbe80850bc8b66509ebef02de0" +dependencies = [ + "crc32fast", + "indexmap", + "memchr", + "typed-path", +] + [[package]] name = "zlib-rs" version = "0.6.8" diff --git a/Cargo.toml b/Cargo.toml index b1882f9b..d59f8f45 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,6 +29,7 @@ lightcraft-merge = { path = "crates/merge" } lightcraft-engine = { path = "crates/engine" } lightcraft-ui-egui = { path = "crates/ui-egui" } lightcraft-mcp = { path = "crates/mcp" } +lightcraft-segment = { path = "crates/segment" } # foundation thiserror = "2" @@ -62,6 +63,10 @@ jxl-oxide = { version = "0.12", default-features = false } ravif = { version = "0.13", default-features = false, features = ["threading"] } # no "asm": it needs nasm (non-Rust) # gpu (crates/gpu; native backends only, no GL: everything stays pure Rust) +# SAM 3 inference (lightcraft-segment). 0.9.2: later releases pull `tokenizers` with the +# Oniguruma C library into candle-core, and the product is pure Rust. +candle-core = { version = "=0.9.2", default-features = false } +candle-nn = { version = "=0.9.2", default-features = false } wgpu = { version = "30", default-features = false, features = ["std", "parking_lot", "metal", "vulkan", "dx12", "wgsl"] } pollster = "1" diff --git a/NOTICE b/NOTICE index e6f95a8a..1fcf11e7 100644 --- a/NOTICE +++ b/NOTICE @@ -15,6 +15,12 @@ list (author, source, license) is in assets/ATTRIBUTION.md. Inter (assets/fonts/Inter-*.ttf) — SIL Open Font License 1.1. See assets/fonts/OFL-Inter.txt. Showcase photographs by Ansel Adams, Dorothea Lange and NASA (docs/images/) — public domain. See assets/ATTRIBUTION.md. +crates/segment ports the SAM 3 model code of Hugging Face Transformers +(https://github.com/huggingface/transformers, src/transformers/models/sam3 and sam3_tracker), +Copyright The HuggingFace Team and Meta Platforms, Inc., licensed under the Apache License 2.0. +The SAM 3 weights are not part of LightCraft: users download the facebook/sam3 checkpoint +themselves, under Meta's SAM License. + Builds made with craft-fonts (the optional CRAFT_FONTS_DIR build input; all official releases) also embed the fonts listed in https://github.com/storytold/craft-fonts/blob/main/ATTRIBUTION.md (the Japanese fonts BIZ UDPGothic, BIZ UDMincho and Shippori Mincho; the web build embeds only BIZ diff --git a/apps/lightcraft/Cargo.toml b/apps/lightcraft/Cargo.toml index 46acbb73..a29299da 100644 --- a/apps/lightcraft/Cargo.toml +++ b/apps/lightcraft/Cargo.toml @@ -10,7 +10,7 @@ default-run = "lightcraft" workspace = true [dependencies] -lightcraft-engine = { workspace = true } +lightcraft-engine = { workspace = true, features = ["sam"] } lightcraft-sysmem = { workspace = true } lightcraft-ui-egui = { workspace = true } lightcraft-codecs = { workspace = true } diff --git a/apps/lightcraft/src/main.rs b/apps/lightcraft/src/main.rs index 2bdfa16f..26a4fa1f 100644 --- a/apps/lightcraft/src/main.rs +++ b/apps/lightcraft/src/main.rs @@ -415,10 +415,40 @@ OPTIONS: ENVIRONMENT: LIGHTCRAFT_GPU_BACKEND=dx12|vulkan|metal|auto|off graphics backend (default: DX12 on Windows, Metal on macOS, Vulkan on Linux; off = render on the CPU); else WGPU_BACKEND. LIGHTCRAFT_GPU=0: CPU rendering. + LIGHTCRAFT_SAM3_DIR=DIR the SAM 3 model for Object / Describe masks (default: /models/sam3) "; +/// Warnings and errors (failed commands, AI mask analysis) on stderr; `LIGHTCRAFT_LOG=info` +/// (or `debug`) shows more. +struct StderrLog(log::LevelFilter); + +impl log::Log for StderrLog { + fn enabled(&self, m: &log::Metadata) -> bool { + m.level() <= self.0 && (m.level() <= log::Level::Warn || m.target().starts_with("lightcraft")) + } + fn log(&self, r: &log::Record) { + if self.enabled(r.metadata()) { + eprintln!("lightcraft: {} {}: {}", r.level(), r.target(), r.args()); + } + } + fn flush(&self) {} +} + +fn install_log() { + let level = match std::env::var("LIGHTCRAFT_LOG").unwrap_or_default().as_str() { + "debug" => log::LevelFilter::Debug, + "info" => log::LevelFilter::Info, + _ => log::LevelFilter::Warn, + }; + static LOGGER: std::sync::OnceLock = std::sync::OnceLock::new(); + if log::set_logger(LOGGER.get_or_init(|| StderrLog(level))).is_ok() { + log::set_max_level(level); + } +} + fn main() -> eframe::Result { lightcraft_engine::guard::install_hook(std::env::temp_dir().join("lightcraft-panics.log")); + install_log(); alloc_release::install(); let mut control_port: Option = std::env::var("LIGHTCRAFT_CONTROL_PORT").ok().and_then(|p| p.parse().ok()); let mut files = Vec::new(); @@ -481,7 +511,10 @@ fn main() -> eframe::Result { "LightCraft", options, Box::new(move |cc| { - let (session, problem) = open_session(in_memory, library_dir, seed_demo && files.is_empty()); + let (mut session, problem) = open_session(in_memory, library_dir, seed_demo && files.is_empty()); + // AI masks: the SAM 3 checkpoint (facebook/sam3) in /models/sam3, or LIGHTCRAFT_SAM3_DIR + session.segmenter.dir = + std::env::var_os("LIGHTCRAFT_SAM3_DIR").map(std::path::PathBuf::from).or_else(|| config_dir().map(|d| d.join("models").join("sam3"))); let mut app = LightcraftApp::new(session, services()); if let Some(ui) = prefs { app.ui = ui; diff --git a/crates/develop/Cargo.toml b/crates/develop/Cargo.toml index c42aaaa3..6d606148 100644 --- a/crates/develop/Cargo.toml +++ b/crates/develop/Cargo.toml @@ -13,6 +13,7 @@ lightcraft-geom = { workspace = true } lightcraft-color = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +miniz_oxide = { workspace = true } [dev-dependencies] proptest = { workspace = true } diff --git a/crates/develop/src/lib.rs b/crates/develop/src/lib.rs index 72660afd..ff4daedd 100644 --- a/crates/develop/src/lib.rs +++ b/crates/develop/src/lib.rs @@ -8,10 +8,12 @@ pub mod controls; pub mod presets; +pub mod segmask; pub mod settings; pub use controls::{CONTROLS, ControlSpec, Section, Track}; pub use presets::{Preset, SettingsGroup, apply_partial, extract_groups}; +pub use segmask::SegMask; pub use settings::*; use serde_json::Value; diff --git a/crates/develop/src/segmask.rs b/crates/develop/src/segmask.rs new file mode 100644 index 00000000..724ebdf5 --- /dev/null +++ b/crates/develop/src/segmask.rs @@ -0,0 +1,131 @@ +//! Model segmentations stored with a mask ([`SegMask`]), so rendering never needs the model. + +use serde::{Deserialize, Serialize}; + +/// Logits are stored in steps of 1/`SCALE`, clamped to ±127/`SCALE` (≈ ±15.9: a sigmoid of +/// 1 − 1.2e-7, i.e. fully in or out). +const SCALE: f32 = 8.0; +/// Largest grid a stored segmentation may claim (hostile input caps the allocation). +pub const MAX_SIDE: usize = 1024; + +/// A segmentation computed by a model (SAM 3 in `lightcraft-segment`): `side × side` logits +/// over the uncropped, oriented image stretched to a square (row `y`, column `x` covers +/// normalized `((x + 0.5) / side, (y + 0.5) / side)`), quantized to `i8`, deflated and base64 +/// encoded. Coverage at a point is the sigmoid of the bilinearly sampled logit, so edges stay +/// smooth at any output size. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct SegMask { + pub side: u32, + pub data: String, + /// The part of the image the grid covers, normalized `[x0, y0, x1, y1]` (`None` = all of + /// it). A zoomed-in pass over one object covers just its surroundings, at a higher + /// resolution than a pass over the whole photo. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub rect: Option<[f64; 4]>, +} + +impl SegMask { + /// From `side × side` logits over the whole image. + pub fn from_logits(side: usize, logits: &[f32]) -> SegMask { + let q: Vec = logits.iter().take(side * side).map(|l| ((l * SCALE).round().clamp(-127.0, 127.0) as i8) as u8).collect(); + let z = miniz_oxide::deflate::compress_to_vec(&q, 8); + SegMask { side: side as u32, data: base64_encode(&z), rect: None } + } + + /// From `side × side` logits over the part `rect` of the image. + pub fn from_logits_in(side: usize, logits: &[f32], rect: [f64; 4]) -> SegMask { + SegMask { rect: Some(rect), ..SegMask::from_logits(side, logits) } + } + + /// The covered part, normalized (the whole image without a `rect`); `None` when damaged. + pub fn bounds(&self) -> Option<[f64; 4]> { + let r = self.rect.unwrap_or([0.0, 0.0, 1.0, 1.0]); + (r.iter().all(|v| v.is_finite()) && r[2] > r[0] && r[3] > r[1]).then_some(r) + } + + /// The logits, or `None` when the data is damaged or claims an absurd size. + pub fn logits(&self) -> Option> { + let side = self.side as usize; + if side == 0 || side > MAX_SIDE { + return None; + } + let z = base64_decode(&self.data)?; + let q = miniz_oxide::inflate::decompress_to_vec_with_limit(&z, side * side).ok()?; + (q.len() == side * side).then(|| q.iter().map(|b| f32::from(*b as i8) / SCALE).collect()) + } +} + +const B64: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; + +fn base64_encode(bytes: &[u8]) -> String { + let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4); + for chunk in bytes.chunks(3) { + let b = [chunk.first().copied().unwrap_or(0), chunk.get(1).copied().unwrap_or(0), chunk.get(2).copied().unwrap_or(0)]; + let n = (u32::from(b[0]) << 16) | (u32::from(b[1]) << 8) | u32::from(b[2]); + for i in 0..4 { + if i <= chunk.len() { + out.push(char::from(B64[((n >> (18 - 6 * i)) & 63) as usize])); + } else { + out.push('='); + } + } + } + out +} + +fn base64_decode(s: &str) -> Option> { + let mut out = Vec::with_capacity(s.len() / 4 * 3); + let (mut acc, mut bits) = (0u32, 0u32); + for c in s.bytes().filter(|c| !c.is_ascii_whitespace()) { + if c == b'=' { + break; + } + let v = B64.iter().position(|b| *b == c)? as u32; + acc = (acc << 6) | v; + bits += 6; + if bits >= 8 { + bits -= 8; + out.push((acc >> bits) as u8); + acc &= (1 << bits) - 1; + } + } + Some(out) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn base64_round_trips() { + for n in 0..10 { + let v: Vec = (0..n).map(|i| (i * 37 + 11) as u8).collect(); + assert_eq!(base64_decode(&base64_encode(&v)).unwrap(), v); + } + assert_eq!(base64_encode(b"Man"), "TWFu"); + assert_eq!(base64_encode(b"Ma"), "TWE="); + assert!(base64_decode("TW!u").is_none()); + } + + #[test] + fn logits_round_trip_quantized() { + let side = 16; + let l: Vec = (0..side * side).map(|i| (i as f32 - 128.0) / 4.0).collect(); + let m = SegMask::from_logits(side, &l); + let back = m.logits().unwrap(); + for (a, b) in l.iter().zip(&back) { + assert!((a.clamp(-127.0 / SCALE, 127.0 / SCALE) - b).abs() <= 0.5 / SCALE + 1e-6); + } + } + + #[test] + fn damaged_data_is_none() { + let mut m = SegMask::from_logits(8, &[1.0; 64]); + m.side = 9; + assert!(m.logits().is_none(), "size mismatch"); + m.side = 100_000; + assert!(m.logits().is_none(), "absurd size"); + assert!(SegMask { side: 8, data: "not base64 !".into(), rect: None }.logits().is_none()); + assert!(SegMask { rect: Some([0.5, 0.0, 0.2, 1.0]), ..SegMask::from_logits(2, &[0.0; 4]) }.bounds().is_none()); + } +} diff --git a/crates/develop/src/settings.rs b/crates/develop/src/settings.rs index eeb6ff04..736586fe 100644 --- a/crates/develop/src/settings.rs +++ b/crates/develop/src/settings.rs @@ -772,8 +772,33 @@ pub enum MaskShape { Subject, Sky, Background, + /// One object picked by clicks (SAM 3 point prompts): `hint` holds the clicks that include, + /// `exclude` the ones that exclude; `seg` the segmentation computed from them. Object { hint: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + exclude: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + seg: Option, + /// Zoomed-in passes over parts of the image (higher resolution than `seg`), used + /// inside their rectangles. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + detail: Vec, + /// Edge −100..100: below 0 harder (a steeper transition), above 0 softer (feathered). + #[serde(default, skip_serializing_if = "is_zero")] + edge: f64, + }, + /// Everything a description names ("sky", "the red car": SAM 3 concept prompts); `seg` is + /// the segmentation computed for it. + Prompt { + text: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + seg: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + detail: Vec, + /// Edge −100..100, as for `Object`. + #[serde(default, skip_serializing_if = "is_zero")] + edge: f64, }, People { person: u32, diff --git a/crates/engine/Cargo.toml b/crates/engine/Cargo.toml index 5c8e3e46..c2d4e8a8 100644 --- a/crates/engine/Cargo.toml +++ b/crates/engine/Cargo.toml @@ -30,6 +30,11 @@ serde_json = { workspace = true } thiserror = { workspace = true } log = { workspace = true } ab_glyph = { workspace = true } +lightcraft-segment = { workspace = true, optional = true } + +[features] +# AI masks with SAM 3 (Object / Describe). Off for the web and CLI builds. +sam = ["dep:lightcraft-segment"] [dev-dependencies] libc = { workspace = true } diff --git a/crates/engine/src/cmd/masks.rs b/crates/engine/src/cmd/masks.rs index ec3ffac0..da96f967 100644 --- a/crates/engine/src/cmd/masks.rs +++ b/crates/engine/src/cmd/masks.rs @@ -5,6 +5,7 @@ use lightcraft_geom::Point; use serde_json::{Value, json}; use super::{CommandSpec, bad, bool_or, cmd, f64_or, has_active, ok, point, str_param}; +use crate::segment::Click; use crate::{Result, Session}; fn shape_from(kind: &str, p: &Value, c: &str) -> Result { @@ -38,10 +39,57 @@ fn shape_from(kind: &str, p: &Value, c: &str) -> Result { .unwrap_or_default(); MaskShape::ColorRange { samples, refine: f64_or(p, "refine", 50.0) } } - other => return Err(bad(c, format!("unknown mask kind `{other}` (brush|linear|radial|sky|subject|background|luminanceRange|colorRange)"))), + "object" => MaskShape::Object { hint: points(p, "points"), exclude: points(p, "exclude"), seg: None, detail: vec![], edge: 0.0 }, + "prompt" => MaskShape::Prompt { text: str_param(p, "text").unwrap_or_default().trim().to_string(), seg: None, detail: vec![], edge: 0.0 }, + other => { + return Err(bad( + c, + format!("unknown mask kind `{other}` (brush|linear|radial|sky|subject|background|luminanceRange|colorRange|object|prompt)"), + )); + } }) } +/// `[[x, y], …]` normalized points under `key`. +fn points(p: &Value, key: &str) -> Vec { + p.get(key) + .and_then(Value::as_array) + .map(|a| a.iter().filter_map(|q| Some(Point::new(q.get(0)?.as_f64()?, q.get(1)?.as_f64()?))).collect()) + .unwrap_or_default() +} + +/// Compute the segmentation of a new AI shape (Object with clicks, Prompt): an Object +/// without clicks stays empty until the photo is clicked; a Prompt that matches nothing is +/// an error (nothing to select). +fn resolve_ai(s: &mut Session, shape: &mut MaskShape, c: &str) -> Result<()> { + let id = s.active().ok_or_else(|| bad(c, "no active photo"))?; + match shape { + MaskShape::Object { hint, exclude, seg, .. } if !hint.is_empty() => { + *seg = Some(s.segment_clicks(id, &clicks_of(hint, exclude)).map_err(ai_err)?); + } + MaskShape::Object { .. } => { + // the clicks come later: start analyzing the photo now, and refuse an Object mask that + // could never be computed (no model, missing file) rather than leave an empty one + s.segment_prepare(id).map_err(ai_err)?; + } + MaskShape::Prompt { text, seg, .. } => { + let found = s.segment_text(id, text).map_err(ai_err)?; + *seg = Some(found.ok_or_else(|| ai_err(format!("nothing matching “{text}” was found in this photo")))?); + } + _ => {} + } + Ok(()) +} + +/// AI mask failures are shown as they are (not as "invalid parameters"). +fn ai_err(msg: impl Into) -> crate::EngineError { + crate::EngineError::Other(msg.into()) +} + +fn clicks_of(include: &[Point], exclude: &[Point]) -> Vec { + include.iter().map(|p| Click { at: *p, include: true }).chain(exclude.iter().map(|p| Click { at: *p, include: false })).collect() +} + fn masks_edit(s: &mut Session, c: &str, label: &str, f: impl FnOnce(&mut Vec, &mut Option) -> Result<()>) -> Result { let id = s.active().ok_or_else(|| bad(c, "no active photo"))?; let mut d = (*s.develop_of(id).unwrap_or_default()).clone(); @@ -67,11 +115,12 @@ pub fn specs() -> Vec { "Create New Mask", [], None, - "{kind: brush|linear|radial|sky|subject|background|luminanceRange|colorRange, ...shape params (start/end, center/rx/ry/angle/feather, lo/hi…), name?}", + "{kind: brush|linear|radial|sky|subject|background|luminanceRange|colorRange|object|prompt, ...shape params (start/end, center/rx/ry/angle/feather, lo/hi…; object: points/exclude [[x,y],…]; prompt: text), name?}", has_active, |s, p| { let kind = str_param(p, "kind").unwrap_or("radial").to_string(); - let shape = shape_from(&kind, p, "mask.add")?; + let mut shape = shape_from(&kind, p, "mask.add")?; + resolve_ai(s, &mut shape, "mask.add")?; let name = str_param(p, "name").map(str::to_string); let next = s.active().and_then(|id| s.develop_of(id)).map(|d| d.next_mask_id()).unwrap_or(1); masks_edit(s, "mask.add", "Add Mask", |masks, active| { @@ -95,7 +144,8 @@ pub fn specs() -> Vec { has_active, |s, p| { let kind = str_param(p, "kind").unwrap_or("brush").to_string(); - let shape = shape_from(&kind, p, "mask.addComponent")?; + let mut shape = shape_from(&kind, p, "mask.addComponent")?; + resolve_ai(s, &mut shape, "mask.addComponent")?; let op: MaskOp = serde_json::from_value(p.get("op").cloned().unwrap_or(json!("add"))).map_err(|e| bad("mask.addComponent", e.to_string()))?; let active = s.active_mask; @@ -213,6 +263,89 @@ pub fn specs() -> Vec { Ok(json!({"samples": out})) } ), + cmd!( + "mask.objectPoint", + "Add Object Click", + [], + None, + "{x, y: normalized image point, exclude?: bool (⌥-click: leave this part out), id?: maskId} — a click on the selected mask's Object selection (its last Object component); SAM 3 re-segments the object → {include, exclude}", + has_active, + |s, p| { + let c = "mask.objectPoint"; + let at = Point::new(super::f64_req(p, "x", c)?, super::f64_req(p, "y", c)?); + if !(0.0..=1.0).contains(&at.x) || !(0.0..=1.0).contains(&at.y) { + return Err(bad(c, "the point is outside the photo")); + } + let exclude = bool_or(p, "exclude", false); + let id = s.active().ok_or_else(|| bad(c, "no active photo"))?; + let mid = mask_id(p, s.active_mask, c)?; + let d = s.develop_of(id).unwrap_or_default(); + let m = d.masks.iter().find(|m| m.id == mid).ok_or_else(|| bad(c, format!("no mask {mid}")))?; + let k = m + .components + .iter() + .rposition(|x| matches!(x.shape, MaskShape::Object { .. })) + .ok_or_else(|| bad(c, "the mask has no Object selection (add one first)"))?; + let Some(MaskShape::Object { hint, exclude: ex, edge, .. }) = m.components.get(k).map(|x| &x.shape) else { + return Err(bad(c, "not an Object component")); + }; + let (mut hint, mut ex, edge) = (hint.clone(), ex.clone(), *edge); + if exclude { + ex.push(at) + } else { + hint.push(at) + } + let seg = if hint.is_empty() { None } else { Some(s.segment_clicks(id, &clicks_of(&hint, &ex)).map_err(ai_err)?) }; + let out = json!({"include": hint.len(), "exclude": ex.len()}); + masks_edit(s, c, "Object Mask", |masks, _| { + let i = find(masks, mid, c)?; + if let Some(comp) = masks[i].components.get_mut(k) { + comp.shape = MaskShape::Object { hint, exclude: ex, seg, detail: vec![], edge }; + } + Ok(()) + })?; + Ok(out) + } + ), + cmd!( + "mask.refineDetail", + "Refine AI Mask Detail", + [], + None, + "{id?: maskId, component?} — a zoomed-in SAM 3 pass over an Object/Describe selection (default: the mask's last one): the photo around it is analyzed again at a higher resolution, in the background; the mask updates when it's done → {started}", + has_active, + |s, p| { + let c = "mask.refineDetail"; + let id = s.active().ok_or_else(|| bad(c, "no active photo"))?; + let mid = mask_id(p, s.active_mask, c)?; + let d = s.develop_of(id).unwrap_or_default(); + let m = d.masks.iter().find(|m| m.id == mid).ok_or_else(|| bad(c, format!("no mask {mid}")))?; + let k = match p.get("component").and_then(Value::as_u64) { + Some(k) => k as usize, + None => m + .components + .iter() + .rposition(|x| matches!(x.shape, MaskShape::Object { .. } | MaskShape::Prompt { .. })) + .ok_or_else(|| bad(c, "the mask has no Object or Describe selection"))?, + }; + let started = s.segment_detail(id, mid, k).map_err(ai_err)?; + Ok(json!({"started": started})) + } + ), + cmd!( + query "segment.prepare", + "Prepare AI Masks", + [], + None, + "{} — load SAM 3 and analyze the active photo in the background, so Object clicks are instant → {busy}", + has_active, + |s, _| { + let c = "segment.prepare"; + let id = s.active().ok_or_else(|| bad(c, "no active photo"))?; + s.segment_prepare(id).map_err(ai_err)?; + Ok(json!({"busy": s.segmenter.busy()})) + } + ), cmd!( "mask.refine", "Refine Mask Edges", diff --git a/crates/engine/src/lib.rs b/crates/engine/src/lib.rs index 68200a35..40305c5f 100644 --- a/crates/engine/src/lib.rs +++ b/crates/engine/src/lib.rs @@ -32,6 +32,7 @@ pub mod preset_import; pub mod preset_luminar; pub mod presets; pub mod rename; +pub mod segment; pub mod sidecar; pub mod smart; mod view; @@ -150,6 +151,8 @@ pub struct Session { depth: u32, /// Selected mask (Masking panel), by mask id. pub active_mask: Option, + /// AI masks (SAM 3): the model and the last photo prepared for it. + pub segmenter: segment::Segmenter, /// Selected spot (Remove panel), by index into the active photo's spots. pub active_spot: Option, /// The persistent library this session writes to (`None` = in-memory only). @@ -237,6 +240,7 @@ impl Session { clock: Box::new(|| "2026-09-30T12:00:00".to_string()), depth: 0, active_mask: None, + segmenter: segment::Segmenter::default(), active_spot: None, library: None, xmp: sidecar::XmpPrefs::default(), diff --git a/crates/engine/src/segment.rs b/crates/engine/src/segment.rs new file mode 100644 index 00000000..41aed09e --- /dev/null +++ b/crates/engine/src/segment.rs @@ -0,0 +1,588 @@ +//! AI masks: Object (clicks) and Describe (text) selections computed with SAM 3 +//! (`lightcraft-segment`, cargo feature `sam`; the desktop app enables it). +//! +//! The model sees the photo as developed, uncropped and without masks, at 1008 px — the frame +//! mask coordinates live in — so a stored segmentation lines up with the photo at any size. +//! Encoding the image is the slow part (a few seconds on a laptop GPU); it runs once per photo +//! and look, on a worker thread when [`Session::segment_prepare`] is called (the Masking panel +//! does this when an AI mask is started), and each click then takes milliseconds. +//! +//! Without the feature (web, CLI) every request reports that this build has no model; stored +//! segmentations still render (they are plain data in the develop settings). + +use std::path::PathBuf; + +use lightcraft_catalog::PhotoId; +use lightcraft_develop::{MaskShape, SegMask}; +use lightcraft_geom::Point; + +use crate::Session; + +/// Long edge of the image the model sees. +pub const INPUT_EDGE: usize = 1008; + +/// One click of an Object selection (normalized image coordinates). +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Click { + pub at: Point, + pub include: bool, +} + +/// What a request runs on: the photo, its look (settings without masks), the model input. +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +struct Input { + key: u64, + rgb: Vec, + w: usize, + h: usize, +} + +#[cfg(feature = "sam")] +mod model { + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::{Arc, Mutex, PoisonError}; + + use lightcraft_segment::{Encoded, Sam3}; + + #[derive(Default)] + pub struct Inner { + pub model: Option, + /// The last encoded image and its key. + pub cache: Option<(u64, Encoded)>, + } + + #[derive(Clone, Default)] + pub struct Shared { + pub inner: Arc>, + pub busy: Arc, + /// A zoomed-in detail pass is running. + pub detail_busy: Arc, + /// Finished detail passes, waiting for [`crate::Session::segment_poll`]. + pub results: Arc>>, + } + + impl Shared { + pub fn lock(&self) -> std::sync::MutexGuard<'_, Inner> { + self.inner.lock().unwrap_or_else(PoisonError::into_inner) + } + + pub fn busy(&self) -> bool { + self.busy.load(Ordering::Relaxed) + } + + pub fn set_busy(&self, v: bool) { + self.busy.store(v, Ordering::Relaxed); + } + } +} + +/// The model (loaded on first use) and the last encoded photo. +#[derive(Default)] +pub struct Segmenter { + /// Where the `facebook/sam3` checkpoint lives (set by the app). + pub dir: Option, + #[cfg(feature = "sam")] + shared: model::Shared, +} + +impl Segmenter { + /// Whether this build can compute AI masks at all. + pub const AVAILABLE: bool = cfg!(feature = "sam"); + + /// Whether a worker is preparing a photo right now. + pub fn busy(&self) -> bool { + #[cfg(feature = "sam")] + return self.shared.busy(); + #[cfg(not(feature = "sam"))] + false + } + + /// Whether a zoomed-in detail pass is running. + pub fn detail_busy(&self) -> bool { + #[cfg(feature = "sam")] + return self.shared.detail_busy.load(std::sync::atomic::Ordering::Relaxed); + #[cfg(not(feature = "sam"))] + false + } + + /// The model directory, if its files are there. + pub fn model_dir(&self) -> Result { + if !Self::AVAILABLE { + return Err("AI masks are not available in this build".into()); + } + let dir = self.dir.clone().ok_or("no folder is set for the SAM 3 model")?; + #[cfg(feature = "sam")] + if !lightcraft_segment::is_model_dir(&dir) { + return Err(format!( + "the SAM 3 model is not installed: put the facebook/sam3 files (model.safetensors, vocab.json, merges.txt) in {}", + dir.display() + )); + } + Ok(dir) + } +} + +/// Mix `v` into a 64-bit hash. +fn mix(h: u64, v: u64) -> u64 { + (h ^ v).wrapping_mul(0x0100_0000_01b3).rotate_left(29) +} + +/// A finished zoomed-in pass: for which component (and the shape it was computed from, so a +/// component changed meanwhile is left alone), and the patch (`None`: nothing found there). +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +pub struct DetailResult { + photo: PhotoId, + mask: u32, + comp: usize, + shape: MaskShape, + patch: Option, +} + +/// What a detail pass re-runs. +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +enum Prompt { + Clicks(Vec), + Text(String), +} + +/// Objects covering more than this much of the photo gain little from a zoomed pass. +const DETAIL_MAX_AREA: f64 = 0.4; + +/// The part of the photo to zoom into for `seg`'s selection: its bounding box with a margin of +/// a quarter of its size, at least 12 % of the long edge, at most 2:1, inside the photo +/// (`aspect` = width / height). `None` when nothing is selected or it is most of the photo. +pub fn detail_region(seg: &SegMask, aspect: f64) -> Option<[f64; 4]> { + let logits = seg.logits()?; + let r = seg.bounds()?; + let side = seg.side as usize; + let (mut x0, mut y0, mut x1, mut y1) = (usize::MAX, usize::MAX, 0, 0); + for (i, l) in logits.iter().enumerate() { + if *l > 0.0 { + let (x, y) = (i % side, i / side); + (x0, y0, x1, y1) = (x0.min(x), y0.min(y), x1.max(x + 1), y1.max(y + 1)); + } + } + if x0 == usize::MAX { + return None; + } + let (rw, rh) = (r[2] - r[0], r[3] - r[1]); + let nx = |c: usize| r[0] + c as f64 / side as f64 * rw; + let ny = |c: usize| r[1] + c as f64 / side as f64 * rh; + let (bx0, by0, bx1, by1) = (nx(x0), ny(y0), nx(x1), ny(y1)); + if (bx1 - bx0) * (by1 - by0) > DETAIL_MAX_AREA { + return None; + } + // in units of the photo's height: x × aspect + let a = aspect.clamp(0.1, 10.0); + let (cx, cy) = ((bx0 + bx1) / 2.0 * a, (by0 + by1) / 2.0); + let size = ((bx1 - bx0) * a).max(by1 - by0); + let m = 0.25 * size; + let mut w = ((bx1 - bx0) * a + 2.0 * m).max(0.12 * a.max(1.0)); + let mut h = ((by1 - by0) + 2.0 * m).max(0.12 * a.max(1.0)); + w = w.max(h / 2.0).min(a); + h = h.max(w / 2.0).min(1.0); + let x = (cx - w / 2.0).clamp(0.0, (a - w).max(0.0)); + let y = (cy - h / 2.0).clamp(0.0, (1.0 - h).max(0.0)); + let out = [x / a, y, (x + w) / a, y + h]; + // a region nearly the whole photo isn't worth a second pass + ((out[2] - out[0]) * (out[3] - out[1]) < 0.8).then_some(out) +} + +impl Session { + /// The key and settings of photo `id` as the model sees it: its look without masks, uncropped. + fn segment_key(&self, id: PhotoId) -> Option<(u64, lightcraft_develop::DevelopSettings)> { + let p = self.catalog.photo(id)?; + let mut d = (*p.develop).clone(); + d.masks.clear(); + let content = crate::media::content_key(p); + let mut h = mix(0xcbf2_9ce4_8422_2325, d.hash64()); + h = mix(h, id.0); + for b in content.bytes() { + h = mix(h, u64::from(b)); + } + Some((h, d)) + } + + /// Render the model input for `id`. + #[cfg_attr(not(feature = "sam"), allow(dead_code))] + fn segment_input(&mut self, id: PhotoId) -> Result { + let (key, d) = self.segment_key(id).ok_or("no such photo")?; + let job = self.preview_job(id, INPUT_EDGE, INPUT_EDGE, false, &d).ok_or("no such photo")?; + let r = job.run(); + self.accept(&r); + let img = match r.rendered { + Ok(r) => r.image, + Err(e) => { + // the usual reason: the file was moved or deleted (and there's no smart preview) + if let Some(lightcraft_catalog::Source::File { path }) = self.catalog.photo(id).map(|p| &p.source) + && !std::path::Path::new(path).exists() + { + return Err(format!( + "This photo's file is missing (moved or deleted), so it can't be analyzed for AI masks: {path}. Library ▸ Find Missing Photos can relink it." + )); + } + return Err(format!("couldn't render the photo for AI masks: {e}")); + } + }; + let rgb = img.data.iter().flat_map(|p| [p[0], p[1], p[2]]).collect(); + Ok(Input { key, rgb, w: img.width, h: img.height }) + } + + /// Start preparing photo `id` for AI masks on a worker thread (load the model, encode the + /// image). Returns at once; does nothing when it is ready or being prepared. + pub fn segment_prepare(&mut self, id: PhotoId) -> Result<(), String> { + let dir = self.segmenter.model_dir()?; + #[cfg(feature = "sam")] + { + let (key, _) = self.segment_key(id).ok_or("no such photo")?; + let shared = self.segmenter.shared.clone(); + if shared.busy() { + return Ok(()); + } + if let Ok(inner) = shared.inner.try_lock() + && inner.cache.as_ref().is_some_and(|c| c.0 == key) + { + return Ok(()); + } + let input = self.segment_input(id)?; + shared.set_busy(true); + let spawned = std::thread::Builder::new().name("sam3-prepare".into()).spawn(move || { + let mut inner = shared.lock(); + if let Err(e) = prepare(&mut inner, &dir, &input) { + log::warn!("SAM 3: {e}"); + } + shared.set_busy(false); + }); + if let Err(e) = spawned { + self.segmenter.shared.set_busy(false); + return Err(format!("could not start the model: {e}")); + } + } + #[cfg(not(feature = "sam"))] + let _ = (dir, id); + Ok(()) + } + + /// Segment the object `clicks` pick on photo `id` (waits for a running prepare). + pub fn segment_clicks(&mut self, id: PhotoId, clicks: &[Click]) -> Result { + let dir = self.segmenter.model_dir()?; + if !clicks.iter().any(|c| c.include) { + return Err("click the object to include it first".into()); + } + #[cfg(feature = "sam")] + { + let input = self.segment_input_if_needed(id)?; + let shared = self.segmenter.shared.clone(); + let mut inner = shared.lock(); + let enc = encoded(&mut inner, &dir, input)?; + let clicks: Vec = + clicks.iter().map(|c| lightcraft_segment::Click { x: c.at.x as f32, y: c.at.y as f32, positive: c.include }).collect(); + let (model, enc) = enc; + let p = model.segment_clicks(enc, &clicks).map_err(|e| e.to_string())?; + Ok(SegMask::from_logits(lightcraft_segment::MASK_SIDE, &p.logits)) + } + #[cfg(not(feature = "sam"))] + { + let _ = (dir, id); + Err("AI masks are not available in this build".into()) + } + } + + /// Segment everything `text` describes on photo `id`; `Ok(None)` when nothing matches. + pub fn segment_text(&mut self, id: PhotoId, text: &str) -> Result, String> { + let dir = self.segmenter.model_dir()?; + let text = text.trim(); + if text.is_empty() { + return Err("describe what to select".into()); + } + #[cfg(feature = "sam")] + { + let input = self.segment_input_if_needed(id)?; + let shared = self.segmenter.shared.clone(); + let mut inner = shared.lock(); + let (model, enc) = encoded(&mut inner, &dir, input)?; + // "car, road": each phrase on its own, merged (the per-pixel maximum) + let mut merged: Option> = None; + for phrase in text.split([',', ';']).map(str::trim).filter(|p| !p.is_empty()) { + if let Some(p) = model.segment_text(enc, phrase, 0.5).map_err(|e| e.to_string())? { + merged = Some(match merged { + None => p.logits, + Some(m) => m.iter().zip(&p.logits).map(|(a, b)| a.max(*b)).collect(), + }); + } + } + Ok(merged.map(|l| SegMask::from_logits(lightcraft_segment::MASK_SIDE, &l))) + } + #[cfg(not(feature = "sam"))] + { + let _ = (dir, id); + Err("AI masks are not available in this build".into()) + } + } + + /// Render the part `region` (normalized) of photo `id` for the model, from a source sharp + /// enough for it. + #[cfg_attr(not(feature = "sam"), allow(dead_code))] + fn segment_region_input(&mut self, id: PhotoId, region: [f64; 4]) -> Result { + let (_, mut d) = self.segment_key(id).ok_or("no such photo")?; + d.crop.geometry = lightcraft_geom::CropGeometry { rect: lightcraft_geom::Rect::new(region[0], region[1], region[2], region[3]), angle: 0.0 }; + d.crop.flip_h = false; + d.crop.flip_v = false; + d.geometry.constrain_crop = false; + let mut job = self.preview_job(id, INPUT_EDGE, INPUT_EDGE, true, &d).ok_or("no such photo")?; + let frac = (region[2] - region[0]).max(region[3] - region[1]).max(0.01); + let p = self.catalog.photo(id).ok_or("no such photo")?.clone(); + job.level = crate::media::SourceLevel::for_size((INPUT_EDGE as f64 / frac) as usize); + job.source = self.media.source_ref(&p, job.level); + let r = job.run(); + self.accept(&r); + let img = r.rendered.map_err(|e| format!("couldn't render the photo for AI masks: {e}"))?.image; + let rgb = img.data.iter().flat_map(|p| [p[0], p[1], p[2]]).collect(); + Ok(Input { key: 0, rgb, w: img.width, h: img.height }) + } + + /// Start a zoomed-in pass for component `comp` of mask `mask` on photo `id` (an Object or + /// Describe selection): the photo around the selection is rendered at a higher resolution + /// and segmented again, on a worker thread; [`Session::segment_poll`] applies the result. + /// `Ok(false)` when there is nothing to refine (no selection yet, or it is most of the photo). + pub fn segment_detail(&mut self, id: PhotoId, mask: u32, comp: usize) -> Result { + let dir = self.segmenter.model_dir()?; + let d = self.develop_of(id).ok_or("no such photo")?; + let shape = d.masks.iter().find(|m| m.id == mask).and_then(|m| m.components.get(comp)).map(|c| c.shape.clone()).ok_or("no such mask")?; + let (seg, prompt) = match &shape { + MaskShape::Object { hint, exclude, seg: Some(seg), .. } if !hint.is_empty() => ( + seg.clone(), + Prompt::Clicks( + hint.iter().map(|p| Click { at: *p, include: true }).chain(exclude.iter().map(|p| Click { at: *p, include: false })).collect(), + ), + ), + MaskShape::Prompt { text, seg: Some(seg), .. } => (seg.clone(), Prompt::Text(text.clone())), + _ => return Ok(false), + }; + let p = self.catalog.photo(id).ok_or("no such photo")?; + let (pw, ph) = (f64::from(p.width.max(1)), f64::from(p.height.max(1))); + let aspect = if d.orientation.swaps_axes() { ph / pw } else { pw / ph }; + let Some(region) = detail_region(&seg, aspect) else { return Ok(false) }; + #[cfg(feature = "sam")] + { + let shared = self.segmenter.shared.clone(); + if shared.detail_busy.swap(true, std::sync::atomic::Ordering::Relaxed) { + return Ok(false); + } + let input = match self.segment_region_input(id, region) { + Ok(i) => i, + Err(e) => { + shared.detail_busy.store(false, std::sync::atomic::Ordering::Relaxed); + return Err(e); + } + }; + let shape = strip_detail(shape); + let spawned = std::thread::Builder::new().name("sam3-detail".into()).spawn({ + let shared = shared.clone(); + move || { + let patch = { + let mut inner = shared.lock(); + detail_pass(&mut inner, &dir, &input, &prompt, region) + }; + match patch { + Ok(patch) => { + let r = DetailResult { photo: id, mask, comp, shape, patch }; + shared.results.lock().unwrap_or_else(std::sync::PoisonError::into_inner).push(r); + } + Err(e) => log::warn!("SAM 3 detail pass: {e}"), + } + shared.detail_busy.store(false, std::sync::atomic::Ordering::Relaxed); + } + }); + if let Err(e) = spawned { + shared.detail_busy.store(false, std::sync::atomic::Ordering::Relaxed); + return Err(format!("could not start the detail pass: {e}")); + } + Ok(true) + } + #[cfg(not(feature = "sam"))] + { + let _ = (dir, prompt, region, mask); + Ok(false) + } + } + + /// Apply finished detail passes (to components that haven't changed since). True when a + /// mask changed. + pub fn segment_poll(&mut self) -> bool { + #[cfg(feature = "sam")] + { + let done: Vec = + std::mem::take(&mut *self.segmenter.shared.results.lock().unwrap_or_else(std::sync::PoisonError::into_inner)); + let mut changed = false; + for r in done { + let Some(patch) = r.patch else { continue }; + let Some(d) = self.develop_of(r.photo) else { continue }; + let mut d = (*d).clone(); + let Some(c) = d.masks.iter_mut().find(|m| m.id == r.mask).and_then(|m| m.components.get_mut(r.comp)) else { continue }; + if strip_detail(c.shape.clone()) != r.shape { + continue; // clicked again meanwhile: a newer pass will follow + } + match &mut c.shape { + MaskShape::Object { detail, .. } | MaskShape::Prompt { detail, .. } => *detail = vec![patch], + _ => continue, + } + match self.set_develop(r.photo, d, "AI Mask Detail") { + Ok(()) => changed = true, + Err(e) => log::warn!("AI mask detail: {e}"), + } + } + changed + } + #[cfg(not(feature = "sam"))] + false + } + + /// The model input for `id`, unless the cache already holds it (or a prepare is running, + /// which will have it). + #[cfg(feature = "sam")] + fn segment_input_if_needed(&mut self, id: PhotoId) -> Result { + let (key, _) = self.segment_key(id).ok_or("no such photo")?; + if self.segmenter.shared.busy() { + return Ok(Either::Key(key)); + } + let cached = self.segmenter.shared.lock().cache.as_ref().is_some_and(|c| c.0 == key); + if cached { Ok(Either::Key(key)) } else { Ok(Either::Input(self.segment_input(id)?)) } + } +} + +#[cfg(feature = "sam")] +enum Either { + /// Already encoded (or being encoded) under this key. + Key(u64), + Input(Input), +} + +#[cfg(feature = "sam")] +fn load(inner: &mut model::Inner, dir: &std::path::Path) -> Result<(), String> { + if inner.model.is_none() { + let t = web_time::Instant::now(); + inner.model = Some(lightcraft_segment::Sam3::load(dir).map_err(|e| e.to_string())?); + log::info!("SAM 3 loaded in {:?}", t.elapsed()); + } + Ok(()) +} + +/// Load the model and encode `input` into the cache (with the click features ready). +#[cfg(feature = "sam")] +fn prepare(inner: &mut model::Inner, dir: &std::path::Path, input: &Input) -> Result<(), String> { + load(inner, dir)?; + let model = inner.model.as_mut().ok_or("the model did not load")?; + let t = web_time::Instant::now(); + let mut enc = model.encode(&input.rgb, input.w, input.h).map_err(|e| e.to_string())?; + model.prepare_clicks(&mut enc).map_err(|e| e.to_string())?; + log::info!("SAM 3 encoded {}×{} in {:?}", input.w, input.h, t.elapsed()); + inner.cache = Some((input.key, enc)); + Ok(()) +} + +/// The model and the encoded image for `want` (encoding it now when needed). +#[cfg(feature = "sam")] +fn encoded<'a>( + inner: &'a mut model::Inner, + dir: &std::path::Path, + want: Either, +) -> Result<(&'a mut lightcraft_segment::Sam3, &'a mut lightcraft_segment::Encoded), String> { + match want { + Either::Input(input) => prepare(inner, dir, &input)?, + Either::Key(key) => { + if !inner.cache.as_ref().is_some_and(|c| c.0 == key) { + return Err("the photo changed while it was being prepared; try again".into()); + } + } + } + let model::Inner { model, cache } = inner; + match (model.as_mut(), cache.as_mut()) { + (Some(m), Some((_, e))) => Ok((m, e)), + _ => Err("the model is not ready".into()), + } +} + +/// `shape` without its detail patches (what a detail pass was computed from) +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +fn strip_detail(mut shape: MaskShape) -> MaskShape { + // (nor its edge setting: moving the slider meanwhile keeps the pass) + if let MaskShape::Object { detail, edge, .. } | MaskShape::Prompt { detail, edge, .. } = &mut shape { + detail.clear(); + *edge = 0.0; + } + shape +} + +/// Segment `input` (the photo's part `region`) with `prompt`, as a patch over `region`. +#[cfg(feature = "sam")] +fn detail_pass(inner: &mut model::Inner, dir: &std::path::Path, input: &Input, prompt: &Prompt, region: [f64; 4]) -> Result, String> { + load(inner, dir)?; + let model = inner.model.as_mut().ok_or("the model did not load")?; + let t = web_time::Instant::now(); + let mut enc = model.encode(&input.rgb, input.w, input.h).map_err(|e| e.to_string())?; + let (rw, rh) = (region[2] - region[0], region[3] - region[1]); + let logits = match prompt { + Prompt::Clicks(clicks) => { + let inside: Vec = clicks + .iter() + .map(|c| ((c.at.x - region[0]) / rw, (c.at.y - region[1]) / rh, c.include)) + .filter(|(x, y, _)| (0.0..=1.0).contains(x) && (0.0..=1.0).contains(y)) + .map(|(x, y, positive)| lightcraft_segment::Click { x: x as f32, y: y as f32, positive }) + .collect(); + if !inside.iter().any(|c| c.positive) { + return Ok(None); + } + Some(model.segment_clicks(&mut enc, &inside).map_err(|e| e.to_string())?.logits) + } + Prompt::Text(text) => { + let mut merged: Option> = None; + for phrase in text.split([',', ';']).map(str::trim).filter(|p| !p.is_empty()) { + if let Some(p) = model.segment_text(&mut enc, phrase, 0.5).map_err(|e| e.to_string())? { + merged = Some(match merged { + None => p.logits, + Some(m) => m.iter().zip(&p.logits).map(|(a, b)| a.max(*b)).collect(), + }); + } + } + merged + } + }; + log::info!("SAM 3 detail pass ({}×{}) in {:?}", input.w, input.h, t.elapsed()); + Ok(logits.map(|l| SegMask::from_logits_in(lightcraft_segment::MASK_SIDE, &l, region))) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn seg_box(side: usize, x0: usize, y0: usize, x1: usize, y1: usize) -> SegMask { + let l: Vec = + (0..side * side).map(|i| if (x0..x1).contains(&(i % side)) && (y0..y1).contains(&(i / side)) { 8.0 } else { -8.0 }).collect(); + SegMask::from_logits(side, &l) + } + + #[test] + fn detail_region_frames_the_selection_with_a_margin() { + // a small object in the middle of a square photo + let r = detail_region(&seg_box(100, 40, 45, 50, 55), 1.0).unwrap(); + assert!(r[0] < 0.4 && r[2] > 0.5 && r[1] < 0.45 && r[3] > 0.55, "contains it: {r:?}"); + assert!(r[2] - r[0] < 0.25 && r[3] - r[1] < 0.25, "zoomed in: {r:?}"); + assert!(r.iter().all(|v| (0.0..=1.0).contains(v))); + // at the corner: shifted inside the photo, not cut + let r = detail_region(&seg_box(100, 0, 0, 5, 5), 1.5).unwrap(); + assert!(r[0] >= 0.0 && r[1] >= 0.0 && r[0] < 1e-9 && r[1] < 1e-9); + // a thin object stays at most 2:1 (in photo pixels, 1.5:1 photo) + let r = detail_region(&seg_box(100, 10, 50, 90, 52), 1.5).unwrap(); + let (w, h) = ((r[2] - r[0]) * 1.5, r[3] - r[1]); + assert!(w / h <= 2.0 + 1e-9, "{w} × {h}"); + } + + #[test] + fn no_region_for_nothing_or_most_of_the_photo() { + assert!(detail_region(&seg_box(50, 0, 0, 0, 0), 1.0).is_none()); + assert!(detail_region(&seg_box(50, 2, 2, 48, 48), 1.0).is_none()); + let damaged = SegMask { side: 50, data: "?".into(), rect: None }; + assert!(detail_region(&damaged, 1.0).is_none()); + } +} diff --git a/crates/engine/src/tests.rs b/crates/engine/src/tests.rs index 58b7e643..65104127 100644 --- a/crates/engine/src/tests.rs +++ b/crates/engine/src/tests.rs @@ -727,3 +727,37 @@ fn curve_reset_by_channel_and_whole() { assert_eq!(active_dev(&s).curve, lightcraft_develop::ToneCurve::default()); assert!(s.execute("curve.reset", &json!({"channel": "alpha"})).is_err()); } + +#[test] +fn ai_masks_without_the_model_explain_themselves() { + let mut s = demo(); + let id = s.active().unwrap(); + // an Object mask that could never be computed is refused, not left empty + let e = s.execute("mask.add", &json!({"kind": "object"})).unwrap_err().to_string(); + assert!(e.contains("SAM 3") || e.contains("not available"), "{e}"); + assert!(!e.contains("invalid parameters"), "shown as is: {e}"); + assert!(s.develop_of(id).unwrap().masks.is_empty()); + // an empty Object component (as an older document may hold) still takes clicks + let shape = lightcraft_develop::MaskShape::Object { hint: vec![], exclude: vec![], seg: None, detail: vec![], edge: 0.0 }; + let mut d = (*s.develop_of(id).unwrap()).clone(); + d.masks.push(lightcraft_develop::Mask { + id: 1, + components: vec![lightcraft_develop::MaskComponent { name: None, op: lightcraft_develop::MaskOp::Add, invert: false, shape }], + ..Default::default() + }); + s.set_develop(id, d, "test").unwrap(); + s.active_mask = Some(1); + // clicks and descriptions need the model: an error that says what's missing, nothing changed + let e = s.execute("mask.objectPoint", &json!({"x": 0.5, "y": 0.5})).unwrap_err().to_string(); + assert!(e.contains("SAM 3") || e.contains("not available"), "{e}"); + let e = s.execute("mask.add", &json!({"kind": "prompt", "text": "sky"})).unwrap_err().to_string(); + assert!(e.contains("SAM 3") || e.contains("not available"), "{e}"); + let empty = std::env::temp_dir().join(format!("lc-no-sam3-{}", std::process::id())); + s.segmenter.dir = Some(empty.clone()); + let e = s.execute("mask.objectPoint", &json!({"x": 0.5, "y": 0.5})).unwrap_err().to_string(); + assert!(e.contains(&empty.display().to_string()) || e.contains("not available"), "names the folder: {e}"); + assert_eq!(s.develop_of(id).unwrap().masks.len(), 1); + // bad input is an error, not a crash + assert!(s.execute("mask.objectPoint", &json!({"x": 1.5, "y": 0.5})).is_err()); + assert!(s.execute("mask.add", &json!({"kind": "prompt", "text": " "})).is_err()); +} diff --git a/crates/pipeline/src/masks.rs b/crates/pipeline/src/masks.rs index 5c6e003c..01985b5b 100644 --- a/crates/pipeline/src/masks.rs +++ b/crates/pipeline/src/masks.rs @@ -3,7 +3,7 @@ //! Shapes are defined in normalized oriented-image coordinates; sizes are in "long-edge units". //! Components combine with Add (max), Subtract (a·(1−c)) and Intersect (a·c). -use lightcraft_develop::{BrushStroke, LocalAdjustments, Mask, MaskOp, MaskShape}; +use lightcraft_develop::{BrushStroke, LocalAdjustments, Mask, MaskOp, MaskShape, SegMask}; use lightcraft_geom::Point; use lightcraft_raster::{Plane, Rgb32f}; @@ -153,6 +153,19 @@ pub fn shape_alpha(shape: &MaskShape, frame: &Frame, w: usize, h: usize, img: &R } smooth_plane(&mut out, 0.01 * frame_px(frame, w)); } + MaskShape::Object { seg, detail, edge, .. } | MaskShape::Prompt { seg, detail, edge, .. } if seg.is_some() || !detail.is_empty() => { + // Edge: below 0 a steeper transition (up to 8× the logits), above 0 feathered (up to + // 2 % of the long edge, so previews and exports match) + let e = (*edge / 100.0).clamp(-1.0, 1.0) as f32; + let gain = if e < 0.0 { 1.0 - 7.0 * e } else { 1.0 }; + sample_seg(seg.as_ref(), detail, gain, frame, w, h, &mut out); + if e > 0.0 { + smooth_plane(&mut out, e * 0.02 * frame_px(frame, w)); + } + } + // not computed yet (no clicks, or a build without the model): nothing selected + MaskShape::Object { hint, .. } if hint.is_empty() => {} + MaskShape::Prompt { .. } => {} MaskShape::Subject | MaskShape::Object { .. } | MaskShape::People { .. } => { // Saliency heuristic: centre-weighted local contrast (replaced by the segmenter in M12). let m = frame.out_to_norm(w, h); @@ -176,6 +189,60 @@ pub fn shape_alpha(shape: &MaskShape, frame: &Frame, w: usize, h: usize, img: &R out } +/// A decoded [`SegMask`]: its logits and the normalized rectangle they cover. +struct SegGrid { + logits: Vec, + side: usize, + r: [f64; 4], +} + +impl SegGrid { + fn new(seg: &SegMask) -> Option { + Some(SegGrid { logits: seg.logits()?, side: seg.side as usize, r: seg.bounds()? }) + } + + fn contains(&self, n: Point) -> bool { + n.x >= self.r[0] && n.x <= self.r[2] && n.y >= self.r[1] && n.y <= self.r[3] + } + + /// The bilinearly sampled logit at normalized image point `n` (cell centres at + /// `(i + 0.5) / side` of the rectangle). + fn logit(&self, n: Point) -> f32 { + let side = self.side; + let at = |x: usize, y: usize| self.logits.get(y.min(side - 1) * side + x.min(side - 1)).copied().unwrap_or(-16.0); + let u = (n.x - self.r[0]) / (self.r[2] - self.r[0]); + let v = (n.y - self.r[1]) / (self.r[3] - self.r[1]); + let (fx, fy) = ((u * side as f64 - 0.5) as f32, (v * side as f64 - 0.5) as f32); + let (fx, fy) = (fx.clamp(0.0, (side - 1) as f32), fy.clamp(0.0, (side - 1) as f32)); + let (x0, y0) = (fx.floor() as usize, fy.floor() as usize); + let (tx, ty) = (fx - x0 as f32, fy - y0 as f32); + let top = at(x0, y0) * (1.0 - tx) + at(x0 + 1, y0) * tx; + let bot = at(x0, y0 + 1) * (1.0 - tx) + at(x0 + 1, y0 + 1) * tx; + top * (1.0 - ty) + bot * ty + } +} + +/// Stored model segmentations at each output pixel: the zoomed-in `detail` passes inside +/// their rectangles (the highest of them where they overlap), the whole-photo `seg` +/// elsewhere; the logit (times `gain`: > 1 is a harder edge) through a sigmoid. Damaged +/// data selects nothing. +fn sample_seg(seg: Option<&SegMask>, detail: &[SegMask], gain: f32, frame: &Frame, w: usize, h: usize, out: &mut Plane) { + let coarse = seg.and_then(SegGrid::new); + let patches: Vec = detail.iter().filter_map(SegGrid::new).collect(); + if coarse.is_none() && patches.is_empty() { + return; + } + let m = frame.out_to_norm(w, h); + for_rows(&mut out.data, w, |y, row| { + for (x, v) in row.iter_mut().enumerate() { + let n = m.apply(Point::new(x as f64 + 0.5, y as f64 + 0.5)); + let fine = patches.iter().filter(|p| p.contains(n)).map(|p| p.logit(n)).reduce(f32::max); + let l = fine.or_else(|| coarse.as_ref().map(|c| c.logit(n))).unwrap_or(-16.0); + *v = 1.0 / (1.0 + (-l * gain).exp()); + } + }); +} + fn frame_px(frame: &Frame, w: usize) -> f32 { frame.px_per_long(w) as f32 } @@ -356,6 +423,78 @@ mod tests { assert!((a.get(50, 25) - 0.5).abs() < 0.05); } + #[test] + fn stored_segmentations_render_at_any_size_and_follow_the_crop() { + // the left half selected, on a 16² logit grid + let side = 16; + let logits: Vec = (0..side * side).map(|i| if i % side < side / 2 { 12.0 } else { -12.0 }).collect(); + let seg = SegMask::from_logits(side, &logits); + for shape in [ + MaskShape::Object { hint: vec![Point::new(0.25, 0.5)], exclude: vec![], seg: Some(seg.clone()), detail: vec![], edge: 0.0 }, + MaskShape::Prompt { text: "left".into(), seg: Some(seg.clone()), detail: vec![], edge: 0.0 }, + ] { + for (w, h) in [(40, 20), (400, 200)] { + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + assert!(a.get(1, h / 2) > 0.99 && a.get(w - 2, h / 2) < 0.01, "{w}×{h}"); + if w >= 400 { + let soft = (0..w).filter(|x| (0.05..0.95).contains(&a.get(*x, h / 2))).count(); + assert!((2..30).contains(&soft), "an anti-aliased edge, not a ramp: {soft} px"); + } + } + } + // cropped to the right half: nothing left selected + let mut d = DevelopSettings::default(); + d.crop.geometry.rect = lightcraft_geom::Rect::from_center(Point::new(0.75, 0.5), 0.5, 1.0); + let f = Frame::new(40, 20, &d, true); + let shape = MaskShape::Prompt { text: "left".into(), seg: Some(seg), detail: vec![], edge: 0.0 }; + let a = shape_alpha(&shape, &f, 20, 20, &Rgb32f::new(20, 20), &Plane::new(20, 20), 0.0); + assert!(a.data.iter().all(|v| *v < 0.05)); + } + + #[test] + fn detail_patches_replace_the_coarse_mask_inside_their_rectangle() { + // coarse: nothing; a patch over the right half selects its own left half (x 0.5..0.75) + let coarse = SegMask::from_logits(4, &[-12.0; 16]); + let side = 16; + let l: Vec = (0..side * side).map(|i| if i % side < side / 2 { 12.0 } else { -12.0 }).collect(); + let patch = SegMask::from_logits_in(side, &l, [0.5, 0.0, 1.0, 1.0]); + let shape = MaskShape::Object { hint: vec![Point::new(0.6, 0.5)], exclude: vec![], seg: Some(coarse), detail: vec![patch], edge: 0.0 }; + let (w, h) = (200, 100); + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + assert!(a.get(20, 50) < 0.01, "outside the patch: the coarse mask"); + assert!(a.get(120, 50) > 0.99, "inside the patch, its selected half"); + assert!(a.get(190, 50) < 0.01, "inside the patch, its unselected half"); + } + + #[test] + fn the_edge_setting_hardens_or_feathers_the_transition() { + let side = 16; + let l: Vec = (0..side * side).map(|i| ((i % side) as f32 - 7.5) * -0.6).collect(); // a gentle ramp + let seg = SegMask::from_logits(side, &l); + let (w, h) = (400, 40); + let soft_px = |edge: f64| { + let shape = MaskShape::Prompt { text: "x".into(), seg: Some(seg.clone()), detail: vec![], edge }; + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + assert!(a.get(2, h / 2) > 0.9 && a.get(w - 3, h / 2) < 0.1, "still the same selection at {edge}"); + (0..w).filter(|x| (0.1..0.9).contains(&a.get(*x, h / 2))).count() + }; + let (hard, plain, soft) = (soft_px(-100.0), soft_px(0.0), soft_px(100.0)); + assert!(hard < plain && plain < soft, "transition widths {hard} < {plain} < {soft}"); + } + + #[test] + fn ai_shapes_without_a_segmentation_select_nothing() { + let f = frame(30, 30); + let (img, l) = (Rgb32f::new(30, 30), Plane::new(30, 30)); + for shape in [ + MaskShape::Object { hint: vec![], exclude: vec![], seg: None, detail: vec![], edge: 0.0 }, + MaskShape::Prompt { text: "sky".into(), seg: None, detail: vec![], edge: 0.0 }, + MaskShape::Prompt { text: "sky".into(), seg: Some(SegMask { side: 4, data: "damaged!".into(), rect: None }), detail: vec![], edge: 0.0 }, + ] { + assert!(shape_alpha(&shape, &f, 30, 30, &img, &l, 0.0).data.iter().all(|v| *v == 0.0)); + } + } + #[test] fn radial_inside_and_invert() { let f = frame(100, 100); diff --git a/crates/segment/Cargo.toml b/crates/segment/Cargo.toml new file mode 100644 index 00000000..336b052f --- /dev/null +++ b/crates/segment/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "lightcraft-segment" +version.workspace = true +edition.workspace = true +license.workspace = true +description = "LightCraft segmentation: SAM 3 (Segment Anything with Concepts) inference in pure Rust (L3)" + +[lints] +workspace = true + +[dependencies] +serde = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } +log = { workspace = true } + +# the model runs natively only (the web build has no AI masks) +[target.'cfg(not(target_arch = "wasm32"))'.dependencies] +candle-core = { workspace = true } +candle-nn = { workspace = true } + +[target.'cfg(target_os = "macos")'.dependencies] +candle-core = { workspace = true, features = ["metal"] } +candle-nn = { workspace = true, features = ["metal"] } diff --git a/crates/segment/examples/sam3_bench.rs b/crates/segment/examples/sam3_bench.rs new file mode 100644 index 00000000..982f482e --- /dev/null +++ b/crates/segment/examples/sam3_bench.rs @@ -0,0 +1,25 @@ +//! Times SAM 3 on this machine (a synthetic gradient image; the first run includes compiling +//! the GPU kernels): `cargo run -p lightcraft-segment --release --example sam3_bench -- `. + +fn main() -> Result<(), Box> { + let mut args = std::env::args().skip(1); + let dir = std::path::PathBuf::from(args.next().ok_or("usage: sam3_bench ")?); + let t = std::time::Instant::now(); + let mut model = lightcraft_segment::Sam3::load(&dir)?; + eprintln!("load: {:?} on {:?}", t.elapsed(), model.device()); + let (w, h) = (1200, 800); + let rgb: Vec = (0..w * h * 3).map(|i| ((i / 3 % w) * 255 / w) as u8).collect(); + for run in 0..3 { + let t = std::time::Instant::now(); + let mut enc = model.encode(&rgb, w, h)?; + lightcraft_segment::sync(&model)?; + let e = t.elapsed(); + let t = std::time::Instant::now(); + let p = model.segment_clicks(&mut enc, &[lightcraft_segment::Click { x: 0.5, y: 0.5, positive: true }])?; + let c = t.elapsed(); + let t = std::time::Instant::now(); + let d = model.segment_text(&mut enc, "sky", 0.5)?; + eprintln!("run {run}: encode {e:?}, click {c:?} (score {:.2}), text {:?} ({})", p.score, t.elapsed(), d.is_some()); + } + Ok(()) +} diff --git a/crates/segment/src/clip.rs b/crates/segment/src/clip.rs new file mode 100644 index 00000000..c392d5c5 --- /dev/null +++ b/crates/segment/src/clip.rs @@ -0,0 +1,77 @@ +//! The CLIP text encoder (24 layers, 1024 wide, causal) that turns a prompt into the token +//! features SAM 3's detector attends to. + +use candle_core::{Device, Module, Tensor}; +use candle_nn::{Embedding, LayerNorm, VarBuilder}; + +use crate::Result; +use crate::nn::{Attention, Mlp, ln}; +use crate::tokenizer::CONTEXT; + +const HIDDEN: usize = 1024; +const LAYERS: usize = 24; +const HEADS: usize = 16; +const EPS: f64 = 1e-5; + +struct Layer { + ln1: LayerNorm, + attn: Attention, + ln2: LayerNorm, + mlp: Mlp, +} + +pub struct TextEncoder { + tokens: Embedding, + positions: Tensor, + layers: Vec, + final_ln: LayerNorm, + device: Device, +} + +impl TextEncoder { + /// `vb` at `detector_model.text_encoder.text_model`. + pub fn new(vb: VarBuilder) -> Result { + let e = vb.pp("embeddings"); + let mut layers = Vec::with_capacity(LAYERS); + for i in 0..LAYERS { + let l = vb.pp(format!("encoder.layers.{i}")); + layers.push(Layer { + ln1: ln(HIDDEN, EPS, l.pp("layer_norm1"))?, + attn: Attention::named(HIDDEN, HIDDEN, HEADS, "out_proj", l.pp("self_attn"))?, + ln2: ln(HIDDEN, EPS, l.pp("layer_norm2"))?, + mlp: Mlp::new(HIDDEN, 4 * HIDDEN, true, l.pp("mlp"))?, + }); + } + Ok(Self { + tokens: candle_nn::embedding(49408, HIDDEN, e.pp("token_embedding"))?, + positions: e.get((CONTEXT, HIDDEN), "position_embedding.weight")?.unsqueeze(0)?, + layers, + final_ln: ln(HIDDEN, EPS, vb.pp("final_layer_norm"))?, + device: vb.device().clone(), + }) + } + + /// `[1, 32, 1024]` features for token `ids` (`valid` marks real tokens). + pub fn forward(&self, ids: &[u32], valid: &[bool]) -> Result { + let ids = Tensor::from_slice(ids, (1, ids.len()), &self.device)?; + let mut x = self.tokens.forward(&ids)?.broadcast_add(&self.positions)?; + // causal, and padding keys masked + let n = valid.len(); + let mut bias = vec![0f32; n * n]; + for (q, row) in bias.chunks_mut(n).enumerate() { + for (k, v) in row.iter_mut().enumerate() { + if k > q || !valid.get(k).copied().unwrap_or(false) { + *v = f32::MIN; + } + } + } + let bias = Tensor::from_vec(bias, (1, 1, n, n), &self.device)?; + for l in &self.layers { + let h = l.ln1.forward(&x)?; + x = (&x + l.attn.forward(&h, &h, &h, Some(&bias))?)?; + let h = l.ln2.forward(&x)?; + x = (&x + l.mlp.forward(&h)?)?; + } + Ok(self.final_ln.forward(&x)?) + } +} diff --git a/crates/segment/src/detector.rs b/crates/segment/src/detector.rs new file mode 100644 index 00000000..b21b32d8 --- /dev/null +++ b/crates/segment/src/detector.rs @@ -0,0 +1,304 @@ +//! Text prompts (SAM 3's concept detector): the prompt's CLIP features fused into the image +//! features by a DETR encoder, 200 object queries refined by a DETR decoder (with box +//! relative-position bias and a presence token), per-query masks from a pixel decoder. + +use candle_core::{D, Device, IndexOp, Module, Tensor}; +use candle_nn::{Conv2d, Conv2dConfig, GroupNorm, LayerNorm, Linear, VarBuilder}; + +use crate::clip::TextEncoder; +use crate::neck::{FPN, Neck}; +use crate::nn::{Attention, DecoderMlp, Mlp, TORCH_LN_EPS, cxcywh_to_xyxy, inverse_sigmoid, linear, ln, padding_bias, sine_grid, to_vec}; +use crate::{Error, Prediction, Result}; + +const HIDDEN: usize = FPN; +const HEADS: usize = 8; +const QUERIES: usize = 200; +const LAYERS: usize = 6; + +/// Detector image features, computed once per image. +pub struct Features { + /// Levels 0..=2 (288, 144, 72 per side). + levels: Vec, + /// Level 2 flattened `[1, 5184, 256]` and its sine position encoding. + flat: Tensor, + pos: Tensor, + side: usize, +} + +/// Everything the detector predicts for one prompt. +pub struct Output { + /// Per-query mask logits `[200, 288, 288]` flattened. + pub masks: Tensor, + /// Per-query logits (before the presence score). + pub logits: Vec, + /// Whether the concept is in the image at all (logit). + pub presence: f32, + /// Boxes (x0, y0, x1, y1), normalized. + pub boxes: Vec<[f32; 4]>, + /// The text features `[1, 32, 256]` (tests). + pub text: Tensor, +} + +struct EncoderLayer { + ln1: LayerNorm, + self_attn: Attention, + ln2: LayerNorm, + cross_attn: Attention, + ln3: LayerNorm, + mlp: Mlp, +} + +struct DecoderLayer { + self_attn: Attention, + self_ln: LayerNorm, + text_attn: Attention, + text_ln: LayerNorm, + vision_attn: Attention, + vision_ln: LayerNorm, + mlp: Mlp, + mlp_ln: LayerNorm, +} + +pub struct Detector { + device: Device, + neck: Neck, + text: TextEncoder, + text_proj: Linear, + encoder: Vec, + decoder: Vec, + out_ln: LayerNorm, + box_head: DecoderMlp, + query_embed: Tensor, + reference_points: Tensor, + presence_token: Tensor, + presence_head: DecoderMlp, + presence_ln: LayerNorm, + ref_point_head: DecoderMlp, + rpb_x: DecoderMlp, + rpb_y: DecoderMlp, + // dot-product scoring + score_text_mlp: DecoderMlp, + score_text_ln: LayerNorm, + score_text_proj: Linear, + score_query_proj: Linear, + // mask decoder + prompt_attn: Attention, + prompt_ln: LayerNorm, + pixel_convs: Vec, + pixel_norms: Vec, + instance_proj: Conv2d, + mask_embedder: Vec, +} + +impl Detector { + /// `vb` at `detector_model`. + pub fn new(vb: VarBuilder) -> Result { + let e = vb.pp("detr_encoder"); + let mut encoder = Vec::with_capacity(LAYERS); + for i in 0..LAYERS { + let l = e.pp(format!("layers.{i}")); + encoder.push(EncoderLayer { + ln1: ln(HIDDEN, TORCH_LN_EPS, l.pp("layer_norm1"))?, + self_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("self_attn"))?, + ln2: ln(HIDDEN, TORCH_LN_EPS, l.pp("layer_norm2"))?, + cross_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("cross_attn"))?, + ln3: ln(HIDDEN, TORCH_LN_EPS, l.pp("layer_norm3"))?, + mlp: Mlp::new(HIDDEN, 2048, false, l.pp("mlp"))?, + }); + } + let d = vb.pp("detr_decoder"); + let mut decoder = Vec::with_capacity(LAYERS); + for i in 0..LAYERS { + let l = d.pp(format!("layers.{i}")); + decoder.push(DecoderLayer { + self_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("self_attn"))?, + self_ln: ln(HIDDEN, TORCH_LN_EPS, l.pp("self_attn_layer_norm"))?, + text_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("text_cross_attn"))?, + text_ln: ln(HIDDEN, TORCH_LN_EPS, l.pp("text_cross_attn_layer_norm"))?, + vision_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("vision_cross_attn"))?, + vision_ln: ln(HIDDEN, TORCH_LN_EPS, l.pp("vision_cross_attn_layer_norm"))?, + mlp: Mlp::new(HIDDEN, 2048, false, l.pp("mlp"))?, + mlp_ln: ln(HIDDEN, TORCH_LN_EPS, l.pp("mlp_layer_norm"))?, + }); + } + let s = vb.pp("dot_product_scoring"); + let m = vb.pp("mask_decoder"); + let p = m.pp("pixel_decoder"); + let conv3 = Conv2dConfig { padding: 1, ..Default::default() }; + Ok(Self { + device: vb.device().clone(), + neck: Neck::new(vb.pp("vision_encoder.neck"))?, + text: TextEncoder::new(vb.pp("text_encoder.text_model"))?, + text_proj: linear(1024, HIDDEN, vb.pp("text_projection"))?, + encoder, + decoder, + out_ln: ln(HIDDEN, TORCH_LN_EPS, d.pp("output_layer_norm"))?, + box_head: DecoderMlp::new(HIDDEN, HIDDEN, 4, 3, d.pp("box_head"))?, + query_embed: d.get((QUERIES, HIDDEN), "query_embed.weight")?.unsqueeze(0)?, + reference_points: d.get((QUERIES, 4), "reference_points.weight")?.unsqueeze(0)?, + presence_token: d.get((1, HIDDEN), "presence_token.weight")?.unsqueeze(0)?, + presence_head: DecoderMlp::new(HIDDEN, HIDDEN, 1, 3, d.pp("presence_head"))?, + presence_ln: ln(HIDDEN, TORCH_LN_EPS, d.pp("presence_layer_norm"))?, + ref_point_head: DecoderMlp::new(2 * HIDDEN, HIDDEN, HIDDEN, 2, d.pp("ref_point_head"))?, + rpb_x: DecoderMlp::new(2, HIDDEN, HEADS, 2, d.pp("box_rpb_embed_x"))?, + rpb_y: DecoderMlp::new(2, HIDDEN, HEADS, 2, d.pp("box_rpb_embed_y"))?, + score_text_mlp: DecoderMlp::new(HIDDEN, 2048, HIDDEN, 2, s.pp("text_mlp"))?, + score_text_ln: ln(HIDDEN, TORCH_LN_EPS, s.pp("text_mlp_out_norm"))?, + score_text_proj: linear(HIDDEN, HIDDEN, s.pp("text_proj"))?, + score_query_proj: linear(HIDDEN, HIDDEN, s.pp("query_proj"))?, + prompt_attn: Attention::new(HIDDEN, HIDDEN, HEADS, m.pp("prompt_cross_attn"))?, + prompt_ln: ln(HIDDEN, TORCH_LN_EPS, m.pp("prompt_cross_attn_norm"))?, + pixel_convs: (0..3).map(|i| Ok(candle_nn::conv2d(HIDDEN, HIDDEN, 3, conv3, p.pp(format!("conv_layers.{i}")))?)).collect::>()?, + pixel_norms: (0..3).map(|i| Ok(candle_nn::group_norm(8, HIDDEN, 1e-5, p.pp(format!("norms.{i}")))?)).collect::>()?, + instance_proj: candle_nn::conv2d(HIDDEN, HIDDEN, 1, Conv2dConfig::default(), m.pp("instance_projection"))?, + mask_embedder: (0..3).map(|i| linear(HIDDEN, HIDDEN, m.pp(format!("mask_embedder.layers.{i}")))).collect::>()?, + }) + } + + /// The detector's feature pyramid for the backbone output `[1, 1024, 72, 72]`. + pub fn features(&self, backbone: &Tensor) -> Result { + let levels = self.neck.forward(backbone, 3)?; + let last = levels.last().ok_or_else(|| Error::Model("detector neck".into()))?; + let (_, _, side, _) = last.dims4()?; + let flat = last.flatten_from(2)?.transpose(1, 2)?.contiguous()?; + let pos = sine_grid(side, side, HIDDEN / 2, &self.device)?.flatten_from(2)?.transpose(1, 2)?.contiguous()?; + Ok(Features { levels, flat, pos, side }) + } + + /// Box relative-position bias `[1, heads, 1 + 200, side²]` for reference boxes (cx, cy, w, h). + fn rpb(&self, boxes: &Tensor, side: usize) -> Result { + let xyxy = cxcywh_to_xyxy(boxes)?.squeeze(0)?; // [200, 4] + let coords = Tensor::arange(0u32, side as u32, &self.device)?.to_dtype(candle_core::DType::F32)?.affine(1.0 / side as f64, 0.0)?; + let coords = coords.reshape((1, side, 1))?; + let log_scale = |d: Tensor| -> Result { + let d = (d * 8.0)?; + Ok(d.sign()?.mul(&((d.abs()? + 1.0)?.log()? / (8f64.ln()))?)?) + }; + let pick = |a: usize, b: usize| -> Result { Ok(Tensor::cat(&[xyxy.narrow(1, a, 1)?, xyxy.narrow(1, b, 1)?], 1)?.unsqueeze(1)?) }; + let dy = log_scale(coords.broadcast_sub(&pick(1, 3)?)?)?; // [200, side, 2] + let dx = log_scale(coords.broadcast_sub(&pick(0, 2)?)?)?; + let ey = self.rpb_y.forward(&dy)?; // [200, side, heads] + let ex = self.rpb_x.forward(&dx)?; + let m = ey.unsqueeze(2)?.broadcast_add(&ex.unsqueeze(1)?)?; // [200, side, side, heads] + let m = m.reshape((QUERIES, side * side, HEADS))?.permute((2, 0, 1))?; // [heads, 200, side²] + let m = m.pad_with_zeros(1, 1, 0)?; // the presence token attends everywhere + Ok(m.unsqueeze(0)?.contiguous()?) + } + + /// Sine embedding of boxes (cx, cy, w, h) `[1, 200, 512]` (y, x, w, h features). + fn box_sine(&self, boxes: &Tensor) -> Result { + let feats = HIDDEN / 2; + let dim_t: Vec = (0..feats).map(|i| 10000f64.powf(2.0 * (i / 2) as f64 / feats as f64) as f32).collect(); + let dim_t = Tensor::from_vec(dim_t, (1, 1, feats), &self.device)?; + let even: Vec = (0..feats).map(|i| if i % 2 == 0 { 1.0 } else { 0.0 }).collect(); + let even = Tensor::from_vec(even, (1, 1, feats), &self.device)?; + let one = |i: usize| -> Result { + let p = (boxes.narrow(D::Minus1, i, 1)? * (2.0 * std::f64::consts::PI))?.broadcast_div(&dim_t)?; + Ok((p.sin()?.broadcast_mul(&even)? + p.cos()?.broadcast_mul(&(1.0 - &even)?)?)?) + }; + Ok(Tensor::cat(&[one(1)?, one(0)?, one(2)?, one(3)?], D::Minus1)?) + } + + /// Run the detector for one tokenized prompt. + pub fn predict(&self, f: &Features, ids: &[u32], valid: &[bool]) -> Result { + let text = self.text_proj.forward(&self.text.forward(ids, valid)?)?; // [1, 32, 256] + let text_bias = padding_bias(valid, &self.device)?; + + // DETR encoder: image tokens attend to themselves and to the prompt + let mut x = f.flat.clone(); + for l in &self.encoder { + let h = l.ln1.forward(&x)?; + let hp = (&h + &f.pos)?; + x = (&x + l.self_attn.forward(&hp, &hp, &h, None)?)?; + let h = l.ln2.forward(&x)?; + x = (&x + l.cross_attn.forward(&h, &text, &text, Some(&text_bias))?)?; + let h = l.ln3.forward(&x)?; + x = (&x + l.mlp.forward(&h)?)?; + } + let memory = x; + + // DETR decoder: presence token + 200 queries with iterative box refinement + let mut boxes = candle_nn::ops::sigmoid(&self.reference_points)?; + let mut h = Tensor::cat(&[self.presence_token.clone(), self.query_embed.clone()], 1)?; + let key_pos = (&memory + &f.pos)?; + let mut queries = None; + let mut presence = None; + for l in &self.decoder { + let qpos = self.ref_point_head.forward(&self.box_sine(&boxes)?)?.pad_with_zeros(1, 1, 0)?; + let bias = self.rpb(&boxes, f.side)?; + let q = (&h + &qpos)?; + h = l.self_ln.forward(&(&h + l.self_attn.forward(&q, &q, &h, None)?)?)?; + let q = (&h + &qpos)?; + h = l.text_ln.forward(&(&h + l.text_attn.forward(&q, &text, &text, Some(&text_bias))?)?)?; + let q = (&h + &qpos)?; + h = l.vision_ln.forward(&(&h + l.vision_attn.forward(&q, &key_pos, &memory, Some(&bias))?)?)?; + h = l.mlp_ln.forward(&(&h + l.mlp.forward(&h)?)?)?; + let qh = self.out_ln.forward(&h.narrow(1, 1, QUERIES)?)?; + boxes = candle_nn::ops::sigmoid(&(self.box_head.forward(&qh)? + inverse_sigmoid(&boxes)?)?)?; + let p = self.presence_head.forward(&self.presence_ln.forward(&h.narrow(1, 0, 1)?)?)?; + presence = Some(p.clamp(-10f32, 10f32)?); + queries = Some(qh); + } + let (Some(queries), Some(presence)) = (queries, presence) else { return Err(Error::Model("decoder has no layers".into())) }; + + // dot-product scoring against the mean of the (refined) prompt tokens + let t = self.score_text_ln.forward(&(self.score_text_mlp.forward(&text)? + &text)?)?; + let w: Vec = valid.iter().map(|v| if *v { 1.0 } else { 0.0 }).collect(); + let count = w.iter().sum::().max(1.0); + let w = Tensor::from_vec(w, (1, valid.len(), 1), &self.device)?; + let pooled = (t.broadcast_mul(&w)?.sum(1)? / f64::from(count))?; // [1, 256] + let pt = self.score_text_proj.forward(&pooled)?; // [1, 256] + let pq = self.score_query_proj.forward(&queries)?; // [1, 200, 256] + let logits = (pq.matmul(&pt.unsqueeze(2)?)? / (HIDDEN as f64).sqrt())?.clamp(-12f32, 12f32)?; + + // masks: the encoder memory attends to the prompt, then a pixel decoder up to 288² + let mem = (&memory + self.prompt_attn.forward(&self.prompt_ln.forward(&memory)?, &text, &text, Some(&text_bias))?)?; + let [l0, l1, _] = &f.levels[..] else { return Err(Error::Model("detector levels".into())) }; + let mut prev = mem.transpose(1, 2)?.reshape((1, HIDDEN, f.side, f.side))?; + for (i, skip) in [l1, l0].into_iter().enumerate() { + let (_, _, sh, sw) = skip.dims4()?; + prev = (prev.upsample_nearest2d(sh, sw)? + skip)?; + let (Some(conv), Some(norm)) = (self.pixel_convs.get(i), self.pixel_norms.get(i)) else { + return Err(Error::Model("pixel decoder".into())); + }; + prev = norm.forward(&conv.forward(&prev)?)?.relu()?; + } + let inst = self.instance_proj.forward(&prev)?; // [1, 256, 288, 288] + let (_, c, mh, mw) = inst.dims4()?; + let mut e = queries.clone(); + for (i, l) in self.mask_embedder.iter().enumerate() { + e = l.forward(&e)?; + if i + 1 < self.mask_embedder.len() { + e = e.relu()?; + } + } + let masks = e.matmul(&inst.reshape((1, c, mh * mw))?)?.reshape((QUERIES, mh, mw))?; + let boxes = cxcywh_to_xyxy(&boxes)?; + let bv = to_vec(&boxes)?; + Ok(Output { + masks, + logits: to_vec(&logits)?, + presence: to_vec(&presence)?.first().copied().unwrap_or(-10.0), + boxes: bv.as_chunks::<4>().0.to_vec(), + text, + }) + } +} + +/// The instances of `out` scoring above `threshold` (sigmoid(logit) × sigmoid(presence)), +/// merged into one mask (the per-pixel maximum of their probabilities, as logits). +pub fn merge(out: &Output, threshold: f32) -> Option { + let s = |x: f32| 1.0 / (1.0 + (-x).exp()); + let presence = s(out.presence); + let keep: Vec<(usize, f32)> = out.logits.iter().enumerate().map(|(i, l)| (i, s(*l) * presence)).filter(|(_, p)| *p > threshold).collect(); + let best = keep.iter().map(|k| k.1).fold(0f32, f32::max); + let mut merged: Option> = None; + for (i, _) in &keep { + let m = out.masks.i(*i).ok().and_then(|t| to_vec(&t).ok())?; + merged = Some(match merged { + None => m, + Some(acc) => acc.iter().zip(&m).map(|(a, b)| a.max(*b)).collect(), + }); + } + merged.map(|logits| Prediction { logits, score: best }) +} diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs new file mode 100644 index 00000000..83bcbbaf --- /dev/null +++ b/crates/segment/src/lib.rs @@ -0,0 +1,252 @@ +//! SAM 3 (Segment Anything with Concepts, Meta, 2025) inference in pure Rust on candle. +//! +//! - **Point prompts** ([`Sam3::segment_clicks`]): positive and negative clicks select one +//! object, via SAM 3's interactive (SAM 2 style) prompt encoder and mask decoder. +//! - **Text prompts** ([`Sam3::segment_text`]): a noun phrase ("the red car", "sky") selects +//! every instance of that concept, via the CLIP text encoder and the DETR detector. +//! +//! The expensive part, the 32-layer ViT over the 1008 × 1008 input, runs once per image +//! ([`Sam3::encode`]); each click or prompt then reuses the [`Encoded`] features. +//! +//! On macOS the model runs on the GPU through Metal; elsewhere (for now) on the CPU. The +//! weights are not part of LightCraft: they are the `facebook/sam3` checkpoint (SAM License), +//! read from a directory the user downloads them to ([`Sam3::load`]). +//! +//! Ported from the Hugging Face `transformers` implementation (Apache-2.0); see `NOTICE`. +//! No UI dependencies (L3). +#![forbid(unsafe_code)] +#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] + +#[cfg(not(target_arch = "wasm32"))] +mod clip; +#[cfg(not(target_arch = "wasm32"))] +mod detector; +pub mod mask; +#[cfg(not(target_arch = "wasm32"))] +mod neck; +#[cfg(not(target_arch = "wasm32"))] +mod nn; +pub mod tokenizer; +#[cfg(not(target_arch = "wasm32"))] +mod tracker; +#[cfg(not(target_arch = "wasm32"))] +mod vit; +#[cfg(not(target_arch = "wasm32"))] +mod weights; + +use std::path::{Path, PathBuf}; +#[cfg(not(target_arch = "wasm32"))] +use std::sync::Arc; + +#[cfg(not(target_arch = "wasm32"))] +use candle_core::{DType, Device, Tensor}; +#[cfg(not(target_arch = "wasm32"))] +use candle_nn::VarBuilder; + +#[cfg(not(target_arch = "wasm32"))] +pub use tracker::Click; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("SAM 3 model: {0}")] + Model(String), + #[error("SAM 3 model files not found in {0} (expected model.safetensors, vocab.json, merges.txt)")] + Missing(PathBuf), + #[cfg(not(target_arch = "wasm32"))] + #[error(transparent)] + Candle(#[from] candle_core::Error), +} + +pub type Result = std::result::Result; + +/// Side of the square model input. +pub const INPUT: usize = 1008; +/// Side of the low-resolution masks the model predicts. +pub const MASK_SIDE: usize = 288; +#[cfg(not(target_arch = "wasm32"))] +const _: () = assert!(MASK_SIDE == tracker::LOW_RES && INPUT == vit::IMAGE); + +/// The files of a model directory. +pub const WEIGHTS_FILE: &str = "model.safetensors"; + +/// Whether `dir` holds a usable checkpoint. +pub fn is_model_dir(dir: &Path) -> bool { + dir.join(WEIGHTS_FILE).is_file() && dir.join("vocab.json").is_file() && dir.join("merges.txt").is_file() +} + +#[cfg(not(target_arch = "wasm32"))] +/// The best device here: Metal on macOS (when a GPU is available), else the CPU. +pub fn best_device() -> Device { + #[cfg(target_os = "macos")] + if let Ok(d) = Device::new_metal(0) { + return d; + } + Device::Cpu +} + +#[cfg(not(target_arch = "wasm32"))] +/// A loaded model. The vision backbone loads with it; the point and text heads load on first use. +pub struct Sam3 { + dir: PathBuf, + device: Device, + weights: Arc, + backbone: vit::Vit, + tracker_neck: Option, + tracker: Option, + detector: Option, + tokenizer: Option, +} + +#[cfg(not(target_arch = "wasm32"))] +/// An image run through the backbone, with the head-specific features computed on demand. +pub struct Encoded { + backbone: Tensor, + tracker: Option, + detector: Option, +} + +/// A predicted mask: `MASK_SIDE²` logits over the model's square input (which the image was +/// stretched to), and a confidence 0..1. +#[derive(Clone, Debug)] +pub struct Prediction { + pub logits: Vec, + pub score: f32, +} + +#[cfg(not(target_arch = "wasm32"))] +impl Sam3 { + /// Load from a directory holding `model.safetensors` (the `facebook/sam3` checkpoint in + /// Hugging Face format), `vocab.json` and `merges.txt`. + pub fn load(dir: &Path) -> Result { + Self::load_on(dir, best_device()) + } + + pub fn load_on(dir: &Path, device: Device) -> Result { + if !is_model_dir(dir) { + return Err(Error::Missing(dir.to_path_buf())); + } + let weights = weights::Weights::open(&dir.join(WEIGHTS_FILE))?; + let backbone = vit::Vit::new(Self::vb(&weights, &device).pp("detector_model.vision_encoder.backbone"))?; + Ok(Self { dir: dir.to_path_buf(), device, weights, backbone, tracker_neck: None, tracker: None, detector: None, tokenizer: None }) + } + + fn vb<'a>(w: &Arc, device: &Device) -> VarBuilder<'a> { + VarBuilder::from_backend(Box::new(weights::Backend(w.clone())), DType::F32, device.clone()) + } + + pub fn device(&self) -> &Device { + &self.device + } + + /// Run the backbone on an RGB image (8-bit, row-major, `w × h`), stretched to the square + /// input like the reference preprocessing. + pub fn encode(&self, rgb: &[u8], w: usize, h: usize) -> Result { + let pixels = mask::preprocess(rgb, w, h, INPUT)?; + let pixels = Tensor::from_vec(pixels, (1, 3, INPUT, INPUT), &self.device)?; + self.encode_pixels(&pixels) + } + + /// Run the backbone on preprocessed pixels `[1, 3, 1008, 1008]` (−1..1). + pub fn encode_pixels(&self, pixels: &Tensor) -> Result { + Ok(Encoded { backbone: self.backbone.forward(&pixels.to_device(&self.device)?)?, tracker: None, detector: None }) + } + + fn tracker_parts(&mut self) -> Result<(&neck::Neck, &tracker::Tracker)> { + if self.tracker_neck.is_none() { + self.tracker_neck = Some(neck::Neck::new(Self::vb(&self.weights, &self.device).pp("tracker_neck"))?); + } + if self.tracker.is_none() { + self.tracker = Some(tracker::Tracker::new(Self::vb(&self.weights, &self.device).pp("tracker_model"))?); + } + match (&self.tracker_neck, &self.tracker) { + (Some(n), Some(t)) => Ok((n, t)), + _ => Err(Error::Model("tracker failed to load".into())), + } + } + + /// Compute the click decoder's features for `enc` now (so the first click is fast). + pub fn prepare_clicks(&mut self, enc: &mut Encoded) -> Result<()> { + let (neck, tracker) = self.tracker_parts()?; + if enc.tracker.is_none() { + let levels = neck.forward(&enc.backbone, 3)?; + enc.tracker = Some(tracker.features(&levels)?); + } + Ok(()) + } + + /// Segment one object from clicks (positions normalized to the image, 0..1). + pub fn segment_clicks(&mut self, enc: &mut Encoded, clicks: &[Click]) -> Result { + if clicks.is_empty() { + return Err(Error::Model("no clicks".into())); + } + let (neck, tracker) = self.tracker_parts()?; + if enc.tracker.is_none() { + let levels = neck.forward(&enc.backbone, 3)?; + enc.tracker = Some(tracker.features(&levels)?); + } + let f = enc.tracker.as_ref().ok_or_else(|| Error::Model("tracker features".into()))?; + // one click: SAM's three alternatives, best by predicted IoU; more: the single mask + let p = tracker.predict(f, clicks, clicks.len() == 1)?; + Ok(Prediction { logits: p.logits, score: p.iou * sigmoid(p.object_score) }) + } + + fn detector_parts(&mut self) -> Result<(&detector::Detector, &tokenizer::Tokenizer)> { + if self.detector.is_none() { + self.detector = Some(detector::Detector::new(Self::vb(&self.weights, &self.device).pp("detector_model"))?); + } + if self.tokenizer.is_none() { + self.tokenizer = Some(tokenizer::Tokenizer::load(&self.dir)?); + } + match (&self.detector, &self.tokenizer) { + (Some(d), Some(t)) => Ok((d, t)), + _ => Err(Error::Model("detector failed to load".into())), + } + } + + /// Segment every instance of the concept `text` names. Instances scoring above + /// `threshold` (0..1, 0.5 is the reference default) are merged into one mask; `None` when + /// nothing matches. + pub fn segment_text(&mut self, enc: &mut Encoded, text: &str, threshold: f32) -> Result> { + let (det, tok) = self.detector_parts()?; + if enc.detector.is_none() { + enc.detector = Some(det.features(&enc.backbone)?); + } + let f = enc.detector.as_ref().ok_or_else(|| Error::Model("detector features".into()))?; + let (ids, valid) = tok.encode(text); + let out = det.predict(f, &ids, &valid)?; + Ok(detector::merge(&out, threshold)) + } + + /// Raw detector outputs for `text` (tests and tools). + pub fn detect(&mut self, enc: &mut Encoded, text: &str) -> Result { + let (det, tok) = self.detector_parts()?; + if enc.detector.is_none() { + enc.detector = Some(det.features(&enc.backbone)?); + } + let f = enc.detector.as_ref().ok_or_else(|| Error::Model("detector features".into()))?; + let (ids, valid) = tok.encode(text); + det.predict(f, &ids, &valid) + } +} + +#[cfg(not(target_arch = "wasm32"))] +impl Encoded { + /// The backbone output `[1, 1024, 72, 72]` (tests). + pub fn backbone(&self) -> &Tensor { + &self.backbone + } +} + +#[cfg(not(target_arch = "wasm32"))] +fn sigmoid(x: f32) -> f32 { + 1.0 / (1.0 + (-x).exp()) +} + +#[cfg(not(target_arch = "wasm32"))] +pub use detector::Output as DetectorOutput; + +#[cfg(not(target_arch = "wasm32"))] +/// Wait for queued GPU work to finish (benchmarks). +pub fn sync(model: &Sam3) -> Result<()> { + Ok(model.device.synchronize()?) +} diff --git a/crates/segment/src/mask.rs b/crates/segment/src/mask.rs new file mode 100644 index 00000000..1e61c1f7 --- /dev/null +++ b/crates/segment/src/mask.rs @@ -0,0 +1,126 @@ +//! Image preprocessing and mask post-processing (CPU, no model needed). + +use crate::{Error, Result}; + +/// Separable resampling weights for `src → dst` samples with an antialiased triangle filter +/// (Pillow's BILINEAR: the support widens with the downscale factor). +fn weights(src: usize, dst: usize) -> Vec<(usize, Vec)> { + let scale = src as f64 / dst as f64; + let ss = scale.max(1.0); + let support = ss; + (0..dst) + .map(|o| { + let center = (o as f64 + 0.5) * scale; + let lo = ((center - support + 0.5).floor().max(0.0)) as usize; + let hi = ((center + support + 0.5).floor() as usize).min(src); + let mut w: Vec = (lo..hi).map(|i| (1.0 - ((i as f64 - center + 0.5) / ss).abs()).max(0.0) as f32).collect(); + let sum: f32 = w.iter().sum(); + if sum > 0.0 { + w.iter_mut().for_each(|v| *v /= sum); + } + (lo, w) + }) + .collect() +} + +/// An 8-bit RGB image (`w × h`, row-major) stretched to `side × side`, scaled to −1..1 and laid +/// out channel-first (`[3, side, side]`), as SAM 3's processor does. +pub fn preprocess(rgb: &[u8], w: usize, h: usize, side: usize) -> Result> { + let n = w.checked_mul(h).and_then(|n| n.checked_mul(3)).ok_or_else(|| Error::Model("image too large".into()))?; + if w == 0 || h == 0 || rgb.len() < n { + return Err(Error::Model(format!("image buffer of {} bytes for {w} × {h}", rgb.len()))); + } + let (wx, wy) = (weights(w, side), weights(h, side)); + // horizontal pass: h rows × side columns × 3 + let mut tmp = vec![0f32; h * side * 3]; + for y in 0..h { + for (ox, (lo, ws)) in wx.iter().enumerate() { + let mut acc = [0f32; 3]; + for (k, wt) in ws.iter().enumerate() { + let i = (y * w + lo + k) * 3; + for (c, a) in acc.iter_mut().enumerate() { + *a += wt * f32::from(rgb.get(i + c).copied().unwrap_or(0)); + } + } + if let Some(dst) = tmp.get_mut((y * side + ox) * 3..(y * side + ox) * 3 + 3) { + dst.copy_from_slice(&acc); + } + } + } + let mut out = vec![0f32; 3 * side * side]; + for (oy, (lo, ws)) in wy.iter().enumerate() { + for ox in 0..side { + let mut acc = [0f32; 3]; + for (k, wt) in ws.iter().enumerate() { + let i = ((lo + k) * side + ox) * 3; + for (c, a) in acc.iter_mut().enumerate() { + *a += wt * tmp.get(i + c).copied().unwrap_or(0.0); + } + } + for (c, a) in acc.iter().enumerate() { + // Pillow rounds to 8 bits; then /255 and (x − 0.5) / 0.5 + let v = a.round().clamp(0.0, 255.0) / 255.0; + if let Some(d) = out.get_mut((c * side + oy) * side + ox) { + *d = (v - 0.5) / 0.5; + } + } + } + } + Ok(out) +} + +/// Mask `logits` (`side × side`) resampled bilinearly (half-pixel centres) to `w × h` and +/// mapped to 0..1 coverage (sigmoid). +pub fn alpha(logits: &[f32], side: usize, w: usize, h: usize) -> Vec { + let mut out = vec![0f32; w * h]; + if side == 0 || logits.len() < side * side { + return out; + } + let sample = |x: usize, y: usize| logits.get(y.min(side - 1) * side + x.min(side - 1)).copied().unwrap_or(-20.0); + let (sx, sy) = (side as f32 / w.max(1) as f32, side as f32 / h.max(1) as f32); + for y in 0..h { + let fy = ((y as f32 + 0.5) * sy - 0.5).max(0.0); + let (y0, ty) = (fy.floor() as usize, fy.fract()); + for x in 0..w { + let fx = ((x as f32 + 0.5) * sx - 0.5).max(0.0); + let (x0, tx) = (fx.floor() as usize, fx.fract()); + let top = sample(x0, y0) * (1.0 - tx) + sample(x0 + 1, y0) * tx; + let bot = sample(x0, y0 + 1) * (1.0 - tx) + sample(x0 + 1, y0 + 1) * tx; + let l = top * (1.0 - ty) + bot * ty; + if let Some(o) = out.get_mut(y * w + x) { + *o = 1.0 / (1.0 + (-l).exp()); + } + } + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn preprocess_is_bounded_and_channel_first() { + let rgb: Vec = (0..20 * 10 * 3).map(|i| (i % 256) as u8).collect(); + let p = preprocess(&rgb, 20, 10, 8).unwrap(); + assert_eq!(p.len(), 3 * 64); + assert!(p.iter().all(|v| (-1.0..=1.0).contains(v))); + // a flat image stays flat + let flat = vec![255u8; 30 * 30 * 3]; + assert!(preprocess(&flat, 30, 30, 7).unwrap().iter().all(|v| (*v - 1.0).abs() < 1e-6)); + } + + #[test] + fn short_buffers_are_errors() { + assert!(preprocess(&[0; 5], 2, 2, 4).is_err()); + assert!(preprocess(&[], 0, 0, 4).is_err()); + } + + #[test] + fn alpha_follows_the_logits() { + let logits: Vec = (0..16).map(|i| if i % 4 < 2 { 10.0 } else { -10.0 }).collect(); + let a = alpha(&logits, 4, 8, 2); + assert!(a[0] > 0.99 && a[7] < 0.01); + assert!(alpha(&[], 4, 3, 3).iter().all(|v| *v == 0.0)); + } +} diff --git a/crates/segment/src/neck.rs b/crates/segment/src/neck.rs new file mode 100644 index 00000000..df4db88d --- /dev/null +++ b/crates/segment/src/neck.rs @@ -0,0 +1,73 @@ +//! The feature pyramid (`Sam3VisionNeck`): the backbone's 72 × 72 map at four scales +//! (288, 144, 72 and 36 per side), each projected to 256 channels. + +use candle_core::{Module, Tensor}; +use candle_nn::{Conv2d, Conv2dConfig, ConvTranspose2d, ConvTranspose2dConfig, VarBuilder}; + +use crate::Result; +use crate::vit::HIDDEN; + +pub const FPN: usize = 256; + +enum Scale { + /// Two 2× transposed convolutions with GELU between (4×). + Up4(ConvTranspose2d, ConvTranspose2d), + Up2(ConvTranspose2d), + Same, + /// 2 × 2 max pooling (½×). + Down2, +} + +struct Level { + scale: Scale, + proj1: Conv2d, + proj2: Conv2d, +} + +pub struct Neck { + levels: Vec, +} + +impl Neck { + /// `vb` at `…neck` (detector) or `tracker_neck`. + pub fn new(vb: VarBuilder) -> Result { + let up = ConvTranspose2dConfig { stride: 2, ..Default::default() }; + let mut levels = Vec::new(); + for (i, factor) in [4.0f32, 2.0, 1.0, 0.5].into_iter().enumerate() { + let l = vb.pp(format!("fpn_layers.{i}")); + let s = l.pp("scale_layers"); + let (scale, inter) = if factor == 4.0 { + let a = candle_nn::conv_transpose2d(HIDDEN, HIDDEN / 2, 2, up, s.pp("0"))?; + let b = candle_nn::conv_transpose2d(HIDDEN / 2, HIDDEN / 4, 2, up, s.pp("2"))?; + (Scale::Up4(a, b), HIDDEN / 4) + } else if factor == 2.0 { + (Scale::Up2(candle_nn::conv_transpose2d(HIDDEN, HIDDEN / 2, 2, up, s.pp("0"))?), HIDDEN / 2) + } else if factor == 1.0 { + (Scale::Same, HIDDEN) + } else { + (Scale::Down2, HIDDEN) + }; + levels.push(Level { + scale, + proj1: candle_nn::conv2d(inter, FPN, 1, Conv2dConfig::default(), l.pp("proj1"))?, + proj2: candle_nn::conv2d(FPN, FPN, 3, Conv2dConfig { padding: 1, ..Default::default() }, l.pp("proj2"))?, + }); + } + Ok(Self { levels }) + } + + /// `x`: the backbone's `[1, 1024, 72, 72]`; `n`: how many levels (from the finest) to compute. + pub fn forward(&self, x: &Tensor, n: usize) -> Result> { + let mut out = Vec::with_capacity(n); + for l in self.levels.iter().take(n) { + let y = match &l.scale { + Scale::Up4(a, b) => b.forward(&a.forward(x)?.gelu_erf()?)?, + Scale::Up2(a) => a.forward(x)?, + Scale::Same => x.clone(), + Scale::Down2 => x.max_pool2d(2)?, + }; + out.push(l.proj2.forward(&l.proj1.forward(&y)?)?); + } + Ok(out) + } +} diff --git a/crates/segment/src/nn.rs b/crates/segment/src/nn.rs new file mode 100644 index 00000000..cec6a58b --- /dev/null +++ b/crates/segment/src/nn.rs @@ -0,0 +1,238 @@ +//! Shared building blocks: attention, MLPs, layer norms, sine position encodings. + +use candle_core::{D, DType, Device, Module, Tensor}; +use candle_nn::{LayerNorm, Linear, VarBuilder}; + +use crate::Result; + +/// `nn.LayerNorm(d)` (PyTorch's default eps). +pub const TORCH_LN_EPS: f64 = 1e-5; + +pub fn ln(d: usize, eps: f64, vb: VarBuilder) -> Result { + Ok(candle_nn::layer_norm(d, eps, vb)?) +} + +pub fn linear(i: usize, o: usize, vb: VarBuilder) -> Result { + Ok(candle_nn::linear(i, o, vb)?) +} + +/// Largest attention-score block computed at once (elements): the 5184-token global layers +/// would otherwise need ~1.7 GB for all 16 heads. +const SCORE_BLOCK: usize = 64 << 20; + +/// Scaled dot-product attention on `[B, H, L, d]` tensors, with an optional additive bias +/// broadcastable to `[B, H, Lq, Lk]`. Heads are processed in groups to bound memory. +pub fn sdpa(q: &Tensor, k: &Tensor, v: &Tensor, bias: Option<&Tensor>, scale: f64) -> Result { + let (b, h, lq, _) = q.dims4()?; + let lk = k.dim(2)?; + let per_head = b.saturating_mul(lq).saturating_mul(lk).max(1); + let group = (SCORE_BLOCK / per_head).clamp(1, h); + let bias_heads = match bias { + Some(t) => t.dim(1)?, + None => 1, + }; + let mut outs = Vec::with_capacity(h.div_ceil(group)); + let mut start = 0; + while start < h { + let n = group.min(h - start); + let qs = q.narrow(1, start, n)?; + let ks = k.narrow(1, start, n)?; + let vs = v.narrow(1, start, n)?; + let mut att = (qs.matmul(&ks.t()?)? * scale)?; + if let Some(bias) = bias { + let bs = if bias_heads > 1 { bias.narrow(1, start, n)? } else { bias.clone() }; + att = att.broadcast_add(&bs)?; + } + let att = candle_nn::ops::softmax_last_dim(&att)?; + outs.push(att.matmul(&vs)?); + start += n; + } + Ok(Tensor::cat(&outs, 1)?) +} + +/// Multi-head attention with q/k/v/o projections (`Sam3Attention`, `Sam3TrackerAttention` +/// with `internal` = hidden / downsample rate). +pub struct Attention { + q: Linear, + k: Linear, + v: Linear, + o: Linear, + heads: usize, + head_dim: usize, +} + +impl Attention { + pub fn new(hidden: usize, internal: usize, heads: usize, vb: VarBuilder) -> Result { + Self::named(hidden, internal, heads, "o_proj", vb) + } + + /// With the output projection called `out` (CLIP's `out_proj`). + pub fn named(hidden: usize, internal: usize, heads: usize, out: &str, vb: VarBuilder) -> Result { + Ok(Self { + q: linear(hidden, internal, vb.pp("q_proj"))?, + k: linear(hidden, internal, vb.pp("k_proj"))?, + v: linear(hidden, internal, vb.pp("v_proj"))?, + o: linear(internal, hidden, vb.pp(out))?, + heads, + head_dim: internal / heads, + }) + } + + fn split(&self, x: &Tensor) -> Result { + let (b, l, _) = x.dims3()?; + Ok(x.reshape((b, l, self.heads, self.head_dim))?.transpose(1, 2)?.contiguous()?) + } + + /// The q, k and v projections of `x` (for attention variants that transform them). + pub fn qkv(&self, x: &Tensor) -> Result<(Tensor, Tensor, Tensor)> { + Ok((self.q.forward(x)?, self.k.forward(x)?, self.v.forward(x)?)) + } + + /// The output projection. + pub fn out(&self, x: &Tensor) -> Result { + Ok(self.o.forward(x)?) + } + + /// `q`, `k`, `v`: `[B, L, hidden]`; `bias` broadcastable to `[B, heads, Lq, Lk]`. + pub fn forward(&self, q: &Tensor, k: &Tensor, v: &Tensor, bias: Option<&Tensor>) -> Result { + let (b, lq, _) = q.dims3()?; + let q = self.split(&self.q.forward(q)?)?; + let k = self.split(&self.k.forward(k)?)?; + let v = self.split(&self.v.forward(v)?)?; + let o = sdpa(&q, &k, &v, bias, (self.head_dim as f64).powf(-0.5))?; + let o = o.transpose(1, 2)?.reshape((b, lq, self.heads * self.head_dim))?; + Ok(self.o.forward(&o)?) + } +} + +/// An additive key-padding bias `[1, 1, 1, L]` (0 for valid tokens, the f32 minimum for padding). +pub fn padding_bias(valid: &[bool], device: &Device) -> Result { + let v: Vec = valid.iter().map(|&ok| if ok { 0.0 } else { f32::MIN }).collect(); + Ok(Tensor::from_vec(v, (1, 1, 1, valid.len()), device)?) +} + +/// `Sam3MLP`: fc1 → activation → fc2. +pub struct Mlp { + fc1: Linear, + fc2: Linear, + gelu: bool, +} + +impl Mlp { + pub fn new(hidden: usize, inter: usize, gelu: bool, vb: VarBuilder) -> Result { + Ok(Self { fc1: linear(hidden, inter, vb.pp("fc1"))?, fc2: linear(inter, hidden, vb.pp("fc2"))?, gelu }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + let x = self.fc1.forward(x)?; + let x = if self.gelu { x.gelu_erf()? } else { x.relu()? }; + Ok(self.fc2.forward(&x)?) + } +} + +/// `Sam3DecoderMLP`: `layer1 … layerN` with ReLU between (2 or 3 layers). +pub struct DecoderMlp { + layers: Vec, +} + +impl DecoderMlp { + pub fn new(input: usize, hidden: usize, output: usize, n: usize, vb: VarBuilder) -> Result { + let mut layers = Vec::with_capacity(n); + for i in 0..n { + let (a, b) = (if i == 0 { input } else { hidden }, if i + 1 == n { output } else { hidden }); + layers.push(linear(a, b, vb.pp(format!("layer{}", i + 1)))?); + } + Ok(Self { layers }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + let mut x = x.clone(); + for (i, l) in self.layers.iter().enumerate() { + x = l.forward(&x)?; + if i + 1 < self.layers.len() { + x = x.relu()?; + } + } + Ok(x) + } +} + +/// `Sam3TrackerFeedForward`: proj_in → act → (layers.i → act)* → proj_out (→ sigmoid). +pub struct FeedForward { + proj_in: Linear, + layers: Vec, + proj_out: Linear, + sigmoid: bool, +} + +impl FeedForward { + pub fn new(input: usize, hidden: usize, output: usize, n: usize, sigmoid: bool, vb: VarBuilder) -> Result { + let mut layers = Vec::new(); + for i in 0..n.saturating_sub(2) { + layers.push(linear(hidden, hidden, vb.pp(format!("layers.{i}")))?); + } + Ok(Self { proj_in: linear(input, hidden, vb.pp("proj_in"))?, layers, proj_out: linear(hidden, output, vb.pp("proj_out"))?, sigmoid }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + let mut x = self.proj_in.forward(x)?.relu()?; + for l in &self.layers { + x = l.forward(&x)?.relu()?; + } + let x = self.proj_out.forward(&x)?; + Ok(if self.sigmoid { candle_nn::ops::sigmoid(&x)? } else { x }) + } +} + +/// Layer norm over the channel axis of an NCHW tensor (`Sam3TrackerLayerNorm`, channels_first). +pub fn ln_channels(x: &Tensor, norm: &LayerNorm) -> Result { + let x = x.permute((0, 2, 3, 1))?; + Ok(norm.forward(&x)?.permute((0, 3, 1, 2))?.contiguous()?) +} + +/// The DETR sine embedding of an `h × w` grid, normalized (`Sam3SinePositionEmbedding` with +/// `normalize=True`, scale 2π): `[1, 2·feats, h, w]`, y features first. +pub fn sine_grid(h: usize, w: usize, feats: usize, device: &Device) -> Result { + let scale = 2.0 * std::f64::consts::PI; + let dim_t: Vec = (0..feats).map(|i| 10000f64.powf(2.0 * (i / 2) as f64 / feats as f64)).collect(); + let mut out = vec![0f32; 2 * feats * h * w]; + for y in 0..h { + let ye = (y + 1) as f64 / (h as f64 + 1e-6) * scale; + for x in 0..w { + let xe = (x + 1) as f64 / (w as f64 + 1e-6) * scale; + for (i, d) in dim_t.iter().enumerate() { + let (py, px) = (ye / d, xe / d); + let (vy, vx) = if i % 2 == 0 { (py.sin(), px.sin()) } else { (py.cos(), px.cos()) }; + if let Some(v) = out.get_mut((i * h + y) * w + x) { + *v = vy as f32; + } + if let Some(v) = out.get_mut(((feats + i) * h + y) * w + x) { + *v = vx as f32; + } + } + } + } + Ok(Tensor::from_vec(out, (1, 2 * feats, h, w), device)?) +} + +/// `inverse_sigmoid` with the reference's clamping (eps 1e-3). +pub fn inverse_sigmoid(x: &Tensor) -> Result { + let x = x.clamp(0f32, 1f32)?; + let x1 = x.clamp(1e-3f32, f32::MAX)?; + let x2 = (1.0 - &x)?.clamp(1e-3f32, f32::MAX)?; + Ok((x1 / x2)?.log()?) +} + +/// `[..., 4]` boxes (cx, cy, w, h) → (x0, y0, x1, y1). +pub fn cxcywh_to_xyxy(b: &Tensor) -> Result { + let cx = b.narrow(D::Minus1, 0, 1)?; + let cy = b.narrow(D::Minus1, 1, 1)?; + let w = (b.narrow(D::Minus1, 2, 1)? * 0.5)?; + let h = (b.narrow(D::Minus1, 3, 1)? * 0.5)?; + Ok(Tensor::cat(&[(&cx - &w)?, (&cy - &h)?, (&cx + &w)?, (&cy + &h)?], D::Minus1)?) +} + +/// A tensor's values as `f32` on the CPU. +pub fn to_vec(t: &Tensor) -> Result> { + Ok(t.to_dtype(DType::F32)?.flatten_all()?.to_vec1::()?) +} diff --git a/crates/segment/src/tokenizer.rs b/crates/segment/src/tokenizer.rs new file mode 100644 index 00000000..c1a9f3c1 --- /dev/null +++ b/crates/segment/src/tokenizer.rs @@ -0,0 +1,173 @@ +//! The CLIP byte-pair-encoding tokenizer (lower-cased, `` word ends), padded to the +//! 32-token context SAM 3's text encoder takes. + +use std::collections::HashMap; +use std::path::Path; + +use crate::{Error, Result}; + +pub const CONTEXT: usize = 32; +const BOS: u32 = 49406; +const EOS: u32 = 49407; + +pub struct Tokenizer { + vocab: HashMap, + ranks: HashMap<(String, String), usize>, + byte_char: [char; 256], +} + +/// GPT-2's reversible byte → printable character map. +fn bytes_to_unicode() -> [char; 256] { + let mut map = ['\0'; 256]; + let mut n = 0u32; + for b in 0..=255u8 { + let printable = (b'!'..=b'~').contains(&b) || (0xA1..=0xAC).contains(&b) || (0xAE..=0xFF).contains(&b); + let c = if printable { + u32::from(b) + } else { + n += 1; + 255 + n + }; + if let Some(slot) = map.get_mut(usize::from(b)) { + *slot = char::from_u32(c).unwrap_or('?'); + } + } + map +} + +impl Tokenizer { + /// From `vocab.json` and `merges.txt` in `dir`. + pub fn load(dir: &Path) -> Result { + let read = |name: &str| std::fs::read_to_string(dir.join(name)).map_err(|e| Error::Model(format!("{name}: {e}"))); + let vocab: HashMap = serde_json::from_str(&read("vocab.json")?).map_err(|e| Error::Model(format!("vocab.json: {e}")))?; + let merges = read("merges.txt")?; + Ok(Self::new(vocab, &merges)) + } + + pub fn new(vocab: HashMap, merges: &str) -> Self { + let ranks = merges + .lines() + .filter(|l| !l.starts_with("#version") && !l.trim().is_empty()) + .filter_map(|l| l.split_once(' ').map(|(a, b)| (a.to_string(), b.to_string()))) + .enumerate() + .map(|(i, p)| (p, i)) + .collect(); + Self { vocab, ranks, byte_char: bytes_to_unicode() } + } + + /// Token ids (start, words, end, padding) and which positions are real tokens. + pub fn encode(&self, text: &str) -> (Vec, Vec) { + let mut ids = vec![BOS]; + for word in pre_tokenize(&text.split_whitespace().collect::>().join(" ").to_lowercase()) { + let mapped: String = word.bytes().map(|b| self.byte_char.get(usize::from(b)).copied().unwrap_or('?')).collect(); + for piece in self.bpe(&mapped) { + ids.push(self.vocab.get(&piece).copied().unwrap_or(EOS)); + } + } + ids.truncate(CONTEXT - 1); + ids.push(EOS); + let n = ids.len(); + ids.resize(CONTEXT, EOS); + let valid = (0..CONTEXT).map(|i| i < n).collect(); + (ids, valid) + } + + fn bpe(&self, word: &str) -> Vec { + let mut syms: Vec = word.chars().map(String::from).collect(); + if let Some(last) = syms.last_mut() { + last.push_str(""); + } + loop { + let best = syms + .windows(2) + .enumerate() + .filter_map(|(i, w)| match w { + [a, b] => self.ranks.get(&(a.clone(), b.clone())).map(|r| (*r, i)), + _ => None, + }) + .min(); + let Some((_, i)) = best else { break }; + let (Some(a), Some(b)) = (syms.get(i).cloned(), syms.get(i + 1).cloned()) else { break }; + let merged = a.clone() + &b; + // merge every occurrence of the pair, left to right + let mut out = Vec::with_capacity(syms.len()); + let mut j = 0; + while j < syms.len() { + if j + 1 < syms.len() && syms.get(j) == Some(&a) && syms.get(j + 1) == Some(&b) { + out.push(merged.clone()); + j += 2; + } else { + out.extend(syms.get(j).cloned()); + j += 1; + } + } + syms = out; + if syms.len() == 1 { + break; + } + } + syms + } +} + +/// CLIP's pre-tokenizer: contractions, letter runs, single digits, and runs of other +/// non-space characters. +fn pre_tokenize(text: &str) -> Vec { + let chars: Vec = text.chars().collect(); + let mut out = Vec::new(); + let mut i = 0; + while let Some(&c) = chars.get(i) { + if c.is_whitespace() { + i += 1; + continue; + } + if c == '\'' { + let rest: String = chars.get(i + 1..(i + 3).min(chars.len())).map(|s| s.iter().collect()).unwrap_or_default(); + if let Some(k) = ["re", "ve", "ll"] + .iter() + .find(|k| rest.starts_with(**k)) + .map(|k| k.len()) + .or_else(|| rest.chars().next().filter(|c| matches!(c, 's' | 't' | 'm' | 'd')).map(|_| 1)) + { + out.push(chars.get(i..i + 1 + k).map(|s| s.iter().collect()).unwrap_or_default()); + i += 1 + k; + continue; + } + } + let start = i; + if c.is_alphabetic() { + while chars.get(i).is_some_and(|c| c.is_alphabetic()) { + i += 1; + } + } else if c.is_numeric() { + i += 1; + } else { + while chars.get(i).is_some_and(|c| !c.is_whitespace() && !c.is_alphabetic() && !c.is_numeric()) { + i += 1; + } + } + out.push(chars.get(start..i).map(|s| s.iter().collect()).unwrap_or_default()); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn splits_like_clip() { + assert_eq!(pre_tokenize("a dog's toy, 25 cats!"), ["a", "dog", "'s", "toy", ",", "2", "5", "cats", "!"]); + } + + #[test] + fn merges_by_rank() { + let vocab: HashMap = + [("c", 1), ("a", 2), ("t", 3), ("ca", 4), ("cat", 5)].into_iter().map(|(k, v)| (k.to_string(), v)).collect(); + let t = Tokenizer::new(vocab, "#version: 0.2\nc a\nca t\n"); + let (ids, valid) = t.encode(" CAT "); + assert_eq!(&ids[..3], &[BOS, 5, EOS]); + assert_eq!(ids.len(), CONTEXT); + assert_eq!(valid.iter().filter(|v| **v).count(), 3); + } +} diff --git a/crates/segment/src/tracker.rs b/crates/segment/src/tracker.rs new file mode 100644 index 00000000..828d56b6 --- /dev/null +++ b/crates/segment/src/tracker.rs @@ -0,0 +1,249 @@ +//! Point prompts (SAM 3's interactive "tracker" head, SAM 2 style): a prompt encoder for +//! positive/negative clicks and a two-way-transformer mask decoder over the image features. + +use candle_core::{DType, Device, IndexOp, Module, Tensor}; +use candle_nn::{Conv2d, Conv2dConfig, ConvTranspose2d, ConvTranspose2dConfig, LayerNorm, VarBuilder}; + +use crate::nn::{Attention, FeedForward, TORCH_LN_EPS, ln, ln_channels}; +use crate::vit::{GRID, IMAGE}; +use crate::{Result, nn::to_vec}; + +const HIDDEN: usize = 256; +const MASK_TOKENS: usize = 4; +/// Low-resolution mask side (4 × the 72 grid). +pub const LOW_RES: usize = 288; + +/// The image features the decoder needs, computed once per image. +pub struct Features { + /// Level 2 (72 × 72, 256 channels). + pub image: Tensor, + /// `conv_s0`(level 0): `[1, 32, 288, 288]`. + pub s0: Tensor, + /// `conv_s1`(level 1): `[1, 64, 144, 144]`. + pub s1: Tensor, +} + +/// One click: position normalized to the image (0..1), and whether it includes or excludes. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Click { + pub x: f32, + pub y: f32, + pub positive: bool, +} + +/// A predicted mask: 288 × 288 logits over the (squared) model input, and its quality. +pub struct Prediction { + pub logits: Vec, + pub iou: f32, + pub object_score: f32, +} + +/// Random-Fourier positional encoding of points normalized to 0..1 (`Sam3TrackerPositionalEmbedding`). +struct Fourier { + pe: Vec, // [2, 128] +} + +impl Fourier { + fn encode(&self, x: f32, y: f32) -> Vec { + let half = self.pe.len() / 2; + let (cx, cy) = (2.0 * x - 1.0, 2.0 * y - 1.0); + let proj: Vec = (0..half) + .map(|i| { + let a = cx * self.pe.get(i).copied().unwrap_or(0.0) + cy * self.pe.get(half + i).copied().unwrap_or(0.0); + 2.0 * std::f32::consts::PI * a + }) + .collect(); + proj.iter().map(|a| a.sin()).chain(proj.iter().map(|a| a.cos())).collect() + } +} + +struct TwoWayBlock { + self_attn: Attention, + ln1: LayerNorm, + t2i: Attention, + ln2: LayerNorm, + mlp: FeedForward, + ln3: LayerNorm, + ln4: LayerNorm, + i2t: Attention, + skip_first_pe: bool, +} + +/// `Sam3TrackerAttention` on `[1, L, C]` (the point batch of one is folded away). +fn attn(a: &Attention, q: &Tensor, k: &Tensor, v: &Tensor) -> Result { + a.forward(q, k, v, None) +} + +impl TwoWayBlock { + fn forward(&self, queries: &Tensor, keys: &Tensor, qpe: &Tensor, kpe: &Tensor) -> Result<(Tensor, Tensor)> { + let queries = if self.skip_first_pe { + attn(&self.self_attn, queries, queries, queries)? + } else { + let q = (queries + qpe)?; + (queries + attn(&self.self_attn, &q, &q, queries)?)? + }; + let queries = self.ln1.forward(&queries)?; + let (q, k) = ((&queries + qpe)?, (keys + kpe)?); + let queries = self.ln2.forward(&(&queries + attn(&self.t2i, &q, &k, keys)?)?)?; + let queries = self.ln3.forward(&(&queries + self.mlp.forward(&queries)?)?)?; + let (q, k) = ((&queries + qpe)?, (keys + kpe)?); + let keys = self.ln4.forward(&(keys + attn(&self.i2t, &k, &q, &queries)?)?)?; + Ok((queries, keys)) + } +} + +pub struct Tracker { + device: Device, + point_pe: Fourier, + image_pe: Tensor, + point_embed: Tensor, + not_a_point: Tensor, + no_mask: Tensor, + conv_s0: Conv2d, + conv_s1: Conv2d, + output_tokens: Tensor, + blocks: Vec, + final_t2i: Attention, + ln_final: LayerNorm, + up1: ConvTranspose2d, + up_ln: LayerNorm, + up2: ConvTranspose2d, + hyper: Vec, + iou_head: FeedForward, + obj_head: FeedForward, +} + +fn read_vec(vb: &VarBuilder, shape: (usize, usize), name: &str) -> Result> { + to_vec(&vb.get(shape, name)?) +} + +impl Tracker { + /// `vb` at `tracker_model`. + pub fn new(vb: VarBuilder) -> Result { + let device = vb.device().clone(); + let pe = vb.pp("prompt_encoder"); + let md = vb.pp("mask_decoder"); + let image_fourier = Fourier { pe: read_vec(&vb, (2, HIDDEN / 2), "shared_image_embedding.positional_embedding")? }; + // image-wide positional embedding: pixel centres of the 72 × 72 grid + let mut grid = Vec::with_capacity(GRID * GRID * HIDDEN); + for r in 0..GRID { + for c in 0..GRID { + grid.extend(image_fourier.encode((c as f32 + 0.5) / GRID as f32, (r as f32 + 0.5) / GRID as f32)); + } + } + let image_pe = Tensor::from_vec(grid, (1, GRID * GRID, HIDDEN), &device)?; + let t = md.pp("transformer"); + let mut blocks = Vec::new(); + for i in 0..2 { + let b = t.pp(format!("layers.{i}")); + blocks.push(TwoWayBlock { + self_attn: Attention::new(HIDDEN, HIDDEN, 8, b.pp("self_attn"))?, + ln1: ln(HIDDEN, TORCH_LN_EPS, b.pp("layer_norm1"))?, + t2i: Attention::new(HIDDEN, HIDDEN / 2, 8, b.pp("cross_attn_token_to_image"))?, + ln2: ln(HIDDEN, TORCH_LN_EPS, b.pp("layer_norm2"))?, + mlp: FeedForward::new(HIDDEN, 2048, HIDDEN, 2, false, b.pp("mlp"))?, + ln3: ln(HIDDEN, TORCH_LN_EPS, b.pp("layer_norm3"))?, + ln4: ln(HIDDEN, TORCH_LN_EPS, b.pp("layer_norm4"))?, + i2t: Attention::new(HIDDEN, HIDDEN / 2, 8, b.pp("cross_attn_image_to_token"))?, + skip_first_pe: i == 0, + }); + } + let output_tokens = Tensor::cat( + &[ + md.get((1, HIDDEN), "obj_score_token.weight")?, + md.get((1, HIDDEN), "iou_token.weight")?, + md.get((MASK_TOKENS, HIDDEN), "mask_tokens.weight")?, + ], + 0, + )?; + let up = ConvTranspose2dConfig { stride: 2, ..Default::default() }; + let one = Conv2dConfig::default(); + Ok(Self { + point_pe: Fourier { pe: read_vec(&pe, (2, HIDDEN / 2), "shared_embedding.positional_embedding")? }, + image_pe, + point_embed: pe.get((4, HIDDEN), "point_embed.weight")?.to_device(&Device::Cpu)?, + not_a_point: pe.get((1, HIDDEN), "not_a_point_embed.weight")?.to_device(&Device::Cpu)?, + no_mask: pe.get((1, HIDDEN), "no_mask_embed.weight")?.reshape((1, HIDDEN, 1, 1))?, + conv_s0: candle_nn::conv2d(HIDDEN, HIDDEN / 8, 1, one, md.pp("conv_s0"))?, + conv_s1: candle_nn::conv2d(HIDDEN, HIDDEN / 4, 1, one, md.pp("conv_s1"))?, + output_tokens, + blocks, + final_t2i: Attention::new(HIDDEN, HIDDEN / 2, 8, t.pp("final_attn_token_to_image"))?, + ln_final: ln(HIDDEN, TORCH_LN_EPS, t.pp("layer_norm_final_attn"))?, + up1: candle_nn::conv_transpose2d(HIDDEN, HIDDEN / 4, 2, up, md.pp("upscale_conv1"))?, + up_ln: ln(HIDDEN / 4, 1e-6, md.pp("upscale_layer_norm"))?, + up2: candle_nn::conv_transpose2d(HIDDEN / 4, HIDDEN / 8, 2, up, md.pp("upscale_conv2"))?, + hyper: (0..MASK_TOKENS) + .map(|i| FeedForward::new(HIDDEN, HIDDEN, HIDDEN / 8, 3, false, md.pp(format!("output_hypernetworks_mlps.{i}")))) + .collect::>()?, + iou_head: FeedForward::new(HIDDEN, HIDDEN, MASK_TOKENS, 3, true, md.pp("iou_prediction_head"))?, + obj_head: FeedForward::new(HIDDEN, HIDDEN, 1, 3, false, md.pp("pred_obj_score_head"))?, + device, + }) + } + + /// The decoder's image features from the tracker neck's levels (288, 144, 72 per side). + pub fn features(&self, levels: &[Tensor]) -> Result { + let [l0, l1, l2, ..] = levels else { return Err(crate::Error::Model("tracker needs three feature levels".into())) }; + Ok(Features { image: l2.clone(), s0: self.conv_s0.forward(l0)?, s1: self.conv_s1.forward(l1)? }) + } + + /// Sparse prompt embeddings `[1, n + 1, 256]` (a padding point closes the list). + fn embed_clicks(&self, clicks: &[Click]) -> Result { + let point_embed = to_vec(&self.point_embed)?; + let not_a_point = to_vec(&self.not_a_point)?; + let mut rows = Vec::with_capacity((clicks.len() + 1) * HIDDEN); + for c in clicks { + // the reference shifts pixel coordinates (in the 1008² input) by half a pixel + let (x, y) = ((c.x * IMAGE as f32 + 0.5) / IMAGE as f32, (c.y * IMAGE as f32 + 0.5) / IMAGE as f32); + let e = self.point_pe.encode(x, y); + let label = usize::from(c.positive); + rows.extend(e.iter().enumerate().map(|(i, v)| v + point_embed.get(label * HIDDEN + i).copied().unwrap_or(0.0))); + } + rows.extend_from_slice(¬_a_point); + Ok(Tensor::from_vec(rows, (1, clicks.len() + 1, HIDDEN), &self.device)?) + } + + /// Predict masks for `clicks`. `multimask`: the three alternatives (best by predicted IoU is + /// returned), else the single-mask output with SAM 2's stability fallback. + pub fn predict(&self, f: &Features, clicks: &[Click], multimask: bool) -> Result { + let sparse = self.embed_clicks(clicks)?; + let tokens = Tensor::cat(&[self.output_tokens.unsqueeze(0)?, sparse], 1)?; // [1, T, 256] + let image = f.image.broadcast_add(&self.no_mask)?; // dense prompt: "no mask" + let (b, c, h, w) = image.dims4()?; + let keys = image.flatten_from(2)?.transpose(1, 2)?.contiguous()?; // [1, HW, 256] + let (mut q, mut k) = (tokens.clone(), keys); + for blk in &self.blocks { + (q, k) = blk.forward(&q, &k, &tokens, &self.image_pe)?; + } + let qf = (&q + &tokens)?; + let kf = (&k + &self.image_pe)?; + let q = self.ln_final.forward(&(&q + attn(&self.final_t2i, &qf, &kf, &k)?)?)?; + let iou_token = q.i((.., 1, ..))?; + let image = k.transpose(1, 2)?.reshape((b, c, h, w))?; + let up = (self.up1.forward(&image)? + &f.s1)?; + let up = ln_channels(&up, &self.up_ln)?.gelu_erf()?; + let up = (self.up2.forward(&up)? + &f.s0)?.gelu_erf()?; // [1, 32, 288, 288] + let (_, uc, uh, uw) = up.dims4()?; + let hyper = (0..MASK_TOKENS) + .map(|i| self.hyper.get(i).map_or(Err(crate::Error::Model("hypernetwork".into())), |m| m.forward(&q.i((.., 2 + i, ..))?))) + .collect::>>()?; + let hyper = Tensor::stack(&hyper, 1)?; // [1, 4, 32] + let masks = hyper.matmul(&up.reshape((b, uc, uh * uw))?)?.reshape((b, MASK_TOKENS, uh, uw))?; + let iou = to_vec(&self.iou_head.forward(&iou_token)?)?; + let object_score = to_vec(&self.obj_head.forward(&q.i((.., 0, ..))?)?)?.first().copied().unwrap_or(0.0); + let pick = |i: usize| -> Result { + Ok(Prediction { logits: to_vec(&masks.i((0, i))?)?, iou: iou.get(i).copied().unwrap_or(0.0), object_score }) + }; + let best_multi = || (1..MASK_TOKENS).max_by(|a, b| iou.get(*a).partial_cmp(&iou.get(*b)).unwrap_or(std::cmp::Ordering::Equal)).unwrap_or(1); + if multimask { + return pick(best_multi()); + } + // single-mask output unless it is unstable (area changes > 2 % when the threshold moves ±0.05) + let single = masks.i((0, 0))?; + let inter = single.gt(0.05)?.to_dtype(DType::F32)?.sum_all()?.to_scalar::()?; + let union = single.gt(-0.05)?.to_dtype(DType::F32)?.sum_all()?.to_scalar::()?; + let stability = if union > 0.0 { inter / union } else { 1.0 }; + if stability >= 0.98 { pick(0) } else { pick(best_multi()) } + } +} diff --git a/crates/segment/src/vit.rs b/crates/segment/src/vit.rs new file mode 100644 index 00000000..ad901826 --- /dev/null +++ b/crates/segment/src/vit.rs @@ -0,0 +1,166 @@ +//! The SAM 3 vision backbone: a 32-layer ViT (Perception Encoder) with 2-D axial rotary +//! position encoding, windowed attention (24 × 24) in most layers and global attention in four. + +use candle_core::{D, Device, Module, Tensor}; +use candle_nn::{Conv2d, Conv2dConfig, LayerNorm, VarBuilder}; + +use crate::nn::{Attention, Mlp, ln, sdpa}; +use crate::{Error, Result}; + +pub const IMAGE: usize = 1008; +pub const PATCH: usize = 14; +/// Patches per side (72). +pub const GRID: usize = IMAGE / PATCH; +pub const HIDDEN: usize = 1024; +const HEADS: usize = 16; +const HEAD_DIM: usize = HIDDEN / HEADS; +const LAYERS: usize = 32; +const WINDOW: usize = 24; +const GLOBAL: [usize; 4] = [7, 15, 23, 31]; +const PRETRAIN_GRID: usize = 336 / PATCH; +const EPS: f64 = 1e-6; + +/// Cosine and sine tables `[N, head_dim]` for a `side × side` grid whose coordinates are scaled +/// by `scale` (`Sam3ViTRotaryEmbedding`: x frequencies then y, each duplicated pairwise). +fn rope_tables(side: usize, scale: f64, device: &Device) -> Result<(Tensor, Tensor)> { + let spatial = HEAD_DIM / 2; + let inv: Vec = (0..spatial).step_by(2).map(|i| 1.0 / 10000f64.powf(i as f64 / spatial as f64)).collect(); + let n = side * side; + let (mut cos, mut sin) = (Vec::with_capacity(n * HEAD_DIM), Vec::with_capacity(n * HEAD_DIM)); + for r in 0..side { + for c in 0..side { + let (x, y) = (c as f64 * scale, r as f64 * scale); + for pos in [x, y] { + for f in &inv { + let a = (pos * f) as f32 as f64; + let (s, co) = (a.sin() as f32, a.cos() as f32); + cos.extend([co, co]); + sin.extend([s, s]); + } + } + } + } + Ok((Tensor::from_vec(cos, (n, HEAD_DIM), device)?, Tensor::from_vec(sin, (n, HEAD_DIM), device)?)) +} + +/// `x·cos + rotate_pairwise(x)·sin`, where rotate_pairwise maps (a, b) pairs to (−b, a). +fn apply_rope(x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result { + let dims = x.dims().to_vec(); + let last = *dims.last().ok_or_else(|| Error::Model("rope on a scalar".into()))?; + let mut paired = dims.clone(); + paired.pop(); + paired.extend([last / 2, 2]); + let p = x.reshape(paired.as_slice())?; + let a = p.narrow(D::Minus1, 0, 1)?; + let b = p.narrow(D::Minus1, 1, 1)?; + let rot = Tensor::cat(&[b.neg()?, a], D::Minus1)?.reshape(dims.as_slice())?; + Ok((x.broadcast_mul(cos)? + rot.broadcast_mul(sin)?)?) +} + +struct Layer { + ln1: LayerNorm, + attn: Attention, + ln2: LayerNorm, + mlp: Mlp, + window: usize, +} + +impl Layer { + /// `x`: `[B, H, W, C]`. + fn forward(&self, x: &Tensor, rope: &(Tensor, Tensor)) -> Result { + let (b, h, w, c) = x.dims4()?; + let y = self.ln1.forward(x)?; + let y = match (h.checked_div(self.window), w.checked_div(self.window)) { + (Some(nh), Some(nw)) => { + // non-overlapping windows (72 is a multiple of 24: no padding) + let ws = self.window; + let y = y.reshape((b, nh, ws, nw, ws, c))?.permute((0, 1, 3, 2, 4, 5))?.reshape((b * nh * nw, ws, ws, c))?; + let y = self.attend(&y, rope)?; + y.reshape((b, nh, nw, ws, ws, c))?.permute((0, 1, 3, 2, 4, 5))?.reshape((b, h, w, c))? + } + // global attention + _ => self.attend(&y, rope)?, + }; + let x = (x + y)?; + let y = self.mlp.forward(&self.ln2.forward(&x)?)?; + Ok((x + y)?) + } + + /// RoPE self-attention over `[B, H, W, C]`. + fn attend(&self, x: &Tensor, rope: &(Tensor, Tensor)) -> Result { + let (b, h, w, c) = x.dims4()?; + let seq = x.reshape((b, h * w, c))?; + Ok(self.attn_rope(&seq, rope)?.reshape((b, h, w, c))?) + } + + fn attn_rope(&self, x: &Tensor, rope: &(Tensor, Tensor)) -> Result { + let (b, n, _) = x.dims3()?; + let split = |t: Tensor| -> Result { Ok(t.reshape((b, n, HEADS, HEAD_DIM))?.transpose(1, 2)?.contiguous()?) }; + let (q, k, v) = self.attn.qkv(x)?; + let q = apply_rope(&split(q)?, &rope.0, &rope.1)?; + let k = apply_rope(&split(k)?, &rope.0, &rope.1)?; + let v = split(v)?; + let o = sdpa(&q, &k, &v, None, (HEAD_DIM as f64).powf(-0.5))?; + let o = o.transpose(1, 2)?.reshape((b, n, HIDDEN))?; + self.attn.out(&o) + } +} + +pub struct Vit { + patch: Conv2d, + pos: Tensor, + ln_pre: LayerNorm, + layers: Vec, + rope_window: (Tensor, Tensor), + rope_global: (Tensor, Tensor), +} + +impl Vit { + /// `vb` at `detector_model.vision_encoder.backbone`. + pub fn new(vb: VarBuilder) -> Result { + let device = vb.device().clone(); + let cfg = Conv2dConfig { stride: PATCH, ..Default::default() }; + let patch = candle_nn::conv2d_no_bias(3, HIDDEN, PATCH, cfg, vb.pp("embeddings.patch_embeddings.projection"))?; + let pos = vb.get((1, PRETRAIN_GRID * PRETRAIN_GRID, HIDDEN), "embeddings.position_embeddings")?; + // tile (not interpolate) the 24 × 24 pretraining table over the 72 × 72 grid + let reps = GRID.div_ceil(PRETRAIN_GRID); + let pos = pos + .reshape((1, PRETRAIN_GRID, PRETRAIN_GRID, HIDDEN))? + .repeat((1, reps, reps, 1))? + .narrow(1, 0, GRID)? + .narrow(2, 0, GRID)? + .reshape((1, GRID * GRID, HIDDEN))?; + let mut layers = Vec::with_capacity(LAYERS); + for i in 0..LAYERS { + let l = vb.pp(format!("layers.{i}")); + layers.push(Layer { + ln1: ln(HIDDEN, EPS, l.pp("layer_norm1"))?, + attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("attention"))?, + ln2: ln(HIDDEN, EPS, l.pp("layer_norm2"))?, + mlp: Mlp::new(HIDDEN, 4736, true, l.pp("mlp"))?, + window: if GLOBAL.contains(&i) { 0 } else { WINDOW }, + }); + } + Ok(Self { + patch, + pos, + ln_pre: ln(HIDDEN, EPS, vb.pp("layer_norm"))?, + layers, + rope_window: rope_tables(WINDOW, 1.0, &device)?, + rope_global: rope_tables(GRID, WINDOW as f64 / GRID as f64, &device)?, + }) + } + + /// `pixels`: `[1, 3, 1008, 1008]` normalized to −1..1. Returns `[1, 1024, 72, 72]`. + pub fn forward(&self, pixels: &Tensor) -> Result { + let x = self.patch.forward(pixels)?; // [1, C, 72, 72] + let b = x.dim(0)?; + let x = x.flatten_from(2)?.transpose(1, 2)?.broadcast_add(&self.pos)?; + let mut x = self.ln_pre.forward(&x.reshape((b, GRID, GRID, HIDDEN))?)?; + for l in &self.layers { + let rope = if l.window > 0 { &self.rope_window } else { &self.rope_global }; + x = l.forward(&x, rope)?; + } + Ok(x.permute((0, 3, 1, 2))?.contiguous()?) + } +} diff --git a/crates/segment/src/weights.rs b/crates/segment/src/weights.rs new file mode 100644 index 00000000..15de564a --- /dev/null +++ b/crates/segment/src/weights.rs @@ -0,0 +1,108 @@ +//! Safetensors weights read on demand. +//! +//! Tensors are read from the file when a module asks for them (positional reads, no memory +//! map: `unsafe` stays in `lightcraft-sysmem`), so building only the point-prompt model never +//! touches the text encoder's bytes. + +use std::collections::HashMap; +use std::fs::File; +use std::io::{Read, Seek, SeekFrom}; +use std::path::Path; +use std::sync::{Arc, Mutex}; + +use candle_core::{DType, Device, Shape, Tensor}; +use candle_nn::var_builder::SimpleBackend; + +use crate::{Error, Result}; + +struct Entry { + shape: Vec, + start: u64, + end: u64, +} + +/// An open safetensors file: its header, and the file to read tensor bytes from. +pub struct Weights { + file: Mutex, + entries: HashMap, + data_start: u64, +} + +impl Weights { + pub fn open(path: &Path) -> Result> { + let mut file = File::open(path).map_err(|e| Error::Model(format!("{}: {e}", path.display())))?; + let mut len = [0u8; 8]; + file.read_exact(&mut len).map_err(|e| Error::Model(format!("{}: {e}", path.display())))?; + let n = u64::from_le_bytes(len); + if n > 100 << 20 { + return Err(Error::Model(format!("{}: not a safetensors file (header of {n} bytes)", path.display()))); + } + let mut header = vec![0u8; n as usize]; + file.read_exact(&mut header).map_err(|e| Error::Model(format!("{}: {e}", path.display())))?; + let header: serde_json::Map = + serde_json::from_slice(&header).map_err(|e| Error::Model(format!("{}: bad header: {e}", path.display())))?; + let mut entries = HashMap::new(); + for (name, v) in header { + if name == "__metadata__" { + continue; + } + let dtype = match v.get("dtype").and_then(|d| d.as_str()) { + Some("F32") => DType::F32, + other => return Err(Error::Model(format!("{name}: unsupported dtype {other:?} (the SAM 3 checkpoint is F32)"))), + }; + let shape: Vec = v + .get("shape") + .and_then(|s| s.as_array()) + .map(|a| a.iter().filter_map(|x| x.as_u64().map(|x| x as usize)).collect()) + .unwrap_or_default(); + let offsets: Vec = + v.get("data_offsets").and_then(|s| s.as_array()).map(|a| a.iter().filter_map(|x| x.as_u64()).collect()).unwrap_or_default(); + let [start, end] = offsets[..] else { return Err(Error::Model(format!("{name}: bad data_offsets"))) }; + let want = shape.iter().try_fold(dtype.size_in_bytes(), |a, d| a.checked_mul(*d)); + if end.checked_sub(start) != want.map(|w| w as u64) { + return Err(Error::Model(format!("{name}: size does not match its shape"))); + } + entries.insert(name, Entry { shape, start, end }); + } + Ok(Arc::new(Self { file: Mutex::new(file), entries, data_start: 8 + n })) + } + + pub fn contains(&self, name: &str) -> bool { + self.entries.contains_key(name) + } + + /// Tensor `name` as stored, on the CPU. + pub fn read(&self, name: &str) -> Result { + let e = self.entries.get(name).ok_or_else(|| Error::Model(format!("missing tensor `{name}`")))?; + let mut bytes = vec![0u8; e.end.saturating_sub(e.start) as usize]; + { + let mut f = self.file.lock().unwrap_or_else(std::sync::PoisonError::into_inner); + f.seek(SeekFrom::Start(self.data_start + e.start)).map_err(|err| Error::Model(format!("{name}: {err}")))?; + f.read_exact(&mut bytes).map_err(|err| Error::Model(format!("{name}: {err}")))?; + } + let v: Vec = bytes.as_chunks::<4>().0.iter().map(|c| f32::from_le_bytes(*c)).collect(); + Ok(Tensor::from_vec(v, e.shape.as_slice(), &Device::Cpu)?) + } +} + +/// [`Weights`] as a candle `VarBuilder` backend. +pub struct Backend(pub Arc); + +impl SimpleBackend for Backend { + fn get(&self, s: Shape, name: &str, _h: candle_nn::Init, dtype: DType, dev: &Device) -> candle_core::Result { + let t = self.get_unchecked(name, dtype, dev)?; + if t.shape() != &s { + candle_core::bail!("shape mismatch for {name}: expected {s:?}, got {:?}", t.shape()) + } + Ok(t) + } + + fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> candle_core::Result { + let t = self.0.read(name).map_err(|e| candle_core::Error::Msg(e.to_string()))?; + t.to_dtype(dtype)?.to_device(dev) + } + + fn contains_tensor(&self, name: &str) -> bool { + self.0.contains(name) + } +} diff --git a/crates/segment/tests/reference.rs b/crates/segment/tests/reference.rs new file mode 100644 index 00000000..a73fcaa7 --- /dev/null +++ b/crates/segment/tests/reference.rs @@ -0,0 +1,97 @@ +//! The port against the reference implementation (Hugging Face `transformers`, fp32). +//! +//! Needs the checkpoint and a reference file made by `tools/sam3_reference.py`: +//! `LIGHTCRAFT_SAM3_DIR= LIGHTCRAFT_SAM3_REF= cargo test -p lightcraft-segment --release -- --nocapture`. +//! Skipped (passes) when either is missing. + +use std::collections::HashMap; +use std::path::PathBuf; + +use candle_core::{Device, IndexOp, Tensor}; +use lightcraft_segment::{Click, Sam3, tokenizer::Tokenizer}; + +fn setup() -> Option<(PathBuf, HashMap)> { + let dir = PathBuf::from(std::env::var_os("LIGHTCRAFT_SAM3_DIR")?); + let reference = PathBuf::from(std::env::var_os("LIGHTCRAFT_SAM3_REF")?); + let r = candle_core::safetensors::load(&reference, &Device::Cpu).ok()?; + Some((dir, r)) +} + +fn vec(t: &Tensor) -> Vec { + t.to_device(&Device::Cpu).unwrap().flatten_all().unwrap().to_vec1::().unwrap() +} + +/// Max |a − b| relative to the reference's max |b|. +fn rel_err(name: &str, a: &[f32], b: &[f32]) -> f32 { + assert_eq!(a.len(), b.len(), "{name}: length"); + let scale = b.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6); + let err = a.iter().zip(b).fold(0f32, |m, (x, y)| m.max((x - y).abs())); + eprintln!("{name:>18}: max abs err {err:.2e}, rel {:.2e} (scale {scale:.2e})", err / scale); + err / scale +} + +/// Fraction of pixels where the two logit maps disagree on the mask (> 0). +fn mask_disagreement(a: &[f32], b: &[f32]) -> f32 { + a.iter().zip(b).filter(|(x, y)| (**x > 0.0) != (**y > 0.0)).count() as f32 / a.len() as f32 +} + +#[test] +fn matches_reference() { + let Some((dir, r)) = setup() else { + eprintln!("skipped: set LIGHTCRAFT_SAM3_DIR and LIGHTCRAFT_SAM3_REF"); + return; + }; + let mut model = Sam3::load(&dir).unwrap(); + eprintln!("device: {:?}", model.device()); + let t = std::time::Instant::now(); + let mut enc = model.encode_pixels(&r["pixel_values"]).unwrap(); + let backbone = enc.backbone().permute((0, 2, 3, 1)).unwrap(); + assert!(rel_err("backbone", &vec(&backbone), &vec(&r["backbone"])) < 1e-3); + eprintln!("backbone: {:?}", t.elapsed()); + + // points: two clicks (single-mask output), one click (best of three) + let clicks = [Click { x: 0.5, y: 0.55, positive: true }, Click { x: 0.2, y: 0.2, positive: false }]; + let t = std::time::Instant::now(); + let p = model.segment_clicks(&mut enc, &clicks).unwrap(); + eprintln!("clicks (incl. tracker load + features): {:?}", t.elapsed()); + let want = vec(&r["trk_masks"]); + assert!(rel_err("2-click logits", &p.logits, &want) < 1e-2); + assert!(mask_disagreement(&p.logits, &want) < 1e-3); + let t = std::time::Instant::now(); + let p = model.segment_clicks(&mut enc, &clicks[..1]).unwrap(); + eprintln!("one more click: {:?}", t.elapsed()); + let iou = vec(&r["trk_multi_iou"]); + let best = (0..3).max_by(|a, b| iou[*a].partial_cmp(&iou[*b]).unwrap()).unwrap(); + let want = vec(&r["trk_multi_masks"].i((0, 0, best)).unwrap()); + assert!(rel_err("1-click logits", &p.logits, &want) < 1e-2); + assert!(mask_disagreement(&p.logits, &want) < 1e-3); + + // tokenizer + let tok = Tokenizer::load(&dir).unwrap(); + let cases = ["person", "Red Car", "a dog's toy", "sky, clouds!", "café 2 people", "tree-line"]; + for (i, c) in cases.iter().enumerate() { + let (ids, _) = tok.encode(c); + let want: Vec = r[&format!("tok{i}")].flatten_all().unwrap().to_vec1().unwrap(); + assert_eq!(ids.iter().map(|v| *v as i64).collect::>(), want, "tokens of {c:?}"); + } + + // text prompts + for (i, text) in ["moon", "sky", "bird"].iter().enumerate() { + let t = std::time::Instant::now(); + let out = model.detect(&mut enc, text).unwrap(); + eprintln!("detect {text:?}: {:?}", t.elapsed()); + assert!(rel_err("text features", &vec(&out.text), &vec(&r[&format!("det{i}_text")])) < 1e-3); + assert!(rel_err("query logits", &out.logits, &vec(&r[&format!("det{i}_logits")])) < 1e-2); + let presence = vec(&r[&format!("det{i}_presence")])[0]; + assert!((out.presence - presence).abs() < 1e-2, "presence {} vs {presence}", out.presence); + let boxes: Vec = out.boxes.iter().flatten().copied().collect(); + assert!(rel_err("boxes", &boxes, &vec(&r[&format!("det{i}_boxes")])) < 1e-2); + // the best instance's mask + let logits = vec(&r[&format!("det{i}_logits")]); + let q = (0..logits.len()).max_by(|a, b| logits[*a].partial_cmp(&logits[*b]).unwrap()).unwrap(); + let ours = vec(&out.masks.i(q).unwrap()); + let want = vec(&r[&format!("det{i}_masks")].i((0, q)).unwrap()); + assert!(rel_err("best mask", &ours, &want) < 2e-2); + assert!(mask_disagreement(&ours, &want) < 2e-3); + } +} diff --git a/crates/ui-egui/src/lib.rs b/crates/ui-egui/src/lib.rs index 8addc61b..0afdd915 100644 --- a/crates/ui-egui/src/lib.rs +++ b/crates/ui-egui/src/lib.rs @@ -260,6 +260,7 @@ impl LightcraftApp { } let r = self.session.execute(id, ¶ms).map_err(|e| e.to_string()); if let Err(e) = &r { + log::warn!("{id}: {e}"); self.ui.status = e.clone(); } r @@ -348,6 +349,34 @@ impl LightcraftApp { self.ui.toast = Some((text.into(), t + 1.4)); } + /// AI masks: start a zoomed-in detail pass once the clicking stops (`ui.detail_due`), and + /// apply finished passes. + fn ai_mask_detail(&mut self, ctx: &egui::Context) { + if self.session.segment_poll() { + ctx.request_repaint(); + } + let now = ctx.input(|i| i.time); + if let Some((due, mask)) = self.ui.detail_due { + if now >= due && !self.session.segmenter.busy() && !self.session.segmenter.detail_busy() { + self.ui.detail_due = None; + if let Err(e) = self.run("mask.refineDetail", serde_json::json!({"id": mask})) { + log::warn!("detail pass: {e}"); + } + } else { + ctx.request_repaint_after(std::time::Duration::from_millis(150)); + } + } + if self.session.segmenter.detail_busy() { + ctx.request_repaint_after(std::time::Duration::from_millis(200)); + } + } + + /// A toast for an error the user has to read and act on (stays 6 s). + pub fn toast_error(&mut self, ctx: &egui::Context, text: impl Into) { + let t = ctx.input(|i| i.time); + self.ui.toast = Some((text.into(), t + 6.0)); + } + fn drain_control(&mut self, ctx: &egui::Context) { let Some(rx) = self.control_rx.take() else { return }; while let Ok(req) = rx.try_recv() { @@ -673,6 +702,7 @@ impl LightcraftApp { } // panels set it again this frame while the pointer rests on a preset or profile self.hover_preview = None; + self.ai_mask_detail(&ctx); if self.ui.fullscreen { // full-screen preview: the photo alone on black egui::CentralPanel::default().frame(egui::Frame::NONE.fill(egui::Color32::BLACK)).show(ui, |ui| panels::detail::show(self, ui)); diff --git a/crates/ui-egui/src/panels/detail.rs b/crates/ui-egui/src/panels/detail.rs index 0697b75c..e8341a2a 100644 --- a/crates/ui-egui/src/panels/detail.rs +++ b/crates/ui-egui/src/panels/detail.rs @@ -526,6 +526,9 @@ fn targeted_drag(app: &mut LightcraftApp, ui: &mut egui::Ui, resp: &egui::Respon } } +/// The colour of a mask hovered in the Masks list. +const HOVER_RED: [u8; 3] = [230, 30, 40]; + /// The diagnostic overlay the loupe shows (the selected mask, Point Color's visualized range, /// Visualize Spots). pub(crate) fn view_overlay(app: &LightcraftApp, d: &DevelopSettings) -> lightcraft_pipeline::Overlay { @@ -533,6 +536,13 @@ pub(crate) fn view_overlay(app: &LightcraftApp, d: &DevelopSettings) -> lightcra if app.ui.fullscreen { return Overlay::None; } + // a mask hovered in the Masks list shows in red; the selected one only with the overlay on (O) + if app.ui.right == RightPanel::Masking + && let Some(m) = app.ui.hover_mask.and_then(|id| d.masks.iter().find(|m| m.id == id)) + && !m.components.is_empty() + { + return Overlay::Mask { id: m.id.min(u16::MAX as u32) as u16, view: MaskView::default(), color: HOVER_RED, opacity: 55 }; + } if app.ui.right == RightPanel::Masking && app.ui.mask_overlay && let Some(m) = app.session.active_mask.and_then(|id| d.masks.iter().find(|m| m.id == id)) @@ -941,6 +951,41 @@ fn mask_overlay(app: &mut LightcraftApp, ui: &mut egui::Ui, resp: &egui::Respons } } } + // Object selections (SAM 3) of the selected mask: each click, green to include, red to exclude + for m in d.masks.iter().filter(|m| Some(m.id) == active) { + for c in &m.components { + if let MaskShape::Object { hint, exclude, .. } = &c.shape { + for (pts, col) in [(hint, Color32::from_rgb(40, 200, 90)), (exclude, Color32::from_rgb(230, 60, 60))] { + for q in pts { + let q = map.screen(*q); + p.circle_filled(q, 5.0, col); + p.circle_stroke(q, 5.0, Stroke::new(1.5, Color32::WHITE)); + } + } + } + } + } + // Object tool: a click includes what's under it, ⌥-click leaves it out + if app.ui.tool == "object" { + if resp.hovered() { + ui.ctx().set_cursor_icon(egui::CursorIcon::Crosshair); + } + if resp.clicked() + && let Some(q) = resp.interact_pointer_pos() + { + let n = map.norm(q); + let exclude = ui.input(|i| i.modifiers.alt); + match app.run("mask.objectPoint", json!({"x": n.x, "y": n.y, "exclude": exclude})) { + Ok(_) => { + if let Some(mid) = app.session.active_mask { + app.ui.detail_due = Some((ui.input(|i| i.time) + 1.0, mid)); + } + } + Err(e) => app.toast_error(ui.ctx(), e), + } + } + return; + } // brush tool if app.ui.tool == "brush" { let r = (app.ui.brush_size as f64 * long) as f32; diff --git a/crates/ui-egui/src/panels/masking.rs b/crates/ui-egui/src/panels/masking.rs index f938dc8d..86fa7aad 100644 --- a/crates/ui-egui/src/panels/masking.rs +++ b/crates/ui-egui/src/panels/masking.rs @@ -53,6 +53,7 @@ fn kind_label(s: &MaskShape) -> (&'static str, Icon) { MaskShape::Sky => ("Sky", Icon::Sky), MaskShape::Background => ("Background", Icon::Subject), MaskShape::Object { .. } => ("Object", Icon::Subject), + MaskShape::Prompt { .. } => ("Describe", Icon::Subject), MaskShape::People { .. } => ("People", Icon::Subject), MaskShape::Landscape { .. } => ("Landscape", Icon::Sky), } @@ -75,7 +76,9 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { egui::Frame::NONE.inner_margin(egui::Margin { left: 24, right: 22, top: 0, bottom: 10 }).show(ui, |ui| { ui.label(egui::RichText::new(crate::i18n::tr("Create New Mask")).color(t.text_dim)); ui.add_space(6.0); - let tiles: [(&str, &str, Icon); 8] = [ + let tiles: [(&str, &str, Icon); 10] = [ + ("object", "Object", Icon::Subject), + ("prompt", "Describe", Icon::Subject), ("subject", "Subject", Icon::Subject), ("sky", "Sky", Icon::Sky), ("background", "Background", Icon::Subject), @@ -102,6 +105,8 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { app.ui.tool = "colorRange".into(); app.toast(ui.ctx(), "Click the photo to pick a colour · ⇧-click adds more"); } + "object" => start_object(app, ui.ctx(), "new"), + "prompt" => app.ui.describe = Some(("new".into(), String::new())), "brush" | "linear" | "radial" => { app.ui.tool = kind.to_string(); if *kind != "brush" { @@ -118,6 +123,19 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { } } }); + describe_field(app, ui, true); + let status = if app.session.segmenter.busy() { + Some("Analyzing the photo for AI masks…") + } else if app.session.segmenter.detail_busy() || app.ui.detail_due.is_some() { + Some("Refining the mask's detail…") + } else { + None + }; + if let Some(status) = status { + ui.add_space(4.0); + ui.label(egui::RichText::new(crate::i18n::tr(status)).color(t.text_dim)); + ui.ctx().request_repaint_after(std::time::Duration::from_millis(250)); + } }); divider(ui); // mask list @@ -126,7 +144,17 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { ui.label(egui::RichText::new(crate::i18n::tr("Masks")).font(t.semibold(13.0)).color(t.text_label)); ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| { let on = app.ui.mask_overlay; - if icon_button(ui, "maskOverlay", Icon::Eye, vec2(24.0, 24.0), on, true, "Show overlay (O)").clicked() { + if icon_button( + ui, + "maskOverlay", + Icon::Eye, + vec2(24.0, 24.0), + on, + true, + "Always show the selected mask (O); hovering a mask shows it", + ) + .clicked() + { app.ui.mask_overlay = !on; } }); @@ -135,6 +163,7 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { ui.label(egui::RichText::new(crate::i18n::tr("No masks yet. Choose a mask type above.")).color(t.text_dim)); } let count = d.masks.len(); + app.ui.hover_mask = None; for (index, m) in d.masks.iter().enumerate() { let sel = app.session.active_mask == Some(m.id); if let Some((rid, name)) = app.ui.renaming_mask.as_mut().filter(|(rid, _)| *rid == m.id) { @@ -158,6 +187,10 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { } let (r, resp) = ui.allocate_exact_size(vec2(ui.available_width(), 30.0), Sense::click()); register(ui.ctx(), format!("mask:{}", m.id), r); + // the mask under the pointer shows on the photo (red overlay) + if ui.rect_contains_pointer(r) { + app.ui.hover_mask = Some(m.id); + } ui.painter().rect_filled( r, 4.0, @@ -201,6 +234,18 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { app.ui.renaming_mask = Some((m.id, m.name.clone())); } resp.context_menu(|ui| mask_menu(app, ui, m.id, &m.name, m.visible, index, count)); + if sel { + // right under the selected mask: ＋ adds to it, − takes away from it (any mask + // type; Describe… opens its field here) + ui.horizontal(|ui| { + ui.add_space(30.0); + let plus = text_button(ui, &format!("maskPlus:{}", m.id), "+", false).on_hover_text(crate::i18n::tr("Add to this mask")); + egui::Popup::menu(&plus).show(|ui| component_menu(app, ui, "add")); + let minus = text_button(ui, &format!("maskMinus:{}", m.id), "−", false).on_hover_text(crate::i18n::tr("Subtract from this mask")); + egui::Popup::menu(&minus).show(|ui| component_menu(app, ui, "subtract")); + }); + describe_field(app, ui, false); + } } if !d.masks.is_empty() { ui.add_space(6.0); @@ -213,7 +258,10 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { egui::Frame::NONE.inner_margin(egui::Margin { left: 24, right: 22, top: 8, bottom: 8 }).show(ui, |ui| { for (i, c) in m.components.iter().enumerate() { let (kind, icon) = kind_label(&c.shape); - let label = c.name.clone().unwrap_or_else(|| kind.to_string()); + let label = c.name.clone().unwrap_or_else(|| match &c.shape { + MaskShape::Prompt { text, .. } => format!("“{text}”"), + _ => kind.to_string(), + }); ui.horizontal(|ui| { let (r, _) = ui.allocate_exact_size(vec2(16.0, 16.0), Sense::hover()); paint(ui.painter(), r, icon, t.text_label); @@ -249,6 +297,9 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { }); let resp = resp.inner; register(ui.ctx(), format!("component:{i}"), resp.rect); + if ui.rect_contains_pointer(resp.rect) { + app.ui.hover_mask = Some(m.id); + } if resp.double_clicked() { app.ui.renaming_component = Some((m.id, i, label.clone())); } @@ -486,6 +537,30 @@ fn range_controls(app: &mut LightcraftApp, ui: &mut egui::Ui, comp: usize, shape } } } + MaskShape::Object { edge, .. } | MaskShape::Prompt { edge, .. } => { + // Edge: how crisp the selection's border is (−100 hard … 0 as computed … 100 soft) + let spec = ControlSpec { + id: "edge", + label: "Edge", + section: Section::Light, + min: -100.0, + max: 100.0, + default: 0.0, + step: 1.0, + decimals: 0, + track: Track::Centered, + }; + let out = slider(ui, &spec, *edge, true, None); + let base = shape.clone(); + apply_slider_out(app, &spec, out, |app, v| { + let mut s = base.clone(); + if let MaskShape::Object { edge, .. } | MaskShape::Prompt { edge, .. } = &mut s { + *edge = v.clamp(-100.0, 100.0); + } + app.run("mask.update", json!({"component": comp, "shape": s})) + }); + ui.label(egui::RichText::new(crate::i18n::tr("− harder border · + softer border")).color(t.text_dim).size(11.0)); + } MaskShape::ColorRange { samples, refine } => { ui.horizontal(|ui| { ui.label( @@ -560,7 +635,84 @@ fn component_row_menu(app: &mut LightcraftApp, ui: &mut egui::Ui, mask: u32, k: } } +/// Start an Object selection (SAM 3 clicks): a new mask, or a component of the selected one +/// combined by `op`; the photo is analyzed meanwhile. +fn start_object(app: &mut LightcraftApp, ctx: &egui::Context, op: &str) { + let r = + if op == "new" { app.run("mask.add", json!({"kind": "object"})) } else { app.run("mask.addComponent", json!({"op": op, "kind": "object"})) }; + match r { + Ok(_) => { + app.ui.tool = "object".into(); + // a missing model is reported now rather than at the first click + app.toast(ctx, "Click the object to select it · ⌥-click leaves a part out"); + } + Err(e) => app.toast_error(ctx, e), + } +} + +/// The Describe field: type what to select ("sky", "the red car") and press Return. +fn describe_field(app: &mut LightcraftApp, ui: &mut egui::Ui, new_mask: bool) { + let Some((op, mut text)) = app.ui.describe.clone() else { return }; + // a new mask's field sits under the tiles; one that combines, under the selected mask + if (op == "new") != new_mask { + return; + } + ui.add_space(8.0); + let prompt = match op.as_str() { + "subtract" => "Describe what to leave out", + "intersect" => "Describe what to keep", + _ => "Describe what to select", + }; + ui.label(crate::i18n::tr(prompt)); + let mut submit = false; + ui.horizontal(|ui| { + let r = + ui.add(egui::TextEdit::singleline(&mut text).hint_text("e.g. sky · the red car · car, road").desired_width(ui.available_width() - 64.0)); + register(ui.ctx(), "maskDescribe", r.rect); + if !r.has_focus() && !r.lost_focus() && text.is_empty() { + r.request_focus(); + } + if r.lost_focus() && ui.input(|i| i.key_pressed(egui::Key::Enter)) { + submit = true; + } + if text_button(ui, "maskDescribeGo", crate::i18n::tr("Select"), false).clicked() { + submit = true; + } + }); + if ui.input(|i| i.key_pressed(egui::Key::Escape)) { + app.ui.describe = None; + return; + } + app.ui.describe = Some((op.clone(), text.clone())); + if submit && !text.trim().is_empty() { + let r = if op == "new" { + app.run("mask.add", json!({"kind": "prompt", "text": text})) + } else { + app.run("mask.addComponent", json!({"op": op, "kind": "prompt", "text": text})) + }; + match r { + Ok(_) => { + app.ui.describe = None; + if let Some(mid) = app.session.active_mask { + app.ui.detail_due = Some((ui.input(|i| i.time) + 0.3, mid)); + } + } + Err(e) => app.toast_error(ui.ctx(), e), + } + } +} + fn component_menu(app: &mut LightcraftApp, ui: &mut egui::Ui, op: &str) { + let b = ui.button(crate::i18n::tr("Object")); + register(ui.ctx(), format!("maskComp:{op}:object"), b.rect); + if b.clicked() { + start_object(app, ui.ctx(), op); + } + let b = ui.button(crate::i18n::tr("Describe…")); + register(ui.ctx(), format!("maskComp:{op}:prompt"), b.rect); + if b.clicked() { + app.ui.describe = Some((op.to_string(), String::new())); + } for (kind, label) in [ ("brush", "Brush"), ("linear", "Linear Gradient"), diff --git a/crates/ui-egui/src/state.rs b/crates/ui-egui/src/state.rs index 1a4446e3..c7e6b225 100644 --- a/crates/ui-egui/src/state.rs +++ b/crates/ui-egui/src/state.rs @@ -283,6 +283,17 @@ pub struct UiState { /// A mask being renamed in the Masks list: its id and the edited name. #[serde(skip)] pub renaming_mask: Option<(u32, String)>, + /// The Describe field (AI mask from a text prompt) while open: how the selection combines + /// (`new` mask, or `add`/`subtract`/`intersect` on the selected one) and the text typed. + #[serde(skip)] + pub describe: Option<(String, String)>, + /// The mask under the pointer in the Masks list (shown on the photo in red). + #[serde(skip)] + pub hover_mask: Option, + /// When to start the zoomed-in detail pass of an AI mask (app time) and which mask: set by + /// each click or description, so the pass runs once the clicking stops. + #[serde(skip)] + pub detail_due: Option<(f64, u32)>, /// A mask component being renamed inline: (mask id, component index, name). pub renaming_component: Option<(u32, usize, String)>, /// Close the window on the next frame (File → Quit). @@ -538,7 +549,7 @@ impl Default for UiState { histogram: true, soft_proof: false, proof: lightcraft_engine::pipeline::Proof { dest_warning: false, ..Default::default() }, - mask_overlay: true, + mask_overlay: false, mask_overlay_mode: "color".into(), mask_overlay_color: [230, 30, 40], mask_overlay_opacity: 50.0, @@ -556,6 +567,9 @@ impl Default for UiState { search: String::new(), focus_search: false, renaming_mask: None, + describe: None, + hover_mask: None, + detail_due: None, renaming_component: None, quit: false, dragging_photos: None, diff --git a/crates/ui-egui/src/tests_masking.rs b/crates/ui-egui/src/tests_masking.rs index 1e4c0d14..5be82af4 100644 --- a/crates/ui-egui/src/tests_masking.rs +++ b/crates/ui-egui/src/tests_masking.rs @@ -47,14 +47,20 @@ fn mask_overlay_keys_and_pins() { exec(&mut h, "mask.add", json!({"kind": "radial", "center": [0.3, 0.4], "rx": 0.1, "ry": 0.1})); exec(&mut h, "mask.add", json!({"kind": "linear", "start": [0.7, 0.2], "end": [0.7, 0.6]})); assert_eq!(h.app.session.active_mask, Some(2)); - // the loupe asks the renderer for the selected mask's overlay + // no overlay until a mask is hovered in the Masks list; then that mask, in red let d = develop(&h); + assert_eq!(crate::panels::detail::view_overlay(&h.app, &d), Overlay::None); + h.app.ui.hover_mask = Some(1); + assert_eq!(crate::panels::detail::view_overlay(&h.app, &d), Overlay::Mask { id: 1, view: MaskView::Color, color: [230, 30, 40], opacity: 55 }); + h.app.ui.hover_mask = None; + // O shows the selected mask all the time (and toggles it off again) + h.request("ui.key", json!({"key": "o"}), T); + assert!(h.app.ui.mask_overlay); let o = crate::panels::detail::view_overlay(&h.app, &d); assert_eq!(o, Overlay::Mask { id: 2, view: MaskView::Color, color: [230, 30, 40], opacity: 50 }); - // O toggles it, Shift+O cycles the colour (and leaves the crop overlay alone) h.request("ui.key", json!({"key": "o"}), T); - assert!(!h.app.ui.mask_overlay); assert_eq!(crate::panels::detail::view_overlay(&h.app, &d), Overlay::None); + // Shift+O cycles the colour (and leaves the crop overlay alone) h.request("ui.key", json!({"key": "o"}), T); let crop = h.app.ui.crop_overlay; let colour = h.app.ui.mask_overlay_color; diff --git a/docs/ai-masks.md b/docs/ai-masks.md new file mode 100644 index 00000000..dc3e6726 --- /dev/null +++ b/docs/ai-masks.md @@ -0,0 +1,103 @@ +# AI masks (Object and Describe) + +LightCraft selects objects with **SAM 3** (Segment Anything with Concepts, Meta 2025), run in +pure Rust by `crates/segment` on [candle](https://github.com/huggingface/candle): on the GPU +through Metal on macOS, on the CPU elsewhere (for now). + +## Using them + +In the Masking panel (M): + +- **Object**: click the thing you want; every click refines the selection. ⌥-click (Alt) a part + to leave it out. Clicks show as green (include) and red (leave out) dots. +- **Describe**: type what to select — `sky`, `trees`, `the red car` — and press Return. Every + instance the model finds (score above 0.5) joins the mask; if nothing matches, nothing is + created and you are told so. +- Several things at once: `car, road` (comma-separated) selects each and merges them. +- **+ / −** under the selected mask (and the Add / Subtract / Intersect menus) combine + selections: Describe `car`, click the mask, + ▸ Describe… `road`; − ▸ Describe… `people` takes + the people out of a sky mask; − ▸ Object then clicks remove one object. +- **Edge** (per selection, −100…100): below 0 a crisper border, above 0 a feathered one (up to 2 % + of the long edge, so previews and exports match). +- Hover a mask in the Masks list to see it in red on the photo; O shows the selected one all the + time. +- **Detail:** about a second after the last click (or right after a description), the photo + around the selection is analyzed again, zoomed in, from a sharper source (~5 s on an M4 Pro, + in the background); the result is kept as a high-resolution patch, so small objects get 5–10× + finer edges than one pass over the whole photo gives. + +The first Object or Describe on a photo analyzes it (the image encoder runs once per photo and +look: ~3 s on an Apple M4 Pro, plus a one-time ~6 s for compiling GPU kernels and loading the +model after launch). That starts in the background as soon as you pick Object; after it, each +click takes ~35 ms and each description ~0.6 s. + +The result is stored with the mask: the model's 288 × 288 mask logits over the uncropped photo +(quantized, compressed, ~10–30 KB). Renders and exports sample them at any resolution and never +need the model, so masks survive moving the library to a machine without it, and the web build +renders them. Editing the photo's look later does not recompute a mask (click again to update). + +## Installing the model + +The weights are not part of LightCraft: they are Meta's `facebook/sam3` checkpoint on Hugging +Face, under the SAM License, and access is gated. Once: + +1. Open , accept the license and wait for approval. +2. Create a read token (Hugging Face ▸ Settings ▸ Access Tokens), or log in with `hf auth login`. +3. Run the installer from the repository: + +```sh +HF_TOKEN=hf_... tools/install-sam3.sh # macOS, Linux +tools/install-sam3.sh --check # verify an installation +``` + +```powershell +$env:HF_TOKEN = "hf_..."; .\tools\install-sam3.ps1 # Windows +``` + +It downloads `model.safetensors` (3.4 GB, fp32; resumes when interrupted, checked against the +official SHA-256), `vocab.json`, `merges.txt` and the small config files into the folder +LightCraft looks in: `~/Library/Application Support/LightCraft/models/sam3/` (macOS), +`%APPDATA%\LightCraft\models\sam3` (Windows), `~/.config/lightcraft/models/sam3` (Linux). +`--dir` (or `LIGHTCRAFT_SAM3_DIR`, which LightCraft also reads) picks another folder; `--repo` +another Hugging Face repository with the same files. Restart LightCraft afterwards. + +Nothing else is needed at run time: the model runs in LightCraft itself (pure Rust, no Python). +Without the files, Object and Describe say where they are expected. + +**Developers:** the accuracy test against the Python reference needs `torch` and +`transformers` (`tools/requirements-sam3-reference.txt`): + +```sh +python3 -m venv .venv-sam3 && .venv-sam3/bin/pip install -r tools/requirements-sam3-reference.txt +.venv-sam3/bin/python tools/sam3_reference.py photo.jpg ref.safetensors +LIGHTCRAFT_SAM3_DIR= LIGHTCRAFT_SAM3_REF=ref.safetensors cargo test -p lightcraft-segment --release -- --nocapture +``` + +## Commands (control channel, MCP) + +| Command | Params | | +|---|---|---| +| `mask.add` | `{kind: "object", points?: [[x,y],…], exclude?: [[x,y],…]}` | an Object mask (empty until clicked) | +| `mask.add` | `{kind: "prompt", text}` | a Describe mask | +| `mask.addComponent` | `{op: add\|subtract\|intersect, kind: "object"\|"prompt", …}` | the same as a component | +| `mask.objectPoint` | `{x, y, exclude?: bool, id?}` | one click on the selected mask's Object selection | +| `mask.refineDetail` | `{id?, component?}` | the zoomed-in detail pass, in the background → `{started}` | +| `segment.prepare` | `{}` | analyze the active photo in the background → `{busy}` | + +Coordinates are normalized to the uncropped, oriented photo, like every mask shape. + +## Implementation notes + +- `crates/segment` ports the Hugging Face `transformers` implementation (Apache-2.0; see + `NOTICE`): the 32-layer ViT backbone with 2-D RoPE and windowed attention, the feature + pyramids, the SAM 2-style prompt encoder and two-way mask decoder for clicks, and the CLIP text + encoder, DETR encoder/decoder (box relative-position bias, presence token) and pixel decoder + for text. The CLIP tokenizer is a small BPE in `tokenizer.rs`. +- candle is pinned at 0.9.2: later releases make `candle-core` depend on `tokenizers` with the + Oniguruma C library, and the product is pure Rust. Weights are read with positional reads (no + memory map: `unsafe` stays in `lightcraft-sysmem`), and only the tensors a path needs. +- Accuracy: `crates/segment/tests/reference.rs` compares every stage against the Python + reference (`tools/sam3_reference.py`); on Metal the backbone matches to 7e-5 relative error + and masks to ~5e-5. f16 was tried and rejected (overflow, ~7 % faster). +- The engine feature `sam` (on in the desktop app, off for web and CLI) gates the model; without + it the commands report that the build has no AI masks. diff --git a/docs/parity.md b/docs/parity.md index fc94d08a..2967d311 100644 --- a/docs/parity.md +++ b/docs/parity.md @@ -69,8 +69,9 @@ Take the first one nobody is working on. then tune tone, highlights, texture/clarity/dehaze, NR and sharpening against it. 5. **LR-EDIT-OPTICS-PROFILE** (P1): a lens-profile database of our own (embedded DNG/maker corrections work today). 6. **AI masks and Enhance** (LR-MASK-SUBJECT / SKY / PEOPLE / OBJECTS, LR-EDIT-DETAIL-DENOISE, SUPERRES, LENSBLUR): - blocked on a model strategy (permissively licensed weights or our own training, pure-Rust inference). A maintainer - decision, not just engineering. + Object and Describe masks run SAM 3 in pure Rust (`crates/segment`) with user-installed weights; Subject / Sky / + People could use the same model with fixed prompts. Bundling or recommending weights (SAM License) and denoise / + super-resolution models remain a maintainer decision. 7. **HDR** (Q. HDR, LR-EXP-HDR), **video** (R. Video), **Classic output modules** (Map view, Book, Slideshow module, Print, publish): large, well understood, lower priority than 1–5. @@ -328,7 +329,7 @@ under the UI command that wraps it (a no-op error); the UI command now wins. | LR-MASK-SUBJECT | Select subject | P2 | 🟡 | `cmd:mask.add` (`subject`), `crates/pipeline/src/masks.rs` | saliency heuristic, no segmentation model | | LR-MASK-SKY | Select sky | P2 | 🟡 | `cmd:mask.add` (`sky`) | heuristic | | LR-MASK-BACKGROUND | Select background | P2 | 🟡 | `cmd:mask.add` (`background`) | inverse of the subject heuristic | -| LR-MASK-OBJECTS | Object selection | P2 | ⬜ | | shape exists (falls back to the subject heuristic); no UI | +| LR-MASK-OBJECTS | Object selection | P2 | 🟡 | `cmd:mask.add` (`object`, `prompt`), `cmd:mask.objectPoint`, `cmd:segment.prepare`, `crates/segment`, `crates/engine/src/segment.rs` | SAM 3 in pure Rust (candle; Metal on macOS, CPU elsewhere): Object tile → click to include, ⌥-click to leave out; Describe tile → a text prompt selects every instance ("sky", "the red car"); both also as Add/Subtract/Intersect components. The segmentation is stored with the mask (288² logits), so renders and exports never need the model. Weights are not bundled: the user installs the `facebook/sam3` checkpoint (SAM License) — see docs/ai-masks.md. + / − under the selected mask combine selections; comma lists (`car, road`); a zoomed-in detail pass (`cmd:mask.refineDetail`, background) for 5–10× finer edges on small objects; per-selection Edge (hard ↔ soft); hovering a mask shows it in red. No brush/box object mode; the first click on a photo waits for its analysis (~4 s on an M4 Pro) | | LR-MASK-PEOPLE | People parts | P2 | ⬜ | | | | LR-MASK-LANDSCAPE | Landscape classes | P2 | ⬜ | | shape exists, evaluates empty | | LR-MASK-BRUSH | Brush mask | P0 | ✅ | `cmd:tool.brush`, `cmd:mask.brushStroke` (`autoMask`), `crates/pipeline/src/masks.rs` | size/feather/flow/density/erase; Auto Mask: dabs weighted by similarity to the colour under the dab centre, refined by a guided filter on luminance (CPU + GPU); no A/B brushes, no pressure | diff --git a/tools/install-sam3.ps1 b/tools/install-sam3.ps1 new file mode 100644 index 00000000..bd43fd98 --- /dev/null +++ b/tools/install-sam3.ps1 @@ -0,0 +1,40 @@ +# Install the SAM 3 model for LightCraft's Object and Describe masks on Windows +# (see docs/ai-masks.md; tools/install-sam3.sh is the macOS/Linux version). +# +# facebook/sam3 is gated: accept the SAM License at https://huggingface.co/facebook/sam3, wait +# for approval, then: $env:HF_TOKEN = "hf_..." ; .\tools\install-sam3.ps1 +param( + [string]$Dir = $(if ($env:LIGHTCRAFT_SAM3_DIR) { $env:LIGHTCRAFT_SAM3_DIR } else { Join-Path $env:APPDATA "LightCraft\models\sam3" }), + [string]$Repo = "facebook/sam3", + [switch]$Check +) +$ErrorActionPreference = "Stop" +$Sha256 = "6d06f0a5f84e435071fe6603e61d0b4cc7b40e0d39d487cfd4d67d8cc11cc14a" +$Files = "config.json", "vocab.json", "merges.txt", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "processor_config.json" + +function Test-Install { + foreach ($f in "model.safetensors", "vocab.json", "merges.txt") { + if (-not (Test-Path (Join-Path $Dir $f))) { throw "missing: $(Join-Path $Dir $f)" } + } + Write-Host "checking model.safetensors (SHA-256 of 3.4 GB)..." + $got = (Get-FileHash (Join-Path $Dir "model.safetensors") -Algorithm SHA256).Hash.ToLower() + if ($got -ne $Sha256) { throw "model.safetensors does not match the official checkpoint (got $got)" } + Write-Host "SAM 3 is installed in $Dir" +} + +if ($Check) { Test-Install; exit 0 } + +$Token = if ($env:HF_TOKEN) { $env:HF_TOKEN } else { $tf = Join-Path $HOME ".cache\huggingface\token"; if (Test-Path $tf) { (Get-Content $tf -Raw).Trim() } } +if (-not $Token -and $Repo -eq "facebook/sam3") { + throw "No Hugging Face token. Accept the license at https://huggingface.co/facebook/sam3, then set HF_TOKEN or run 'hf auth login'." +} +New-Item -ItemType Directory -Force -Path $Dir | Out-Null +$Headers = @{} +if ($Token) { $Headers["Authorization"] = "Bearer $Token" } +$Base = "https://huggingface.co/$Repo/resolve/main" +foreach ($f in $Files + "model.safetensors") { + Write-Host "fetching $f" + Invoke-WebRequest -Uri "$Base/$f" -Headers $Headers -OutFile (Join-Path $Dir $f) +} +Test-Install +Write-Host "Restart LightCraft; Object and Describe in the Masking panel now use it." diff --git a/tools/install-sam3.sh b/tools/install-sam3.sh new file mode 100755 index 00000000..86554c39 --- /dev/null +++ b/tools/install-sam3.sh @@ -0,0 +1,110 @@ +#!/bin/sh +# Install the SAM 3 model for LightCraft's Object and Describe masks (see docs/ai-masks.md). +# +# The weights are not part of LightCraft: they are Meta's facebook/sam3 checkpoint on Hugging +# Face, under the SAM License. Access is gated: open https://huggingface.co/facebook/sam3, accept +# the license, wait for approval, then run this with a token that can read it: +# +# HF_TOKEN=hf_... tools/install-sam3.sh +# +# (or log in once with `hf auth login`; the token it saves is used). Downloads resume when +# interrupted, and the 3.4 GB weights are checked against the official SHA-256. +# +# Options: +# --dir DIR where to install (default: the folder LightCraft looks in; or LIGHTCRAFT_SAM3_DIR) +# --repo NAME another Hugging Face repo with the same files (default: facebook/sam3) +# --check only verify an existing installation +set -eu + +REPO="facebook/sam3" +CHECK_ONLY=0 +SHA256="6d06f0a5f84e435071fe6603e61d0b4cc7b40e0d39d487cfd4d67d8cc11cc14a" +SIZE=3439938512 +FILES="config.json vocab.json merges.txt tokenizer.json tokenizer_config.json special_tokens_map.json processor_config.json" + +case "$(uname -s)" in + Darwin) DEFAULT_DIR="$HOME/Library/Application Support/LightCraft/models/sam3" ;; + *) DEFAULT_DIR="${XDG_CONFIG_HOME:-$HOME/.config}/lightcraft/models/sam3" ;; +esac +DIR="${LIGHTCRAFT_SAM3_DIR:-$DEFAULT_DIR}" + +while [ $# -gt 0 ]; do + case "$1" in + --dir) DIR="$2"; shift 2 ;; + --repo) REPO="$2"; shift 2 ;; + --check) CHECK_ONLY=1; shift ;; + -h|--help) sed -n '2,20p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;; + *) echo "unknown option: $1 (see --help)" >&2; exit 2 ;; + esac +done + +sha256() { + if command -v shasum >/dev/null 2>&1; then shasum -a 256 "$1" | cut -d' ' -f1 + else sha256sum "$1" | cut -d' ' -f1; fi +} + +verify() { + for f in model.safetensors vocab.json merges.txt; do + [ -s "$DIR/$f" ] || { echo "missing: $DIR/$f" >&2; return 1; } + done + echo "checking model.safetensors (SHA-256 of 3.4 GB, takes a few seconds)…" + got=$(sha256 "$DIR/model.safetensors") + if [ "$got" != "$SHA256" ]; then + echo "model.safetensors does not match the official checkpoint (got $got)" >&2 + return 1 + fi + echo "SAM 3 is installed in $DIR" +} + +if [ "$CHECK_ONLY" = 1 ]; then verify; exit $?; fi + +command -v curl >/dev/null 2>&1 || { echo "curl is required" >&2; exit 1; } + +TOKEN="${HF_TOKEN:-${HUGGING_FACE_HUB_TOKEN:-}}" +if [ -z "$TOKEN" ]; then + for t in "${HF_HOME:-$HOME/.cache/huggingface}/token" "$HOME/.huggingface/token"; do + [ -s "$t" ] && TOKEN=$(cat "$t") && break + done +fi +if [ -z "$TOKEN" ] && [ "$REPO" = "facebook/sam3" ]; then + echo "No Hugging Face token found. facebook/sam3 is gated: accept the license at" >&2 + echo "https://huggingface.co/facebook/sam3, then set HF_TOKEN or run \`hf auth login\`." >&2 + exit 1 +fi + +mkdir -p "$DIR" +BASE="https://huggingface.co/$REPO/resolve/main" + +fetch() { # fetch FILE DEST [resume] + set -- "$1" "$2" "${3:-}" + if [ -n "$TOKEN" ]; then auth="Authorization: Bearer $TOKEN"; else auth="X-No-Auth: 1"; fi + code=$(curl -sSL ${3:+-C -} -H "$auth" -w '%{http_code}' -o "$2" "$BASE/$1") || code=$? + case "$code" in + 200|206|416) return 0 ;; + 401|403) echo "access denied to $REPO/$1 (has your request for facebook/sam3 been approved?)" >&2; return 1 ;; + *) echo "download of $REPO/$1 failed ($code)" >&2; return 1 ;; + esac +} + +for f in $FILES; do + echo "fetching $f" + fetch "$f" "$DIR/$f" +done + +if [ -s "$DIR/model.safetensors" ] && [ "$(sha256 "$DIR/model.safetensors")" = "$SHA256" ]; then + echo "model.safetensors already installed" +else + echo "fetching model.safetensors (3.4 GB; resumes if interrupted)" + part="$DIR/model.safetensors.part" + [ -s "$DIR/model.safetensors" ] && mv -f "$DIR/model.safetensors" "$part" + fetch model.safetensors "$part" resume + have=$(wc -c < "$part" | tr -d ' ') + if [ "$have" != "$SIZE" ]; then + echo "incomplete download ($have of $SIZE bytes); run the script again to resume" >&2 + exit 1 + fi + mv -f "$part" "$DIR/model.safetensors" +fi + +verify +echo "Restart LightCraft; Object and Describe in the Masking panel now use it." diff --git a/tools/requirements-sam3-reference.txt b/tools/requirements-sam3-reference.txt new file mode 100644 index 00000000..cf8f263a --- /dev/null +++ b/tools/requirements-sam3-reference.txt @@ -0,0 +1,8 @@ +# Python packages for tools/sam3_reference.py (reference tensors for crates/segment/tests/reference.rs). +# python3 -m venv .venv-sam3 && .venv-sam3/bin/pip install -r tools/requirements-sam3-reference.txt +torch>=2.6 +torchvision>=0.21 +transformers>=4.57 +pillow +safetensors +numpy diff --git a/tools/sam3_reference.py b/tools/sam3_reference.py new file mode 100644 index 00000000..66d9def9 --- /dev/null +++ b/tools/sam3_reference.py @@ -0,0 +1,71 @@ +"""Reference tensors for the Rust SAM 3 port (crates/segment/tests/reference.rs). + +Runs the Hugging Face transformers implementation (fp32, CPU) on one photo and writes the +shared pixel input plus intermediate and final outputs of the click (tracker) and text +(detector) paths. Usage (needs torch, torchvision, transformers, pillow, safetensors): + LIGHTCRAFT_SAM3_DIR= python tools/sam3_reference.py photo.jpg ref.safetensors +The text prompts below suit a photo of the moon in a blue sky; change them for other photos. +""" +import os, sys, time +import numpy as np +import torch +from PIL import Image +from safetensors.torch import save_file +from transformers import Sam3TrackerModel, Sam3Model, CLIPTokenizer + +MODEL = os.environ.get("LIGHTCRAFT_SAM3_DIR", os.path.expanduser("~/Library/Application Support/LightCraft/models/sam3")) +img_path, out_path = sys.argv[1], sys.argv[2] +torch.set_grad_enabled(False) + +img = Image.open(img_path).convert("RGB") +W, H = img.size +# Same preprocessing as the Rust side: bilinear resize to 1008², /255, (x-0.5)/0.5. +r = img.resize((1008, 1008), Image.BILINEAR) +px = torch.from_numpy(np.asarray(r, dtype=np.float32) / 255.0).permute(2, 0, 1)[None] +px = (px - 0.5) / 0.5 +out = {"pixel_values": px.contiguous()} + +t = time.time() +trk = Sam3TrackerModel.from_pretrained(MODEL).eval() +vis = trk.vision_encoder(px) +out["backbone"] = vis.last_hidden_state.contiguous() +for i, f in enumerate(vis.fpn_hidden_states): + out[f"trk_fpn{i}"] = f.contiguous() +# points in 1008 space: one positive, one negative +pts = torch.tensor([[[[0.5 * 1008, 0.55 * 1008], [0.2 * 1008, 0.2 * 1008]]]], dtype=torch.float32) +lbl = torch.tensor([[[1, 0]]], dtype=torch.int32) +o = trk(pixel_values=px, input_points=pts, input_labels=lbl, multimask_output=False) +out["trk_points"] = pts.contiguous() +out["trk_masks"] = o.pred_masks.contiguous() +out["trk_iou"] = o.iou_scores.contiguous() +o3 = trk(pixel_values=px, input_points=pts[:, :, :1], input_labels=lbl[:, :, :1], multimask_output=True) +out["trk_multi_masks"] = o3.pred_masks.contiguous() +out["trk_multi_iou"] = o3.iou_scores.contiguous() +print("tracker", time.time() - t, o.pred_masks.shape, o.iou_scores, file=sys.stderr) +del trk + +t = time.time() +det = Sam3Model.from_pretrained(MODEL).eval() +tok = CLIPTokenizer.from_pretrained(MODEL) +for i, text in enumerate(["moon", "sky", "bird"]): + enc = tok(text, padding="max_length", max_length=32, return_tensors="pt") + out[f"det{i}_ids"] = enc.input_ids.to(torch.int64).contiguous() + out[f"det{i}_mask"] = enc.attention_mask.to(torch.int64).contiguous() + tf = det.get_text_features(input_ids=enc.input_ids, attention_mask=enc.attention_mask) + out[f"det{i}_text"] = tf.pooler_output.contiguous() + o = det(pixel_values=px, input_ids=enc.input_ids, attention_mask=enc.attention_mask) + out[f"det{i}_logits"] = o.pred_logits.contiguous() + out[f"det{i}_presence"] = o.presence_logits.contiguous() + out[f"det{i}_boxes"] = o.pred_boxes.contiguous() + out[f"det{i}_masks"] = o.pred_masks.contiguous() + sc = (o.pred_logits.sigmoid() * o.presence_logits.sigmoid())[0] + print(text, "presence", o.presence_logits.sigmoid().item(), "top", sc.topk(3).values.tolist(), file=sys.stderr) +for i, f in enumerate(det.vision_encoder(px).fpn_hidden_states): + out[f"det_fpn{i}"] = f.contiguous() +print("detector", time.time() - t, file=sys.stderr) +# tokenizer cases for the Rust BPE +cases = ["person", "Red Car", "a dog's toy", "sky, clouds!", "café 2 people", "tree-line"] +for i, c in enumerate(cases): + out[f"tok{i}"] = tok(c, padding="max_length", max_length=32, return_tensors="pt").input_ids.to(torch.int64).contiguous() +save_file(out, out_path, metadata={"image": img_path, "W": str(W), "H": str(H), "cases": "|".join(cases)}) +print("saved", out_path, file=sys.stderr) diff --git a/xtask/src/layers.rs b/xtask/src/layers.rs index b5c63142..28c99462 100644 --- a/xtask/src/layers.rs +++ b/xtask/src/layers.rs @@ -49,6 +49,7 @@ pub const TABLE: &[(&str, Class)] = &[ ("preview", Class::Layer(3)), ("export", Class::Layer(3)), ("merge", Class::Layer(3)), + ("segment", Class::Layer(3)), ("engine", Class::Layer(4)), ("ui-egui", Class::Layer(5)), ("mcp", Class::Layer(5)),