From 6468bfc06bf96ad3a62578e549da927512f987a1 Mon Sep 17 00:00:00 2001 From: Guinther Kovalski Date: Tue, 6 Oct 2026 20:15:04 -0300 Subject: [PATCH 1/4] 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)), From b52f43d1d50f156f262d61356d50d3010f4acaf7 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 7 Oct 2026 09:56:49 -0400 Subject: [PATCH 2/4] M12: lightcraft-segment: Apache-2.0, checked weights, pure-Rust model download MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - The crate ports Apache-2.0-only Hugging Face transformers code: it is now `license = "Apache-2.0"`, and every ported file carries the modification notice Apache-2.0 §4(b) asks for. - Weights::open checks every tensor's byte range (and the header) against the file's length: a truncated or damaged download is an error at load, never a short read later; regression test with hostile offsets. - fetch/: the model download, only when asked for. An ordered mirror list (env LIGHTCRAFT_SAM3_MIRRORS, a mirrors file, DEFAULT_MIRRORS — empty until LightCraft's CDN exists, TODO for the maintainer), fallback to the next mirror, resume via Range, redirects, connect and stall timeouts, cancel, size caps, exact size + SHA-256 (pinned from tools/install-sam3.sh) before an atomic rename; damaged files are deleted, HTTP error bodies never written. - HTTP/1.1 client over std sockets + rustls with the RustCrypto provider (rustls-rustcrypto) and webpki-roots: no ring / aws-lc-rs, nothing compiled from C (no new `cc` users in the app's build graph). - Tests against a local TcpListener server (no internet): fallback, redirects, chunked bodies, resume, 416, hash mismatch, too much data, stalls, cancel. - *.safetensors are gitignored. Co-Authored-By: Claude Opus 5.5 --- .gitignore | 3 + Cargo.lock | 755 +++++++++++++++++++++++++++++- crates/segment/Cargo.toml | 6 +- crates/segment/src/clip.rs | 4 + crates/segment/src/detector.rs | 4 + crates/segment/src/fetch/http.rs | 541 +++++++++++++++++++++ crates/segment/src/fetch/mod.rs | 424 +++++++++++++++++ crates/segment/src/fetch/tests.rs | 314 +++++++++++++ crates/segment/src/lib.rs | 13 +- crates/segment/src/mask.rs | 4 + crates/segment/src/neck.rs | 4 + crates/segment/src/nn.rs | 4 + crates/segment/src/tokenizer.rs | 4 + crates/segment/src/tracker.rs | 4 + crates/segment/src/vit.rs | 4 + crates/segment/src/weights.rs | 47 +- 16 files changed, 2119 insertions(+), 16 deletions(-) create mode 100644 crates/segment/src/fetch/http.rs create mode 100644 crates/segment/src/fetch/mod.rs create mode 100644 crates/segment/src/fetch/tests.rs diff --git a/.gitignore b/.gitignore index 823f9276..51531a34 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,6 @@ plan/ .DS_Store .claude/ /craft-fonts +# model weights are never committed (SAM 3: docs/ai-masks.md) +*.safetensors +*.safetensors.part diff --git a/Cargo.lock b/Cargo.lock index 6a07ead6..3f086161 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -146,6 +146,41 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "aead" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d122413f284cf2d62fb1b7db97e02edb8cda96d769b16e443a4f6195e35662b0" +dependencies = [ + "crypto-common", + "generic-array", +] + +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures", +] + +[[package]] +name = "aes-gcm" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "831010a0f742e1209b3bcea8fab6a8e149051ba6099432c8cb2cc117dec3ead1" +dependencies = [ + "aead", + "aes", + "cipher", + "ctr", + "ghash", + "subtle", +] + [[package]] name = "ahash" version = "0.8.12" @@ -543,6 +578,18 @@ dependencies = [ "windows-link", ] +[[package]] +name = "base16ct" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c7f02d4ea65f2c1853089ffd8d2787bdbc63de2f0d29dedbcf8ccdfa0ccd4cf" + +[[package]] +name = "base64ct" +version = "1.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" + [[package]] name = "bit-set" version = "0.8.0" @@ -600,6 +647,15 @@ version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d8c1fef690941d3e7788d328517591fecc684c084084702d6ff1641e993699a" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "block2" version = "0.5.1" @@ -771,7 +827,7 @@ dependencies = [ "num_cpus", "objc2-foundation 0.3.2", "objc2-metal 0.3.2", - "rand", + "rand 0.9.5", "rand_distr", "rayon", "safetensors 0.7.0", @@ -856,6 +912,41 @@ dependencies = [ "libc", ] +[[package]] +name = "chacha20" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3613f74bd2eac03dad61bd53dbe620703d4371614fe0bc3b9f04dd36fe4e818" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures", +] + +[[package]] +name = "chacha20poly1305" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10cd79432192d1c0f4e1a0fef9527696cc039165d729fb41b3f4f4f354c2dc35" +dependencies = [ + "aead", + "chacha20", + "cipher", + "poly1305", + "zeroize", +] + +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common", + "inout", + "zeroize", +] + [[package]] name = "clipboard-win" version = "5.4.1" @@ -910,6 +1001,12 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "const-oid" +version = "0.9.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" + [[package]] name = "core-foundation" version = "0.9.4" @@ -950,6 +1047,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.5.2" @@ -999,12 +1105,81 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +[[package]] +name = "crypto-bigint" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0dc92fb57ca44df6db8059111ab3af99a63d5d0f8375d9972e319a379c6bab76" +dependencies = [ + "generic-array", + "rand_core 0.6.4", + "subtle", + "zeroize", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "ctr" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835" +dependencies = [ + "cipher", +] + [[package]] name = "cursor-icon" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f27ae1dd37df86211c42e150270f82743308803d90a6f6e6651cd730d5e1732f" +[[package]] +name = "curve25519-dalek" +version = "4.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97fb8b7c4503de7d6ae7b42ab72a5a59857b4c937ec27a3d4539dba95b5ab2be" +dependencies = [ + "cfg-if", + "cpufeatures", + "curve25519-dalek-derive", + "digest", + "fiat-crypto", + "rustc_version", + "subtle", + "zeroize", +] + +[[package]] +name = "curve25519-dalek-derive" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "der" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" +dependencies = [ + "const-oid", + "pem-rfc7468", + "zeroize", +] + [[package]] name = "dhat" version = "0.3.3" @@ -1021,6 +1196,18 @@ dependencies = [ "thousands", ] +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "const-oid", + "crypto-common", + "subtle", +] + [[package]] name = "dispatch" version = "0.2.0" @@ -1085,6 +1272,20 @@ version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e1d926b4d407d372f141f93bb444696142c29d32962ccbd3531117cf3aa0bfa9" +[[package]] +name = "ecdsa" +version = "0.16.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee27f32b5c5292967d2d4a9d7f1e0b0aed2c15daded5a60300e4abb9d8020bca" +dependencies = [ + "der", + "digest", + "elliptic-curve", + "rfc6979", + "signature", + "spki", +] + [[package]] name = "ecolor" version = "0.36.2" @@ -1095,6 +1296,30 @@ dependencies = [ "emath", ] +[[package]] +name = "ed25519" +version = "2.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53" +dependencies = [ + "pkcs8", + "signature", +] + +[[package]] +name = "ed25519-dalek" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9" +dependencies = [ + "curve25519-dalek", + "ed25519", + "serde", + "sha2", + "subtle", + "zeroize", +] + [[package]] name = "eframe" version = "0.36.2" @@ -1212,6 +1437,27 @@ version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" +[[package]] +name = "elliptic-curve" +version = "0.13.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5e6043086bf7973472e0c7dff2142ea0b680d30e18d9cc40f267efbf222bd47" +dependencies = [ + "base16ct", + "crypto-bigint", + "digest", + "ff", + "generic-array", + "group", + "hkdf", + "pem-rfc7468", + "pkcs8", + "rand_core 0.6.4", + "sec1", + "subtle", + "zeroize", +] + [[package]] name = "emath" version = "0.36.2" @@ -1389,6 +1635,22 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b97b65636e5b9ef369943878ac74335ba1c55c1cb6adbf1e2c293c624248d693" +[[package]] +name = "ff" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0b50bfb653653f9ca9095b427bed08ab8d75a137839d9ad64eb11810d5b6393" +dependencies = [ + "rand_core 0.6.4", + "subtle", +] + +[[package]] +name = "fiat-crypto" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d" + [[package]] name = "find-msvc-tools" version = "0.1.14" @@ -1414,7 +1676,7 @@ checksum = "719a903cc23e4a89e87962c2a80fdb45cdaad0983a89bd150bb57b4c8571a7d5" dependencies = [ "half", "num-traits", - "rand", + "rand 0.9.5", "rand_distr", ] @@ -1759,6 +2021,17 @@ dependencies = [ "seq-macro", ] +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", + "zeroize", +] + [[package]] name = "gethostname" version = "1.1.0" @@ -1769,6 +2042,17 @@ dependencies = [ "windows-link", ] +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + [[package]] name = "getrandom" version = "0.3.4" @@ -1792,6 +2076,16 @@ dependencies = [ "r-efi 6.0.0", ] +[[package]] +name = "ghash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0d8a4362ccb29cb0b265253fb0a2728f592895ee6854fd9bc13f2ffda266ff1" +dependencies = [ + "opaque-debug", + "polyval", +] + [[package]] name = "gif" version = "0.14.2" @@ -1927,6 +2221,17 @@ dependencies = [ "windows", ] +[[package]] +name = "group" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0f9ef7462f7c099f518d754361858f86d8a07af53ba9af0fe635bbccb151a63" +dependencies = [ + "ff", + "rand_core 0.6.4", + "subtle", +] + [[package]] name = "guillotiere" version = "0.7.0" @@ -1946,7 +2251,7 @@ dependencies = [ "cfg-if", "crunchy", "num-traits", - "rand", + "rand 0.9.5", "rand_distr", "zerocopy", ] @@ -2003,6 +2308,24 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "hkdf" +version = "0.12.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b5f8eb2ad728638ea2c7d47a21db23b7b58a72ed6a38256b8a1849f15fbbdf7" +dependencies = [ + "hmac", +] + +[[package]] +name = "hmac" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" +dependencies = [ + "digest", +] + [[package]] name = "image" version = "0.25.10" @@ -2048,6 +2371,15 @@ dependencies = [ "hashbrown 0.17.1", ] +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "generic-array", +] + [[package]] name = "interpolate_name" version = "0.2.4" @@ -2397,6 +2729,9 @@ name = "lazy_static" version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" +dependencies = [ + "spin", +] [[package]] name = "libc" @@ -2705,9 +3040,13 @@ dependencies = [ "candle-core", "candle-nn", "log", + "rustls", + "rustls-rustcrypto", "serde", "serde_json", + "sha2", "thiserror 2.0.21", + "webpki-roots", ] [[package]] @@ -3065,6 +3404,22 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-bigint-dig" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e661dda6640fad38e827a6d4a310ff4763082116fe217f279885c97f511bb0b7" +dependencies = [ + "lazy_static", + "libm", + "num-integer", + "num-iter", + "num-traits", + "rand 0.8.8", + "smallvec", + "zeroize", +] + [[package]] name = "num-complex" version = "0.4.6" @@ -3496,6 +3851,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "opaque-debug" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" + [[package]] name = "orbclient" version = "0.3.55" @@ -3534,6 +3895,30 @@ dependencies = [ "ttf-parser", ] +[[package]] +name = "p256" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9863ad85fa8f4460f9c48cb909d38a0d689dba1f6f6988a5e3e0d31071bcd4b" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2", +] + +[[package]] +name = "p384" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe42f1670a52a47d448f14b6a5c61dd78fce51856e68edaa38f7ae3a46b8d6b6" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2", +] + [[package]] name = "parking" version = "2.2.1" @@ -3575,6 +3960,15 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35fb2e5f958ec131621fdd531e9fc186ed768cbe395337403ae56c17a74c68ec" +[[package]] +name = "pem-rfc7468" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88b39c9bfcfc231068454382784bb460aae594343fb030d46e9f50a645418412" +dependencies = [ + "base64ct", +] + [[package]] name = "peniko" version = "0.6.1" @@ -3674,6 +4068,38 @@ dependencies = [ "futures-io", ] +[[package]] +name = "pkcs1" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8ffb9f10fa047879315e6625af03c164b16962a5368d724ed16323b68ace47f" +dependencies = [ + "der", + "pkcs8", + "spki", +] + +[[package]] +name = "pkcs5" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e847e2c91a18bfa887dd028ec33f2fe6f25db77db3619024764914affe8b69a6" +dependencies = [ + "der", + "spki", +] + +[[package]] +name = "pkcs8" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" +dependencies = [ + "der", + "pkcs5", + "spki", +] + [[package]] name = "pkg-config" version = "0.3.34" @@ -3725,6 +4151,17 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bc6355899e1c9462875b6757c79f3caa011a1fdae12bbb1a2e72dd1f234f8336" +[[package]] +name = "poly1305" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8159bd90725d2df49889a078b54f4f79e87f1f8a8444194cdca81d38f5393abf" +dependencies = [ + "cpufeatures", + "opaque-debug", + "universal-hash", +] + [[package]] name = "polycool" version = "0.4.0" @@ -3734,6 +4171,18 @@ dependencies = [ "arrayvec", ] +[[package]] +name = "polyval" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1fe60d06143b2430aa532c94cfe9e29783047f06c0d7fd359a9a51b729fa25" +dependencies = [ + "cfg-if", + "cpufeatures", + "opaque-debug", + "universal-hash", +] + [[package]] name = "portable-atomic" version = "1.15.0" @@ -3764,6 +4213,15 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8cf8e6a8aa66ce33f63993ffc4ea4271eb5b0530a9002db8455ea6050c77bfa" +[[package]] +name = "primeorder" +version = "0.13.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "353e1ca18966c16d9deb1c69278edbc5f194139612772bd9537af60ac231e1e6" +dependencies = [ + "elliptic-curve", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -3811,8 +4269,8 @@ dependencies = [ "bit-vec 0.8.0", "bitflags 2.13.2", "num-traits", - "rand", - "rand_chacha", + "rand 0.9.5", + "rand_chacha 0.9.0", "rand_xorshift", "regex-syntax", "rusty-fork", @@ -3914,14 +4372,34 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "rand" +version = "0.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e058c7de0b26af77780c769414d6257830bb240f3c38477dbc2c16e5f54d6d4c" +dependencies = [ + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + [[package]] name = "rand" version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" dependencies = [ - "rand_chacha", - "rand_core", + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", ] [[package]] @@ -3931,7 +4409,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", ] [[package]] @@ -3950,7 +4437,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" dependencies = [ "num-traits", - "rand", + "rand 0.9.5", ] [[package]] @@ -3959,7 +4446,7 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a" dependencies = [ - "rand_core", + "rand_core 0.9.5", ] [[package]] @@ -3995,8 +4482,8 @@ dependencies = [ "num-traits", "paste", "profiling", - "rand", - "rand_chacha", + "rand 0.9.5", + "rand_chacha 0.9.0", "simd_helpers", "thiserror 2.0.21", "v_frame", @@ -4121,6 +4608,16 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19b30a45b0cd0bcca8037f3d0dc3421eaf95327a17cad11964fb8179b4fc4832" +[[package]] +name = "rfc6979" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dd2a808d456c4a54e300a23e9f5a67e122c3024119acbfd73e3bf664491cb2" +dependencies = [ + "hmac", + "subtle", +] + [[package]] name = "rfd" version = "0.17.2" @@ -4154,6 +4651,41 @@ version = "0.8.53" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47b34b781b31e5d73e9fbc8689c70551fd1ade9a19e3e28cfec8580a79290cc4" +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "rsa" +version = "0.9.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8573f03f5883dcaebdfcf4725caa1ecb9c15b2ef50c43a07b816e06799bb12d" +dependencies = [ + "const-oid", + "digest", + "num-bigint-dig", + "num-integer", + "num-traits", + "pkcs1", + "pkcs8", + "rand_core 0.6.4", + "sha2", + "signature", + "spki", + "subtle", + "zeroize", +] + [[package]] name = "rustc-demangle" version = "0.1.28" @@ -4207,6 +4739,80 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "rustls" +version = "0.23.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" +dependencies = [ + "once_cell", + "rustls-pki-types", + "rustls-webpki 0.103.15", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-pki-types" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96" +dependencies = [ + "zeroize", +] + +[[package]] +name = "rustls-rustcrypto" +version = "0.0.2-alpha" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f12052947763ab8515f753315357599e9b0b4dab3b8ba15f30f725fe6d025557" +dependencies = [ + "aead", + "aes-gcm", + "chacha20poly1305", + "crypto-common", + "der", + "digest", + "ecdsa", + "ed25519-dalek", + "hmac", + "p256", + "p384", + "paste", + "pkcs8", + "rand_core 0.6.4", + "rsa", + "rustls", + "rustls-pki-types", + "rustls-webpki 0.102.8", + "sec1", + "sha2", + "signature", + "x25519-dalek", +] + +[[package]] +name = "rustls-webpki" +version = "0.102.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64ca1bc8749bd4cf37b5ce386cc146580777b4e8572c7b97baf22c83f444bee9" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + [[package]] name = "rustversion" version = "1.0.23" @@ -4267,6 +4873,20 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "sec1" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3e97a565f76233a6003f9f5c54be1d9c5bdfa3eccfb189469f11ec4901c47dc" +dependencies = [ + "base16ct", + "der", + "generic-array", + "pkcs8", + "subtle", + "zeroize", +] + [[package]] name = "self_cell" version = "1.3.0" @@ -4348,6 +4968,17 @@ dependencies = [ "serde_core", ] +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "shlex" version = "2.0.1" @@ -4364,6 +4995,16 @@ dependencies = [ "libc", ] +[[package]] +name = "signature" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" +dependencies = [ + "digest", + "rand_core 0.6.4", +] + [[package]] name = "simd-adler32" version = "0.3.10" @@ -4504,6 +5145,12 @@ dependencies = [ "serde", ] +[[package]] +name = "spin" +version = "0.9.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3763264f6b73151db08c50ff20d7d8a0b8796e021cdea7ceedad07b80155fa0e" + [[package]] name = "spirv" version = "0.4.0+sdk-1.4.341.0" @@ -4513,6 +5160,16 @@ dependencies = [ "bitflags 2.13.2", ] +[[package]] +name = "spki" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d" +dependencies = [ + "base64ct", + "der", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -4525,6 +5182,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + [[package]] name = "syn" version = "2.0.119" @@ -4778,6 +5441,12 @@ version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e28f89b80c87b8fb0cf04ab448d5dd0dd0ade2f8891bae878de66a75a28600e" +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + [[package]] name = "uds_windows" version = "1.2.1" @@ -4854,6 +5523,22 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" +[[package]] +name = "universal-hash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea" +dependencies = [ + "crypto-common", + "subtle", +] + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + [[package]] name = "uuid" version = "1.26.1" @@ -4928,6 +5613,12 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + [[package]] name = "wasip2" version = "1.0.4+wasi-0.2.12" @@ -5148,6 +5839,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki-roots" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "weezl" version = "0.1.12" @@ -5729,6 +6429,17 @@ version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ea6fc2961e4ef194dcbfe56bb845534d0dc8098940c7e5c012a258bfec6701bd" +[[package]] +name = "x25519-dalek" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7e468321c81fb07fa7f4c636c3972b9100f0346e5b6a9f2bd0603a52f7ed277" +dependencies = [ + "curve25519-dalek", + "rand_core 0.6.4", + "zeroize", +] + [[package]] name = "xcursor" version = "0.3.11" @@ -5969,6 +6680,26 @@ dependencies = [ "synstructure 0.14.0", ] +[[package]] +name = "zeroize" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13084392c5e4bc371903e2935a5eaeed24905a7511356b883835e18a78f6879" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c50655cbb0fe3fc43170059e702f1ce5e19b84cec58dc87b037a09935c2f328" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "zip" version = "7.2.0" diff --git a/crates/segment/Cargo.toml b/crates/segment/Cargo.toml index 336b052f..bd7b0d66 100644 --- a/crates/segment/Cargo.toml +++ b/crates/segment/Cargo.toml @@ -2,7 +2,7 @@ name = "lightcraft-segment" version.workspace = true edition.workspace = true -license.workspace = true +license = "Apache-2.0" description = "LightCraft segmentation: SAM 3 (Segment Anything with Concepts) inference in pure Rust (L3)" [lints] @@ -18,6 +18,10 @@ log = { workspace = true } [target.'cfg(not(target_arch = "wasm32"))'.dependencies] candle-core = { workspace = true } candle-nn = { workspace = true } +rustls = { version = "0.23", default-features = false, features = ["std", "tls12"] } +rustls-rustcrypto = "0.0.2-alpha" +webpki-roots = "1" +sha2 = { version = "0.10", default-features = false, features = ["std"] } [target.'cfg(target_os = "macos")'.dependencies] candle-core = { workspace = true, features = ["metal"] } diff --git a/crates/segment/src/clip.rs b/crates/segment/src/clip.rs index c392d5c5..b69c148a 100644 --- a/crates/segment/src/clip.rs +++ b/crates/segment/src/clip.rs @@ -1,5 +1,9 @@ //! The CLIP text encoder (24 layers, 1024 wide, causal) that turns a prompt into the token //! features SAM 3's detector attends to. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`, CLIP text model), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use candle_core::{Device, Module, Tensor}; use candle_nn::{Embedding, LayerNorm, VarBuilder}; diff --git a/crates/segment/src/detector.rs b/crates/segment/src/detector.rs index b21b32d8..7c4eccbd 100644 --- a/crates/segment/src/detector.rs +++ b/crates/segment/src/detector.rs @@ -1,6 +1,10 @@ //! 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. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use candle_core::{D, Device, IndexOp, Module, Tensor}; use candle_nn::{Conv2d, Conv2dConfig, GroupNorm, LayerNorm, Linear, VarBuilder}; diff --git a/crates/segment/src/fetch/http.rs b/crates/segment/src/fetch/http.rs new file mode 100644 index 00000000..332f0147 --- /dev/null +++ b/crates/segment/src/fetch/http.rs @@ -0,0 +1,541 @@ +//! A small HTTP/1.1 GET client for model downloads, in pure Rust: `std::net` sockets and, for +//! `https://`, rustls with the RustCrypto provider (`rustls-rustcrypto`; no `ring` or +//! `aws-lc-rs`, so nothing is compiled from C or assembly) and the Mozilla root certificates +//! (`webpki-roots`). +//! +//! Only what a large-file download needs: one request per connection (`Connection: close`), +//! `Range` requests for resuming, redirects (followed by the caller), `Content-Length`, +//! chunked and read-to-close bodies. Every read waits at most `stall` for data and checks the +//! cancel flag at least twice a second, so a dead server or a cancelled download never hangs +//! the downloading thread. No proxies (see docs/ai-masks.md). + +use std::io::{ErrorKind, Read, Write}; +use std::net::{TcpStream, ToSocketAddrs}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::{Duration, Instant}; + +/// Why a request failed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum HttpError { + /// The URL can't be used (scheme, host, characters). + BadUrl(String), + /// Name resolution or connecting failed (or took longer than the connect timeout). + Connect(String), + /// No data arrived for the stall timeout. + Stalled, + /// The cancel flag was raised. + Cancelled, + /// TLS failed (certificate, handshake). + Tls(String), + /// The server's response could not be understood. + Protocol(String), + /// Any other I/O error. + Io(String), +} + +impl std::fmt::Display for HttpError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + HttpError::BadUrl(e) => write!(f, "unusable URL: {e}"), + HttpError::Connect(e) => write!(f, "could not connect: {e}"), + HttpError::Stalled => write!(f, "the server stopped sending data"), + HttpError::Cancelled => write!(f, "cancelled"), + HttpError::Tls(e) => write!(f, "secure connection failed: {e}"), + HttpError::Protocol(e) => write!(f, "unexpected response: {e}"), + HttpError::Io(e) => write!(f, "{e}"), + } + } +} + +/// A parsed `http://` or `https://` URL. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Url { + pub tls: bool, + pub host: String, + pub port: u16, + /// Path and query, starting with `/`. + pub path: String, +} + +impl Url { + pub fn parse(s: &str) -> Result { + let bad = |why: &str| HttpError::BadUrl(format!("{s}: {why}")); + let s = s.trim(); + if s.chars().any(|c| c.is_control() || c == ' ') { + return Err(bad("contains spaces or control characters")); + } + let (tls, rest) = if let Some(r) = s.strip_prefix("https://") { + (true, r) + } else if let Some(r) = s.strip_prefix("http://") { + (false, r) + } else { + return Err(bad("only http:// and https:// are supported")); + }; + let rest = rest.split('#').next().unwrap_or_default(); + let (authority, path) = match rest.find(['/', '?']) { + Some(i) => (rest.get(..i).unwrap_or_default(), rest.get(i..).unwrap_or_default()), + None => (rest, ""), + }; + let path = if path.starts_with('/') { path.to_string() } else { format!("/{path}") }; + if authority.contains('@') { + return Err(bad("user names and passwords in URLs are not supported")); + } + let (host, port) = if let Some(v6) = authority.strip_prefix('[') { + let end = v6.find(']').ok_or_else(|| bad("unterminated IPv6 address"))?; + let host = v6.get(..end).unwrap_or_default(); + let after = v6.get(end + 1..).unwrap_or_default(); + (host.to_string(), after.strip_prefix(':')) + } else { + match authority.rsplit_once(':') { + Some((h, p)) => (h.to_string(), Some(p)), + None => (authority.to_string(), None), + } + }; + let port = match port { + Some(p) => p.parse::().ok().filter(|p| *p != 0).ok_or_else(|| bad("bad port"))?, + None if tls => 443, + None => 80, + }; + if host.is_empty() || !host.chars().all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '.' | ':' | '_')) { + return Err(bad("bad host name")); + } + Ok(Url { tls, host: host.to_ascii_lowercase(), port, path }) + } + + /// `Location` of a redirect, resolved against this URL. + pub fn join(&self, location: &str) -> Result { + let l = location.trim(); + if l.starts_with("http://") || l.starts_with("https://") { + return Url::parse(l); + } + // another scheme (`ftp:`, `file:`, `data:`…) + if let Some((scheme, _)) = l.split_once(':') + && !scheme.is_empty() + && scheme.chars().all(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '-' | '.')) + { + return Err(HttpError::BadUrl(format!("redirect to an unsupported location `{scheme}:`"))); + } + let scheme = if self.tls { "https" } else { "http" }; + let host = if self.host.contains(':') { format!("[{}]", self.host) } else { self.host.clone() }; + let base = format!("{scheme}://{host}:{}", self.port); + if let Some(rest) = l.strip_prefix("//") { + return Url::parse(&format!("{scheme}://{rest}")); + } + if l.starts_with('/') { + return Url::parse(&format!("{base}{l}")); + } + // relative to the current directory + let dir = self.path.split('?').next().unwrap_or("/"); + let dir = dir.rsplit_once('/').map_or("", |(d, _)| d); + Url::parse(&format!("{base}{dir}/{l}")) + } + + fn host_header(&self) -> String { + let host = if self.host.contains(':') { format!("[{}]", self.host) } else { self.host.clone() }; + let default = if self.tls { 443 } else { 80 }; + if self.port == default { host } else { format!("{host}:{}", self.port) } + } + + /// The URL without its query (for messages: a query may hold a signed token). + pub fn display(&self) -> String { + let scheme = if self.tls { "https" } else { "http" }; + format!("{scheme}://{}{}", self.host_header(), self.path.split('?').next().unwrap_or("/")) + } +} + +/// Timeouts and the cancel flag for one request. +pub struct Limits<'a> { + pub connect: Duration, + pub stall: Duration, + pub cancel: &'a AtomicBool, +} + +/// How often a blocked read wakes up to check the cancel flag. +const TICK: Duration = Duration::from_millis(250); +/// Longest status line + headers accepted. +const MAX_HEAD: usize = 64 * 1024; + +enum Stream { + Plain(TcpStream), + Tls(Box, TcpStream), +} + +fn is_timeout(e: &std::io::Error) -> bool { + matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut | ErrorKind::Interrupted) +} + +fn tls_config() -> Result, HttpError> { + let roots = rustls::RootCertStore { roots: webpki_roots::TLS_SERVER_ROOTS.to_vec() }; + let config = rustls::ClientConfig::builder_with_provider(Arc::new(rustls_rustcrypto::provider())) + .with_safe_default_protocol_versions() + .map_err(|e| HttpError::Tls(e.to_string()))? + .with_root_certificates(roots) + .with_no_client_auth(); + Ok(Arc::new(config)) +} + +impl Stream { + fn connect(url: &Url, limits: &Limits) -> Result { + let addrs: Vec<_> = (url.host.as_str(), url.port).to_socket_addrs().map_err(|e| HttpError::Connect(format!("{}: {e}", url.host)))?.collect(); + if addrs.is_empty() { + return Err(HttpError::Connect(format!("{}: no address", url.host))); + } + let mut last = String::new(); + let mut sock = None; + for a in addrs { + if limits.cancel.load(Ordering::Relaxed) { + return Err(HttpError::Cancelled); + } + match TcpStream::connect_timeout(&a, limits.connect) { + Ok(s) => { + sock = Some(s); + break; + } + Err(e) => last = format!("{}: {e}", url.host), + } + } + let sock = sock.ok_or(HttpError::Connect(last))?; + sock.set_read_timeout(Some(TICK)).map_err(|e| HttpError::Io(e.to_string()))?; + sock.set_write_timeout(Some(TICK)).map_err(|e| HttpError::Io(e.to_string()))?; + let _ = sock.set_nodelay(true); + if !url.tls { + return Ok(Stream::Plain(sock)); + } + let name = rustls::pki_types::ServerName::try_from(url.host.clone()).map_err(|e| HttpError::Tls(format!("{}: {e}", url.host)))?; + let conn = rustls::ClientConnection::new(tls_config()?, name).map_err(|e| HttpError::Tls(e.to_string()))?; + let mut s = Stream::Tls(Box::new(conn), sock); + s.handshake(limits)?; + Ok(s) + } + + /// Wait for `f` to make progress, retrying timeouts until `stall` passes without any. + fn retry(limits: &Limits, mut f: impl FnMut() -> std::io::Result) -> Result { + let start = Instant::now(); + loop { + if limits.cancel.load(Ordering::Relaxed) { + return Err(HttpError::Cancelled); + } + match f() { + Ok(v) => return Ok(v), + Err(e) if is_timeout(&e) => { + if start.elapsed() >= limits.stall { + return Err(HttpError::Stalled); + } + } + Err(e) if e.kind() == ErrorKind::InvalidData => return Err(HttpError::Tls(e.to_string())), + Err(e) => return Err(HttpError::Io(e.to_string())), + } + } + } + + fn handshake(&mut self, limits: &Limits) -> Result<(), HttpError> { + let Stream::Tls(conn, sock) = self else { return Ok(()) }; + while conn.is_handshaking() { + Self::retry(limits, || conn.complete_io(sock))?; + } + Ok(()) + } + + fn write_all(&mut self, data: &[u8], limits: &Limits) -> Result<(), HttpError> { + match self { + Stream::Plain(sock) => { + let mut rest = data; + while !rest.is_empty() { + let n = Self::retry(limits, || sock.write(rest))?; + if n == 0 { + return Err(HttpError::Io("connection closed while sending the request".into())); + } + rest = rest.get(n..).unwrap_or_default(); + } + Ok(()) + } + Stream::Tls(conn, sock) => { + // buffered by rustls, then sent as TLS records + conn.writer().write_all(data).map_err(|e| HttpError::Tls(e.to_string()))?; + while conn.wants_write() { + Self::retry(limits, || conn.write_tls(sock))?; + } + Ok(()) + } + } + } + + /// Read some bytes (0 at the end of the stream), waiting at most `stall` for them. + fn read(&mut self, buf: &mut [u8], limits: &Limits) -> Result { + match self { + Stream::Plain(sock) => Self::retry(limits, || sock.read(buf)), + Stream::Tls(conn, sock) => { + let start = Instant::now(); + loop { + if limits.cancel.load(Ordering::Relaxed) { + return Err(HttpError::Cancelled); + } + match conn.reader().read(buf) { + Ok(n) => return Ok(n), + // closed without close_notify: the caller checks the length it got + Err(e) if e.kind() == ErrorKind::UnexpectedEof => return Ok(0), + Err(e) if e.kind() == ErrorKind::WouldBlock => {} + Err(e) => return Err(HttpError::Tls(e.to_string())), + } + // no plaintext buffered: more TLS records from the socket + match conn.read_tls(sock) { + Ok(0) => return Ok(0), + Ok(_) => { + conn.process_new_packets().map_err(|e| HttpError::Tls(e.to_string()))?; + while conn.wants_write() { + Self::retry(limits, || conn.write_tls(sock))?; + } + } + Err(e) if is_timeout(&e) => { + if start.elapsed() >= limits.stall { + return Err(HttpError::Stalled); + } + } + Err(e) => return Err(HttpError::Io(e.to_string())), + } + } + } + } + } +} + +/// The connection with a read buffer (headers and chunk sizes are read line by line). +struct Conn { + stream: Stream, + buf: Vec, + pos: usize, +} + +impl Conn { + fn fill(&mut self, limits: &Limits) -> Result { + if self.pos >= self.buf.len() { + self.buf.clear(); + self.pos = 0; + } + let mut tmp = [0u8; 64 * 1024]; + let n = self.stream.read(&mut tmp, limits)?; + self.buf.extend_from_slice(tmp.get(..n).unwrap_or_default()); + Ok(n) + } + + /// One line without its CRLF; at most `max` bytes. + fn line(&mut self, max: usize, limits: &Limits) -> Result { + loop { + let avail = self.buf.get(self.pos..).unwrap_or_default(); + if let Some(i) = avail.iter().position(|b| *b == b'\n') { + let line = avail.get(..i).unwrap_or_default(); + let line = line.strip_suffix(b"\r").unwrap_or(line); + let s = String::from_utf8_lossy(line).into_owned(); + self.pos = self.pos.saturating_add(i + 1); + return Ok(s); + } + if avail.len() > max { + return Err(HttpError::Protocol("header line too long".into())); + } + // keep the partial line, read more after it + if self.pos > 0 { + self.buf.drain(..self.pos.min(self.buf.len())); + self.pos = 0; + } + let mut tmp = [0u8; 16 * 1024]; + let n = self.stream.read(&mut tmp, limits)?; + if n == 0 { + return Err(HttpError::Protocol("connection closed in the middle of the response".into())); + } + self.buf.extend_from_slice(tmp.get(..n).unwrap_or_default()); + } + } + + /// Up to `out.len()` body bytes (0 at the end of the stream). + fn read(&mut self, out: &mut [u8], limits: &Limits) -> Result { + if self.pos >= self.buf.len() && self.fill(limits)? == 0 { + return Ok(0); + } + let avail = self.buf.get(self.pos..).unwrap_or_default(); + let n = avail.len().min(out.len()); + if let (Some(dst), Some(src)) = (out.get_mut(..n), avail.get(..n)) { + dst.copy_from_slice(src); + } + self.pos = self.pos.saturating_add(n); + Ok(n) + } +} + +/// A response's status, headers and body. +pub struct Response { + pub status: u16, + headers: Vec<(String, String)>, + conn: Conn, + body: BodyMode, +} + +enum BodyMode { + Length(u64), + Chunked { left: u64, done: bool }, + Close, + Done, +} + +impl Response { + /// The first header named `name` (case-insensitive). + pub fn header(&self, name: &str) -> Option<&str> { + self.headers.iter().find(|(k, _)| k.eq_ignore_ascii_case(name)).map(|(_, v)| v.as_str()) + } + + /// The body's length when the server announced it. + pub fn content_length(&self) -> Option { + match self.body { + BodyMode::Length(n) => Some(n), + _ => None, + } + } + + /// Read body bytes into `out` (0 at the end of the body). + pub fn read(&mut self, out: &mut [u8], limits: &Limits) -> Result { + match &mut self.body { + BodyMode::Done => Ok(0), + BodyMode::Length(left) => { + if *left == 0 { + self.body = BodyMode::Done; + return Ok(0); + } + let want = usize::try_from(*left).unwrap_or(usize::MAX).min(out.len()); + let n = self.conn.read(out.get_mut(..want).unwrap_or_default(), limits)?; + if n == 0 { + return Err(HttpError::Protocol("the connection closed before the end of the file".into())); + } + *left = left.saturating_sub(n as u64); + Ok(n) + } + BodyMode::Close => { + let n = self.conn.read(out, limits)?; + if n == 0 { + self.body = BodyMode::Done; + } + Ok(n) + } + BodyMode::Chunked { left, done } => { + if *done { + return Ok(0); + } + if *left == 0 { + let line = self.conn.line(1024, limits)?; + let hex = line.split(';').next().unwrap_or_default().trim(); + let size = u64::from_str_radix(hex, 16).map_err(|_| HttpError::Protocol(format!("bad chunk size `{hex}`")))?; + if size == 0 { + // trailers, up to the empty line + let mut n = 0; + while !self.conn.line(MAX_HEAD, limits)?.is_empty() { + n += 1; + if n > 100 { + return Err(HttpError::Protocol("too many trailers".into())); + } + } + *done = true; + return Ok(0); + } + *left = size; + } + let want = usize::try_from(*left).unwrap_or(usize::MAX).min(out.len()); + let n = self.conn.read(out.get_mut(..want).unwrap_or_default(), limits)?; + if n == 0 { + return Err(HttpError::Protocol("the connection closed in the middle of a chunk".into())); + } + *left = left.saturating_sub(n as u64); + if *left == 0 && !self.conn.line(16, limits)?.is_empty() { + return Err(HttpError::Protocol("missing CRLF after a chunk".into())); + } + Ok(n) + } + } + } +} + +/// Send `GET url` with extra `headers` and read the status line and headers. +pub fn get(url: &Url, headers: &[(&str, String)], limits: &Limits) -> Result { + let mut req = format!( + "GET {} HTTP/1.1\r\nHost: {}\r\nUser-Agent: LightCraft/{}\r\nAccept: */*\r\nAccept-Encoding: identity\r\nConnection: close\r\n", + url.path, + url.host_header(), + env!("CARGO_PKG_VERSION") + ); + for (k, v) in headers { + if v.chars().any(|c| c == '\r' || c == '\n') || k.chars().any(|c| c == '\r' || c == '\n' || c == ':') { + return Err(HttpError::BadUrl("header with a line break".into())); + } + req.push_str(&format!("{k}: {v}\r\n")); + } + req.push_str("\r\n"); + let mut stream = Stream::connect(url, limits)?; + stream.write_all(req.as_bytes(), limits)?; + let mut conn = Conn { stream, buf: Vec::new(), pos: 0 }; + let status_line = conn.line(MAX_HEAD, limits)?; + let mut parts = status_line.split_whitespace(); + let version = parts.next().unwrap_or_default(); + if !version.starts_with("HTTP/1.") { + return Err(HttpError::Protocol(format!("not an HTTP/1.x response: {}", status_line.chars().take(60).collect::()))); + } + let status = parts.next().and_then(|s| s.parse::().ok()).ok_or_else(|| HttpError::Protocol("bad status line".into()))?; + let mut headers = Vec::new(); + let mut total = status_line.len(); + loop { + let line = conn.line(MAX_HEAD, limits)?; + if line.is_empty() { + break; + } + total = total.saturating_add(line.len()); + if total > MAX_HEAD || headers.len() > 200 { + return Err(HttpError::Protocol("response headers too large".into())); + } + if let Some((k, v)) = line.split_once(':') { + headers.push((k.trim().to_string(), v.trim().to_string())); + } + } + let find = |name: &str| headers.iter().find(|(k, _)| k.eq_ignore_ascii_case(name)).map(|(_, v)| v.as_str()); + let body = if status == 204 || status == 304 || (100..200).contains(&status) { + BodyMode::Done + } else if find("transfer-encoding").is_some_and(|v| v.to_ascii_lowercase().contains("chunked")) { + BodyMode::Chunked { left: 0, done: false } + } else if let Some(n) = find("content-length") { + BodyMode::Length(n.trim().parse::().map_err(|_| HttpError::Protocol(format!("bad Content-Length `{n}`")))?) + } else { + BodyMode::Close + }; + Ok(Response { status, headers, conn, body }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_urls() { + let u = Url::parse("https://cdn.example.com/models/sam3/model.safetensors?sig=1#x").unwrap(); + assert_eq!((u.tls, u.host.as_str(), u.port, u.path.as_str()), (true, "cdn.example.com", 443, "/models/sam3/model.safetensors?sig=1")); + assert_eq!(u.display(), "https://cdn.example.com/models/sam3/model.safetensors"); + let u = Url::parse("http://127.0.0.1:8080").unwrap(); + assert_eq!((u.tls, u.port, u.path.as_str()), (false, 8080, "/")); + let u = Url::parse("http://[::1]:9/a").unwrap(); + assert_eq!((u.host.as_str(), u.port), ("::1", 9)); + for bad in ["ftp://x/y", "https://", "https://a b/", "https://u:p@h/", "http://h:0/", "http://h:99999/", "https://h/\r\nX: y", "file:///etc"] + { + assert!(Url::parse(bad).is_err(), "{bad}"); + } + } + + #[test] + fn resolves_redirects() { + let u = Url::parse("https://a.example/x/y/file?q").unwrap(); + assert_eq!(u.join("https://b.example/z").unwrap().host, "b.example"); + assert_eq!(u.join("/root").unwrap().path, "/root"); + assert_eq!(u.join("other").unwrap().path, "/x/y/other"); + assert_eq!(u.join("//c.example/p").unwrap().host, "c.example"); + assert!(u.join("ftp://nope").is_err()); + } + + #[test] + fn the_tls_config_builds_with_the_pure_rust_provider() { + assert!(tls_config().is_ok()); + } +} diff --git a/crates/segment/src/fetch/mod.rs b/crates/segment/src/fetch/mod.rs new file mode 100644 index 00000000..92c5b6ea --- /dev/null +++ b/crates/segment/src/fetch/mod.rs @@ -0,0 +1,424 @@ +//! Downloading the SAM 3 model, only when the user asked for it: from an ordered list of mirrors +//! (the next one when a mirror fails), resuming partial files, with connect and stall timeouts, +//! a size cap, progress and cancellation, and checked (exact size, SHA-256 where pinned) before +//! each file is moved into the model folder. Runs on whatever thread calls [`download`] (the +//! engine uses a background thread; the UI thread never waits for the network). +//! +//! Files are written to `.part` next to their final place and renamed when verified, so a +//! model folder never holds a half-written file under a real name, and an interrupted download +//! resumes from where it stopped. A file that fails its check is deleted, never kept. +//! +//! The weights are never part of LightCraft (SAM License, see docs/ai-masks.md); where they are +//! downloaded from is configured by [`DEFAULT_MIRRORS`] and the user's own list. + +pub mod http; + +use std::fs::{File, OpenOptions}; +use std::io::{Read, Write}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; + +use sha2::{Digest, Sha256}; + +use http::{HttpError, Limits, Url}; + +/// One file of the model. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct FileSpec { + pub name: &'static str, + /// Exact size in bytes, when pinned (then the download resumes and is checked against it). + pub size: Option, + /// SHA-256 (lower-case hex), when pinned. + pub sha256: Option<&'static str>, + /// Largest size accepted when `size` is not pinned. + pub max: u64, +} + +/// The `facebook/sam3` checkpoint files [`crate::Sam3::load`] needs. +/// +/// `model.safetensors` is pinned to the official checkpoint (Hugging Face LFS SHA-256). The two +/// tokenizer files are small text files without a pinned hash yet: they are only parsed by the +/// tokenizer (never executed) and capped in size. +// TODO(maintainer): pin `vocab.json` and `merges.txt` (size + SHA-256) when the files are +// uploaded to LightCraft's CDN. +pub const SAM3_FILES: &[FileSpec] = &[ + FileSpec { name: "vocab.json", size: None, sha256: None, max: 16 << 20 }, + FileSpec { name: "merges.txt", size: None, sha256: None, max: 16 << 20 }, + FileSpec { + name: crate::WEIGHTS_FILE, + size: Some(SAM3_WEIGHTS_SIZE), + sha256: Some("6d06f0a5f84e435071fe6603e61d0b4cc7b40e0d39d487cfd4d67d8cc11cc14a"), + max: SAM3_WEIGHTS_SIZE, + }, +]; + +/// Size of the official `model.safetensors` (3.44 GB). +pub const SAM3_WEIGHTS_SIZE: u64 = 3_439_938_512; + +/// Where LightCraft downloads the model from, in order: each is a base URL, and file `f` is at +/// `/`. **Empty for now**: LightCraft has no download location of its own yet, so the +/// in-app download needs the user's list (`LIGHTCRAFT_SAM3_MIRRORS` or the mirrors file, see +/// [`mirrors`] and docs/ai-masks.md) until this is filled in. +/// +/// Hugging Face's `facebook/sam3` can't be a default: it is gated (each user must accept the +/// SAM License there and download with their own token). +// TODO(maintainer): add LightCraft's CDN locations (https, primary first), e.g. +// "https:///models/sam3/" — the files there must match `SAM3_FILES`. +pub const DEFAULT_MIRRORS: &[&str] = &[]; + +/// Environment variable with extra mirrors (base URLs separated by commas, spaces or newlines), +/// tried before the defaults. +pub const MIRRORS_ENV: &str = "LIGHTCRAFT_SAM3_MIRRORS"; + +/// Most mirrors used (a hostile list can't make a download loop forever). +const MAX_MIRRORS: usize = 16; +/// Redirects followed per request. +const MAX_REDIRECTS: usize = 8; + +/// The mirrors to try, in order: the environment variable's, then the mirrors file's (one base +/// URL per line, `#` comments), then [`DEFAULT_MIRRORS`]; duplicates and unusable URLs dropped. +pub fn mirrors(env: Option<&str>, file: Option<&Path>) -> Vec { + let mut out: Vec = Vec::new(); + let mut add = |s: &str| { + let s = s.trim().trim_end_matches('/'); + if !s.is_empty() && Url::parse(s).is_ok() && !out.iter().any(|m| m == s) && out.len() < MAX_MIRRORS { + out.push(s.to_string()); + } + }; + if let Some(env) = env { + env.split([',', ' ', '\n', '\t']).for_each(&mut add); + } + // a small text file; anything larger isn't a mirror list + if let Some(text) = file.and_then(|f| std::fs::metadata(f).ok().filter(|m| m.len() < 64 * 1024).and_then(|_| std::fs::read_to_string(f).ok())) { + text.lines().map(|l| l.split('#').next().unwrap_or_default()).for_each(&mut add); + } + DEFAULT_MIRRORS.iter().for_each(|m| add(m)); + out +} + +/// Download settings. +#[derive(Clone, Debug)] +pub struct Options { + /// Longest wait to connect to a mirror. + pub connect_timeout: Duration, + /// Longest wait for the next bytes once connected. + pub stall_timeout: Duration, + /// Tries per mirror and file (each resumes where the last one stopped). + pub attempts: u32, +} + +impl Default for Options { + fn default() -> Self { + Options { connect_timeout: Duration::from_secs(15), stall_timeout: Duration::from_secs(30), attempts: 3 } + } +} + +/// Where a download is. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct Progress { + pub file: String, + /// Bytes of all files on disk so far, and the total expected (pinned sizes + announced ones). + pub done: u64, + pub total: u64, + /// The mirror in use (without query). + pub mirror: String, +} + +/// Why a download failed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum DownloadError { + /// No mirror is configured. + NoMirrors, + Cancelled, + /// A local file problem (folder not writable, disk full…). + Disk(String), + /// Every mirror failed for `file`; the reason for each. + AllMirrorsFailed { + file: String, + errors: Vec, + }, +} + +impl std::fmt::Display for DownloadError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + DownloadError::NoMirrors => write!( + f, + "no download location is configured for the SAM 3 model in this build (set {MIRRORS_ENV}, or put the files in the model folder yourself; see docs/ai-masks.md)" + ), + DownloadError::Cancelled => write!(f, "the download was cancelled"), + DownloadError::Disk(e) => write!(f, "{e}"), + DownloadError::AllMirrorsFailed { file, errors } => write!(f, "could not download {file}: {}", errors.join("; ")), + } + } +} + +/// One mirror's failure for one file; `Retry` resumes on the same mirror. +enum Fail { + Retry(String), + NextMirror(String), + Cancelled, + Disk(String), +} + +/// Download `files` into `dir` from `mirrors` (in order; see [`mirrors`]). Files already there +/// and verified are kept. `progress` is called as bytes arrive (often: throttle in it). +pub fn download( + files: &[FileSpec], + mirrors: &[String], + dir: &Path, + opts: &Options, + cancel: &AtomicBool, + progress: &mut dyn FnMut(&Progress), +) -> Result<(), DownloadError> { + std::fs::create_dir_all(dir).map_err(|e| DownloadError::Disk(format!("can't create {}: {e}", dir.display())))?; + // what's left to fetch + let mut todo = Vec::new(); + for f in files { + if !verified(f, &dir.join(f.name))? { + todo.push(f); + } + } + if todo.is_empty() { + return Ok(()); + } + if mirrors.is_empty() { + return Err(DownloadError::NoMirrors); + } + let total: u64 = files.iter().map(|f| f.size.unwrap_or(0)).sum(); + let mut before: u64 = files.iter().filter(|f| !todo.contains(f)).filter_map(|f| f.size).sum(); + for f in todo { + let mut errors = Vec::new(); + let mut ok = false; + 'mirrors: for m in mirrors { + let url = format!("{m}/{}", f.name); + for _ in 0..opts.attempts.max(1) { + if cancel.load(Ordering::Relaxed) { + return Err(DownloadError::Cancelled); + } + let mut report = |n: u64, announced: u64| { + let mirror = Url::parse(m).map(|u| u.display()).unwrap_or_default(); + progress(&Progress { + file: f.name.to_string(), + done: before.saturating_add(n), + total: total.max(before.saturating_add(announced)), + mirror, + }) + }; + match fetch_file(f, &url, dir, opts, cancel, &mut report) { + Ok(()) => { + ok = true; + break 'mirrors; + } + Err(Fail::Cancelled) => return Err(DownloadError::Cancelled), + Err(Fail::Disk(e)) => return Err(DownloadError::Disk(e)), + Err(Fail::Retry(e)) => { + log::warn!("SAM 3 download of {url}: {e} (retrying)"); + errors.push(format!("{}: {e}", short(m))); + } + Err(Fail::NextMirror(e)) => { + log::warn!("SAM 3 download of {url}: {e}"); + errors.push(format!("{}: {e}", short(m))); + continue 'mirrors; + } + } + } + } + if !ok { + return Err(DownloadError::AllMirrorsFailed { file: f.name.to_string(), errors }); + } + before = before.saturating_add(std::fs::metadata(dir.join(f.name)).map(|m| m.len()).unwrap_or(0)); + } + Ok(()) +} + +/// A mirror for messages (no query string: it may hold a token). +fn short(m: &str) -> String { + Url::parse(m).map(|u| u.display()).unwrap_or_else(|_| "mirror".into()) +} + +/// Whether `path` is already the right file (exact size and hash when pinned). +fn verified(f: &FileSpec, path: &Path) -> Result { + let Ok(meta) = std::fs::metadata(path) else { return Ok(false) }; + if !meta.is_file() || meta.len() == 0 { + return Ok(false); + } + if f.size.is_some_and(|s| s != meta.len()) || meta.len() > f.max { + return Ok(false); + } + match f.sha256 { + Some(want) => Ok(sha256_file(path).map_err(DownloadError::Disk)? == want), + None => Ok(true), + } +} + +fn sha256_file(path: &Path) -> Result { + let mut file = File::open(path).map_err(|e| format!("{}: {e}", path.display()))?; + let mut h = Sha256::new(); + let mut buf = vec![0u8; 1 << 20]; + loop { + let n = file.read(&mut buf).map_err(|e| format!("{}: {e}", path.display()))?; + if n == 0 { + break; + } + h.update(buf.get(..n).unwrap_or_default()); + } + Ok(h.finalize().iter().map(|b| format!("{b:02x}")).collect()) +} + +fn part_path(dir: &Path, name: &str) -> PathBuf { + dir.join(format!("{name}.part")) +} + +/// Fetch one file from one mirror into `.part` (resuming it when its size is pinned), +/// check it and move it into place. +fn fetch_file(f: &FileSpec, url: &str, dir: &Path, opts: &Options, cancel: &AtomicBool, report: &mut dyn FnMut(u64, u64)) -> Result<(), Fail> { + let part = part_path(dir, f.name); + let limit = f.size.unwrap_or(f.max); + let disk = |e: std::io::Error| Fail::Disk(format!("{}: {e}", part.display())); + // resume only a pinned file (an unpinned one can't be checked after resuming) + let mut have = match std::fs::metadata(&part) { + Ok(m) if f.size.is_some() && m.len() <= limit => m.len(), + Ok(_) => { + std::fs::remove_file(&part).map_err(disk)?; + 0 + } + Err(_) => 0, + }; + let fetched = f.size != Some(have) || have == 0; + if fetched { + let limits = Limits { connect: opts.connect_timeout, stall: opts.stall_timeout, cancel }; + let mut url = Url::parse(url).map_err(|e| Fail::NextMirror(e.to_string()))?; + let origin = url.host.clone(); + let mut redirects = 0; + let mut resp = loop { + let mut headers = vec![]; + if have > 0 { + headers.push(("Range", format!("bytes={have}-"))); + } + // a token for the configured host only (never sent on to a redirect's host) + if let Ok(t) = std::env::var("LIGHTCRAFT_SAM3_TOKEN") + && !t.trim().is_empty() + && url.tls + && url.host == origin + { + headers.push(("Authorization", format!("Bearer {}", t.trim()))); + } + let r = http::get(&url, &headers, &limits).map_err(http_fail)?; + match r.status { + 301 | 302 | 303 | 307 | 308 => { + redirects += 1; + if redirects > MAX_REDIRECTS { + return Err(Fail::NextMirror("too many redirects".into())); + } + let loc = r.header("location").ok_or_else(|| Fail::NextMirror("redirect without a location".into()))?; + url = url.join(loc).map_err(|e| Fail::NextMirror(e.to_string()))?; + } + _ => break r, + } + }; + let mut file = match resp.status { + 200 => { + // the whole file (from the start, also when a range was asked for) + if resp.content_length().is_some_and(|n| n > limit || f.size.is_some_and(|s| s != n)) { + return Err(Fail::NextMirror(format!("the server's file has the wrong size ({:?} bytes)", resp.content_length()))); + } + have = 0; + File::create(&part).map_err(disk)? + } + 206 if have > 0 => { + let range = resp.header("content-range").unwrap_or_default(); + let (start, total) = parse_content_range(range).ok_or_else(|| Fail::NextMirror(format!("bad Content-Range `{range}`")))?; + if start != have || total.is_some_and(|t| Some(t) != f.size) { + return Err(Fail::NextMirror(format!("the server resumed at the wrong place ({range})"))); + } + OpenOptions::new().append(true).open(&part).map_err(disk)? + } + 416 => { + // the partial file doesn't fit what the server has: start over + std::fs::remove_file(&part).map_err(disk)?; + return Err(Fail::Retry("the partial download didn't match; starting over".into())); + } + 401 | 403 => return Err(Fail::NextMirror(format!("access denied (HTTP {})", resp.status))), + 404 | 410 => return Err(Fail::NextMirror(format!("not found (HTTP {})", resp.status))), + s if s >= 500 => return Err(Fail::Retry(format!("server error (HTTP {s})"))), + s => return Err(Fail::NextMirror(format!("HTTP {s}"))), + }; + let announced = f.size.unwrap_or_else(|| resp.content_length().unwrap_or(0)); + report(have, announced); + let mut buf = vec![0u8; 256 * 1024]; + loop { + let n = match resp.read(&mut buf, &limits) { + Ok(n) => n, + Err(e) => { + let _ = file.flush(); + if f.size.is_none() { + drop(file); + let _ = std::fs::remove_file(&part); + } + return Err(http_fail(e)); + } + }; + if n == 0 { + break; + } + if have.saturating_add(n as u64) > limit { + drop(file); + let _ = std::fs::remove_file(&part); + return Err(Fail::NextMirror(format!("the server sent more than the expected {limit} bytes"))); + } + file.write_all(buf.get(..n).unwrap_or_default()).map_err(disk)?; + have = have.saturating_add(n as u64); + report(have, announced); + } + file.sync_all().map_err(disk)?; + } + // checks: exact size, hash + let len = std::fs::metadata(&part).map_err(disk)?.len(); + if f.size.is_some_and(|s| s != len) { + return Err(Fail::Retry(format!("incomplete ({len} of {} bytes)", f.size.unwrap_or(0)))); + } + if len == 0 { + let _ = std::fs::remove_file(&part); + return Err(Fail::NextMirror("the server sent an empty file".into())); + } + if let Some(want) = f.sha256 { + let got = sha256_file(&part).map_err(Fail::Disk)?; + if got != want { + std::fs::remove_file(&part).map_err(disk)?; + let why = format!("the file is damaged or not the official one (SHA-256 {got})"); + // a damaged partial file from an earlier run isn't this mirror's fault + return Err(if fetched { Fail::NextMirror(why) } else { Fail::Retry(why) }); + } + } + std::fs::rename(&part, dir.join(f.name)).map_err(disk)?; + Ok(()) +} + +fn http_fail(e: HttpError) -> Fail { + match e { + HttpError::Cancelled => Fail::Cancelled, + HttpError::Stalled | HttpError::Io(_) | HttpError::Protocol(_) => Fail::Retry(e.to_string()), + HttpError::BadUrl(_) | HttpError::Connect(_) | HttpError::Tls(_) => Fail::NextMirror(e.to_string()), + } +} + +/// `bytes a-b/total` → (a, total) (`*` total → None). +fn parse_content_range(v: &str) -> Option<(u64, Option)> { + let v = v.trim().strip_prefix("bytes")?.trim(); + let (range, total) = v.split_once('/')?; + let (a, b) = range.split_once('-')?; + let (a, b) = (a.trim().parse::().ok()?, b.trim().parse::().ok()?); + if b < a { + return None; + } + let total = match total.trim() { + "*" => None, + t => Some(t.parse::().ok()?), + }; + Some((a, total)) +} + +#[cfg(test)] +mod tests; diff --git a/crates/segment/src/fetch/tests.rs b/crates/segment/src/fetch/tests.rs new file mode 100644 index 00000000..81e38f01 --- /dev/null +++ b/crates/segment/src/fetch/tests.rs @@ -0,0 +1,314 @@ +//! The downloader against a local HTTP server (std `TcpListener`; never the internet): mirror +//! fallback, redirects, resuming, chunked bodies, stalls, hash mismatches, size caps, cancel. + +use std::io::{BufRead, BufReader, Write}; +use std::net::{TcpListener, TcpStream}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; + +use super::*; + +/// The served file: 300 kB of a deterministic pattern. +fn payload() -> Vec { + (0..300_000u32).map(|i| (i.wrapping_mul(2_654_435_761) >> 13) as u8).collect() +} + +fn sha_hex(b: &[u8]) -> String { + Sha256::digest(b).iter().map(|b| format!("{b:02x}")).collect() +} + +fn leak(s: String) -> &'static str { + Box::leak(s.into_boxed_str()) +} + +fn pinned() -> FileSpec { + let p = payload(); + FileSpec { name: "model.bin", size: Some(p.len() as u64), sha256: Some(leak(sha_hex(&p))), max: p.len() as u64 } +} + +struct Server { + base: String, + hits: Arc, +} + +/// Serve on 127.0.0.1; the first path segment picks the behaviour. +fn server() -> Server { + let l = TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}", l.local_addr().unwrap()); + let hits = Arc::new(AtomicUsize::new(0)); + let h = hits.clone(); + std::thread::spawn(move || { + for s in l.incoming().flatten() { + let h = h.clone(); + std::thread::spawn(move || handle(s, &h)); + } + }); + Server { base, hits } +} + +fn handle(mut s: TcpStream, hits: &AtomicUsize) { + let n = hits.fetch_add(1, Ordering::SeqCst); + let mut r = BufReader::new(s.try_clone().unwrap()); + let mut first = String::new(); + if r.read_line(&mut first).is_err() { + return; + } + let mut range: Option = None; + loop { + let mut l = String::new(); + if r.read_line(&mut l).unwrap_or(0) == 0 || l.trim().is_empty() { + break; + } + if let Some(v) = l.to_ascii_lowercase().strip_prefix("range: bytes=") { + range = v.trim().trim_end_matches('-').parse().ok(); + } + } + let path = first.split_whitespace().nth(1).unwrap_or("/").to_string(); + let mode = path.trim_start_matches('/').split('/').next().unwrap_or("").to_string(); + let body = payload(); + let send = |s: &mut TcpStream, status: &str, headers: &str, data: &[u8]| { + let _ = s.write_all(format!("HTTP/1.1 {status}\r\n{headers}Connection: close\r\n\r\n").as_bytes()); + let _ = s.write_all(data); + }; + match mode.as_str() { + "missing" => send(&mut s, "404 Not Found", "Content-Length: 9\r\n", b"not found"), + "error" => send(&mut s, "500 Oops", "Content-Length: 25\r\n", b"server error"), + "redirect" => send(&mut s, "302 Found", &format!("Location: {}\r\nContent-Length: 0\r\n", path.replacen("/redirect", "/good", 1)), b""), + "wrong" => { + let mut b = body.clone(); + b[1000] ^= 0xff; + send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", b.len()), &b) + } + "toolong" => { + // no length announced, more bytes than the file can have + let mut b = body.clone(); + b.extend_from_slice(&[7u8; 5000]); + send(&mut s, "200 OK", "", &b) + } + "chunked" => { + let _ = s.write_all(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n"); + for c in body.chunks(70_001) { + let _ = s.write_all(format!("{:x};ext=1\r\n", c.len()).as_bytes()); + let _ = s.write_all(c); + let _ = s.write_all(b"\r\n"); + } + let _ = s.write_all(b"0\r\nX-Trailer: 1\r\n\r\n"); + } + "stall" => { + // half the file, then nothing + let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes()); + let _ = s.write_all(&body[..body.len() / 2]); + std::thread::sleep(Duration::from_secs(5)); + } + "slow" => { + let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes()); + for c in body.chunks(1000) { + if s.write_all(c).is_err() { + return; + } + std::thread::sleep(Duration::from_millis(20)); + } + } + "flaky" if n.is_multiple_of(2) && range.is_none() => { + // the first try breaks off after a third + let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes()); + let _ = s.write_all(&body[..body.len() / 3]); + } + "range416" if range.is_some() => send(&mut s, "416 Range Not Satisfiable", "Content-Length: 0\r\n", b""), + "norange" | "range416" => send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", body.len()), &body), + _ => match range { + Some(from) if from as usize >= body.len() => send(&mut s, "416 Range Not Satisfiable", "Content-Length: 0\r\n", b""), + Some(from) => send( + &mut s, + "206 Partial Content", + &format!("Content-Range: bytes {from}-{}/{}\r\nContent-Length: {}\r\n", body.len() - 1, body.len(), body.len() - from as usize), + &body[from as usize..], + ), + None => send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", body.len()), &body), + }, + } +} + +fn tmp(name: &str) -> PathBuf { + let d = std::env::temp_dir().join(format!("lc-sam3-fetch-{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&d); + d +} + +fn quick() -> Options { + Options { connect_timeout: Duration::from_secs(2), stall_timeout: Duration::from_millis(600), attempts: 2 } +} + +fn run(files: &[FileSpec], mirrors: &[String], dir: &Path, cancel: &AtomicBool) -> (Result<(), DownloadError>, Vec) { + let mut seen = Vec::new(); + let r = download(files, mirrors, dir, &quick(), cancel, &mut |p| seen.push(p.clone())); + (r, seen) +} + +#[test] +fn falls_back_to_the_next_mirror_and_verifies() { + let srv = server(); + let dir = tmp("fallback"); + let mirrors = vec![ + format!("{}/missing", srv.base), + format!("{}/error", srv.base), + "http://127.0.0.1:1/unreachable".to_string(), + format!("{}/redirect", srv.base), + ]; + let (r, seen) = run(&[pinned()], &mirrors, &dir, &AtomicBool::new(false)); + r.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + assert!(!part_path(&dir, "model.bin").exists()); + let last = seen.last().unwrap(); + assert_eq!((last.done, last.total), (300_000, 300_000)); + // an error page never ends up in the file: nothing left over from the failed mirrors + assert_eq!(std::fs::read_dir(&dir).unwrap().count(), 1); + // installed: nothing is fetched again + let hits = srv.hits.load(Ordering::SeqCst); + run(&[pinned()], &mirrors, &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(srv.hits.load(Ordering::SeqCst), hits); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn a_hash_mismatch_is_deleted_and_the_next_mirror_used() { + let srv = server(); + let dir = tmp("hash"); + let (r, _) = run(&[pinned()], &[format!("{}/wrong", srv.base)], &dir, &AtomicBool::new(false)); + let e = r.unwrap_err().to_string(); + assert!(e.contains("SHA-256"), "{e}"); + assert!(!dir.join("model.bin").exists() && !part_path(&dir, "model.bin").exists(), "a damaged file is never kept"); + run(&[pinned()], &[format!("{}/wrong", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn resumes_partial_files_and_restarts_bad_ones() { + let srv = server(); + let dir = tmp("resume"); + std::fs::create_dir_all(&dir).unwrap(); + // a third already there + std::fs::write(part_path(&dir, "model.bin"), &payload()[..100_000]).unwrap(); + run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + // a partial file with the wrong bytes: the hash fails, it's deleted, the next try starts over + std::fs::remove_file(dir.join("model.bin")).unwrap(); + std::fs::write(part_path(&dir, "model.bin"), vec![0u8; 100_000]).unwrap(); + run(&[pinned()], &[format!("{}/good", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + // a server ignoring Range: the whole file again, not appended + std::fs::remove_file(dir.join("model.bin")).unwrap(); + std::fs::write(part_path(&dir, "model.bin"), &payload()[..5]).unwrap(); + run(&[pinned()], &[format!("{}/norange", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + // a connection breaking off is resumed on the same mirror + std::fs::remove_file(dir.join("model.bin")).unwrap(); + run(&[pinned()], &[format!("{}/flaky", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn a_416_or_a_bad_complete_part_starts_over() { + let srv = server(); + let dir = tmp("416"); + std::fs::create_dir_all(&dir).unwrap(); + // the server refuses the range: the part is dropped and the file fetched whole (no loop) + std::fs::write(part_path(&dir, "model.bin"), &payload()[..100]).unwrap(); + run(&[pinned()], &[format!("{}/range416", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + // as long as the file but wrong: hash fails → deleted → downloaded again from the same mirror + std::fs::remove_file(dir.join("model.bin")).unwrap(); + std::fs::write(part_path(&dir, "model.bin"), vec![1u8; 300_000]).unwrap(); + run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn chunked_bodies_and_unpinned_files() { + let srv = server(); + let dir = tmp("chunked"); + let f = FileSpec { name: "vocab.json", size: None, sha256: None, max: 1 << 20 }; + run(std::slice::from_ref(&f), &[format!("{}/chunked", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("vocab.json")).unwrap(), payload()); + // more than the cap: refused, nothing kept + let small = FileSpec { max: 1000, ..f }; + let dir2 = tmp("cap"); + let e = run(&[small], &[format!("{}/good", srv.base)], &dir2, &AtomicBool::new(false)).0.unwrap_err().to_string(); + assert!(e.contains("wrong size") || e.contains("more than"), "{e}"); + assert!(!dir2.join("vocab.json").exists() && !part_path(&dir2, "vocab.json").exists()); + let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::remove_dir_all(&dir2); +} + +#[test] +fn a_server_sending_too_much_is_cut_off() { + let srv = server(); + let dir = tmp("toolong"); + let e = run(&[pinned()], &[format!("{}/toolong", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap_err().to_string(); + assert!(e.contains("more than"), "{e}"); + assert!(!dir.join("model.bin").exists() && !part_path(&dir, "model.bin").exists()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn a_stalled_server_times_out_and_the_next_mirror_finishes() { + let srv = server(); + let dir = tmp("stall"); + let t = Instant::now(); + let (r, _) = run(&[pinned()], &[format!("{}/stall", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)); + r.unwrap(); + assert!(t.elapsed() < Duration::from_secs(4), "{:?}", t.elapsed()); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn cancel_stops_quickly_and_keeps_the_part_for_later() { + let srv = server(); + let dir = tmp("cancel"); + let cancel = Arc::new(AtomicBool::new(false)); + let c = cancel.clone(); + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(400)); + c.store(true, Ordering::SeqCst); + }); + let t = Instant::now(); + let (r, _) = run(&[pinned()], &[format!("{}/slow", srv.base)], &dir, &cancel); + assert_eq!(r, Err(DownloadError::Cancelled)); + assert!(t.elapsed() < Duration::from_secs(3), "{:?}", t.elapsed()); + assert!(!dir.join("model.bin").exists()); + let kept = std::fs::metadata(part_path(&dir, "model.bin")).map(|m| m.len()).unwrap_or(0); + assert!(kept > 0 && kept < 300_000, "{kept}"); + // and the next run resumes it + run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn no_mirrors_is_a_clear_error_and_mirror_lists_are_parsed() { + let dir = tmp("nomirror"); + let e = run(&[pinned()], &[], &dir, &AtomicBool::new(false)).0.unwrap_err(); + assert_eq!(e, DownloadError::NoMirrors); + assert!(e.to_string().contains(MIRRORS_ENV)); + let file = std::env::temp_dir().join(format!("lc-sam3-mirrors-{}.txt", std::process::id())); + std::fs::write(&file, "# mine\nhttps://b.example/m/ # second\n\nnot a url\nhttps://a.example/x\n").unwrap(); + let m = mirrors(Some("https://a.example/x/, ftp://no"), Some(&file)); + assert_eq!(m[..2], ["https://a.example/x".to_string(), "https://b.example/m".to_string()]); + assert_eq!(m.len(), 2 + DEFAULT_MIRRORS.len()); + assert_eq!(mirrors(None, Some(Path::new("/nonexistent/mirrors.txt"))).len(), DEFAULT_MIRRORS.len()); + let _ = std::fs::remove_file(&file); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn content_ranges() { + assert_eq!(parse_content_range("bytes 100-299/300"), Some((100, Some(300)))); + assert_eq!(parse_content_range("bytes 0-0/*"), Some((0, None))); + assert_eq!(parse_content_range("bytes 9-1/300"), None); + assert_eq!(parse_content_range("items 1-2/3"), None); +} diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index 83bcbbaf..42aed538 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -13,7 +13,16 @@ //! 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). +//! This crate is therefore licensed under the Apache License 2.0 only (not the MIT option of +//! the rest of LightCraft). No UI dependencies (L3). +//! +//! [`fetch`] downloads the model, only when the user asks for it (from configurable mirrors, +//! verified before use). Nothing in LightCraft requires the model: without it, AI masks report +//! that it isn't installed and everything else works. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`src/transformers/models/sam3` and `sam3_tracker`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. #![forbid(unsafe_code)] #![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] @@ -21,6 +30,8 @@ mod clip; #[cfg(not(target_arch = "wasm32"))] mod detector; +#[cfg(not(target_arch = "wasm32"))] +pub mod fetch; pub mod mask; #[cfg(not(target_arch = "wasm32"))] mod neck; diff --git a/crates/segment/src/mask.rs b/crates/segment/src/mask.rs index 1e61c1f7..af349a26 100644 --- a/crates/segment/src/mask.rs +++ b/crates/segment/src/mask.rs @@ -1,4 +1,8 @@ //! Image preprocessing and mask post-processing (CPU, no model needed). +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/image_processing_sam3*.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use crate::{Error, Result}; diff --git a/crates/segment/src/neck.rs b/crates/segment/src/neck.rs index df4db88d..c4f716dd 100644 --- a/crates/segment/src/neck.rs +++ b/crates/segment/src/neck.rs @@ -1,5 +1,9 @@ //! 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. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use candle_core::{Module, Tensor}; use candle_nn::{Conv2d, Conv2dConfig, ConvTranspose2d, ConvTranspose2dConfig, VarBuilder}; diff --git a/crates/segment/src/nn.rs b/crates/segment/src/nn.rs index cec6a58b..540196e5 100644 --- a/crates/segment/src/nn.rs +++ b/crates/segment/src/nn.rs @@ -1,4 +1,8 @@ //! Shared building blocks: attention, MLPs, layer norms, sine position encodings. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`, `models/sam3_tracker`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use candle_core::{D, DType, Device, Module, Tensor}; use candle_nn::{LayerNorm, Linear, VarBuilder}; diff --git a/crates/segment/src/tokenizer.rs b/crates/segment/src/tokenizer.rs index c1a9f3c1..0819aa48 100644 --- a/crates/segment/src/tokenizer.rs +++ b/crates/segment/src/tokenizer.rs @@ -1,5 +1,9 @@ //! The CLIP byte-pair-encoding tokenizer (lower-cased, `` word ends), padded to the //! 32-token context SAM 3's text encoder takes. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (the CLIP tokenizer used by `models/sam3`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use std::collections::HashMap; use std::path::Path; diff --git a/crates/segment/src/tracker.rs b/crates/segment/src/tracker.rs index 828d56b6..d51219c8 100644 --- a/crates/segment/src/tracker.rs +++ b/crates/segment/src/tracker.rs @@ -1,5 +1,9 @@ //! 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. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3_tracker/modeling_sam3_tracker.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use candle_core::{DType, Device, IndexOp, Module, Tensor}; use candle_nn::{Conv2d, Conv2dConfig, ConvTranspose2d, ConvTranspose2dConfig, LayerNorm, VarBuilder}; diff --git a/crates/segment/src/vit.rs b/crates/segment/src/vit.rs index ad901826..16b09a72 100644 --- a/crates/segment/src/vit.rs +++ b/crates/segment/src/vit.rs @@ -1,5 +1,9 @@ //! 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. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use candle_core::{D, Device, Module, Tensor}; use candle_nn::{Conv2d, Conv2dConfig, LayerNorm, VarBuilder}; diff --git a/crates/segment/src/weights.rs b/crates/segment/src/weights.rs index 15de564a..70ac92b2 100644 --- a/crates/segment/src/weights.rs +++ b/crates/segment/src/weights.rs @@ -31,10 +31,11 @@ pub struct Weights { impl Weights { pub fn open(path: &Path) -> Result> { let mut file = File::open(path).map_err(|e| Error::Model(format!("{}: {e}", path.display())))?; + let file_len = file.metadata().map_err(|e| Error::Model(format!("{}: {e}", path.display())))?.len(); 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 { + if n > 100 << 20 || n.saturating_add(8) > file_len { return Err(Error::Model(format!("{}: not a safetensors file (header of {n} bytes)", path.display()))); } let mut header = vec![0u8; n as usize]; @@ -62,6 +63,10 @@ impl Weights { if end.checked_sub(start) != want.map(|w| w as u64) { return Err(Error::Model(format!("{name}: size does not match its shape"))); } + // the bytes must be in the file (a damaged or truncated download fails here, not later) + if (8 + n).checked_add(end).is_none_or(|e| e > file_len) { + return Err(Error::Model(format!("{}: tensor {name} lies outside the file (incomplete or damaged download?)", path.display()))); + } entries.insert(name, Entry { shape, start, end }); } Ok(Arc::new(Self { file: Mutex::new(file), entries, data_start: 8 + n })) @@ -74,7 +79,9 @@ impl Weights { /// 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]; + // checked against the file's length in `open` + let len = usize::try_from(e.end.saturating_sub(e.start)).map_err(|_| Error::Model(format!("{name}: too large")))?; + let mut bytes = vec![0u8; len]; { 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}")))?; @@ -106,3 +113,39 @@ impl SimpleBackend for Backend { self.0.contains(name) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn write(name: &str, header: &str, data_len: usize) -> std::path::PathBuf { + let p = std::env::temp_dir().join(format!("lc-weights-{name}-{}.safetensors", std::process::id())); + let mut b = (header.len() as u64).to_le_bytes().to_vec(); + b.extend_from_slice(header.as_bytes()); + b.extend(std::iter::repeat_n(0u8, data_len)); + std::fs::write(&p, b).unwrap(); + p + } + + #[test] + fn tensors_must_lie_inside_the_file() { + let ok = write("ok", r#"{"a":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]}}"#, 16); + let w = Weights::open(&ok).unwrap(); + assert_eq!(w.read("a").unwrap().dims(), &[2, 2]); + // a truncated download: the header promises bytes that aren't there + let short = write("short", r#"{"a":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]}}"#, 8); + assert!(Weights::open(&short).err().unwrap().to_string().contains("outside the file")); + // hostile offsets: huge, or beyond u64 + let huge = write("huge", r#"{"a":{"dtype":"F32","shape":[1073741824,1073741824],"data_offsets":[0,4611686018427387904]}}"#, 16); + assert!(Weights::open(&huge).is_err()); + let wrap = write("wrap", r#"{"a":{"dtype":"F32","shape":[1],"data_offsets":[18446744073709551612,18446744073709551616]}}"#, 16); + assert!(Weights::open(&wrap).is_err()); + // a header longer than the file + let p = std::env::temp_dir().join(format!("lc-weights-hdr-{}.safetensors", std::process::id())); + std::fs::write(&p, 1000u64.to_le_bytes()).unwrap(); + assert!(Weights::open(&p).is_err()); + for f in [ok, short, huge, wrap, p] { + let _ = std::fs::remove_file(f); + } + } +} From d379e780cc75fea7c31af0dcf6abab38e95ab7de Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 7 Oct 2026 09:56:50 -0400 Subject: [PATCH 3/4] =?UTF-8?q?M12:=20AI=20masks=20never=20require=20SAM?= =?UTF-8?q?=203=20=E2=80=94=20background=20worker,=20consented=20download,?= =?UTF-8?q?=20review=20fixes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Engine (crates/engine/src/segment/): - One worker thread owns the model and the encoded photo; the session sends jobs (with the render job for the model input) and gets results over channels, never locking what the model holds. In the app (Segmenter::background) clicks, descriptions, prepares and detail passes return at once and Session::segment_poll applies results (one undo step per settled selection, pending clicks shown on the photo); the CLI/MCP wait for them. Newer requests for the same selection supersede queued ones. - Every job runs under catch_unwind: a candle panic becomes an error and the model is reloaded next time; busy counters are reset by drop guards; a dead worker is restarted. The model unloads after 10 min idle. - segment.model.status / .download {acknowledged: true} / .cancel: the download runs on its own thread; status via atomics (try_lock only). Without the model every AI request fails at once with "The SAM 3 model is not installed…"; renders/exports never need it; `seg` params make AI masks from stored segmentations without it. - Hostile documents: at most 4 detail patches per component (capped on deserialize and in the renderer), at most 64 Object clicks. UI: Object / Describe without the model open "Download the SAM 3 Model?" (size, folder, SAM License link; Not Now / Download, progress, cancel, Try Again with the reason, Continue once installed); progress in the Masking panel when the dialog is closed. Reverted the PR's unrelated default change (mask overlay off + hover-red) to main's behaviour. tools/install-sam3.*: --fail so error pages never land in files, 416 restarts instead of looping, the weights are verified before they are moved into place (a bad file is deleted, never resumed from). Docs: docs/ai-masks.md (in-app download, mirrors, SAM License summary, no weights in the repo), NOTICE, parity LR-MASK-OBJECTS, What's New, ROADMAP. Co-Authored-By: Claude Opus 5.5 --- NOTICE | 12 +- ROADMAP.md | 6 +- apps/lightcraft/src/main.rs | 7 +- crates/develop/src/segmask.rs | 52 ++ crates/develop/src/settings.rs | 4 +- crates/engine/src/cmd/masks.rs | 170 +++++- crates/engine/src/lib.rs | 2 + crates/engine/src/segment.rs | 588 --------------------- crates/engine/src/segment/download.rs | 146 ++++++ crates/engine/src/segment/mod.rs | 715 ++++++++++++++++++++++++++ crates/engine/src/segment/worker.rs | 426 +++++++++++++++ crates/engine/src/tests_segment.rs | 232 +++++++++ crates/pipeline/src/masks.rs | 18 +- crates/ui-egui/src/control.rs | 4 +- crates/ui-egui/src/lib.rs | 59 ++- crates/ui-egui/src/panels/detail.rs | 20 +- crates/ui-egui/src/panels/dialogs.rs | 116 ++++- crates/ui-egui/src/panels/masking.rs | 105 ++-- crates/ui-egui/src/state.rs | 17 +- crates/ui-egui/src/tests_masking.rs | 62 ++- docs/ai-masks.md | 191 +++++-- docs/parity.md | 9 +- docs/whats-new.md | 6 + tools/install-sam3.ps1 | 24 +- tools/install-sam3.sh | 36 +- 25 files changed, 2280 insertions(+), 747 deletions(-) delete mode 100644 crates/engine/src/segment.rs create mode 100644 crates/engine/src/segment/download.rs create mode 100644 crates/engine/src/segment/mod.rs create mode 100644 crates/engine/src/segment/worker.rs create mode 100644 crates/engine/src/tests_segment.rs diff --git a/NOTICE b/NOTICE index 1fcf11e7..f3f6499f 100644 --- a/NOTICE +++ b/NOTICE @@ -15,11 +15,17 @@ 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 +crates/segment (lightcraft-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. +The port was modified by the LightCraft contributors (2026): translated from Python/PyTorch to +Rust on candle and restructured; each ported file says so. Because of this, crates/segment is +licensed under the Apache License 2.0 only (LICENSE-APACHE), not under the MIT option. + +The SAM 3 weights are not part of LightCraft and are never distributed with it. A user downloads +the facebook/sam3 checkpoint only when they choose to (in the app, or by hand), under Meta's SAM +License (https://github.com/facebookresearch/sam3/blob/main/LICENSE), a separate, non-OSI +licence. See docs/ai-masks.md. 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 diff --git a/ROADMAP.md b/ROADMAP.md index fdfbec0c..f2dcb9bc 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -39,7 +39,7 @@ duplicate Local entries #22, black GPU exports on an Intel iGPU #78). | Nikon / Sony / older-Canon raw shooter | ~65% | Camera colour fidelity and coverage (ARW and NEF preview estimates are only a starting point) | | Canon CR3 / Fujifilm / Olympus shooter | ~35% | Their raws open as embedded previews only (CR3: the full-size JPEG with full metadata) | | Lightroom Classic power user | ~45% | Print, Book, Map, publish, tethering | -| Relies on AI (masks, denoise) | ~25% | No segmentation or AI denoise models | +| Relies on AI (masks, denoise) | ~25% | Object / Describe masks via optional SAM 3 (download not yet hosted); no AI denoise | ## Where we're going @@ -57,7 +57,9 @@ Priorities, in order. Each points at tracker rows in [`docs/parity.md`](docs/par (references stay in the local `plan/`), then tune against the numbers. 5. **Lens profiles of our own** (LR-EDIT-OPTICS-PROFILE, P1). 6. **AI model strategy** (maintainer decision): which permissively licensed models (or our own training) for - segmentation masks and denoise; then pure-Rust inference. Unblocks M12 and Enhance. + segmentation masks and denoise; then pure-Rust inference. Unblocks M12 and Enhance. Object / Describe masks now + run Meta's SAM 3 in pure Rust (optional, user-consented download under the non-OSI SAM License; CDN mirrors still + to be set up, docs/ai-masks.md); Subject / Sky / People and Enhance are still open. 7. **Then:** HDR (Q), the Classic output modules (Print first, then Map view, Book, Slideshow), video (R), localisation and accessibility. diff --git a/apps/lightcraft/src/main.rs b/apps/lightcraft/src/main.rs index 7617cba1..e46fd2ee 100644 --- a/apps/lightcraft/src/main.rs +++ b/apps/lightcraft/src/main.rs @@ -406,7 +406,9 @@ 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) + LIGHTCRAFT_SAM3_DIR=DIR the SAM 3 model for Object / Describe masks (default: /models/sam3; + optional: LightCraft offers to download it when first needed) + LIGHTCRAFT_SAM3_MIRRORS=URL,… where to download the SAM 3 model from (base URLs, tried in order) "; /// Warnings and errors (failed commands, AI mask analysis) on stderr; `LIGHTCRAFT_LOG=info` @@ -504,8 +506,11 @@ fn main() -> eframe::Result { Box::new(move |cc| { 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 + // (never required: without it, AI masks offer to download it; see docs/ai-masks.md) 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"))); + // the user's own download locations, one base URL per line (LIGHTCRAFT_SAM3_MIRRORS too) + session.segmenter.mirrors_file = config_dir().map(|d| d.join("models").join("sam3-mirrors.txt")); let mut app = LightcraftApp::new(session, services()); if let Some(ui) = prefs { app.ui = ui; diff --git a/crates/develop/src/segmask.rs b/crates/develop/src/segmask.rs index 724ebdf5..172d53fc 100644 --- a/crates/develop/src/segmask.rs +++ b/crates/develop/src/segmask.rs @@ -55,6 +55,34 @@ impl SegMask { } } +/// Most zoomed-in detail patches one AI mask component keeps (the app makes one; each costs a +/// decode per render, so a hostile document can't make rendering explode). +pub const MAX_DETAIL: usize = 4; + +/// Deserialize a component's detail patches, keeping at most [`MAX_DETAIL`] (the rest are +/// skipped without being decoded). +pub(crate) fn de_detail<'de, D: serde::Deserializer<'de>>(d: D) -> Result, D::Error> { + struct Capped; + impl<'de> serde::de::Visitor<'de> for Capped { + type Value = Vec; + fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + f.write_str("a list of segmentations") + } + fn visit_seq>(self, mut seq: A) -> Result, A::Error> { + let mut out = Vec::new(); + while out.len() < MAX_DETAIL { + match seq.next_element::()? { + Some(m) => out.push(m), + None => return Ok(out), + } + } + while seq.next_element::()?.is_some() {} + Ok(out) + } + } + d.deserialize_seq(Capped) +} + const B64: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; fn base64_encode(bytes: &[u8]) -> String { @@ -118,6 +146,30 @@ mod tests { } } + #[test] + fn a_hostile_document_keeps_at_most_max_detail_patches() { + let patch = serde_json::to_value(SegMask::from_logits(4, &[1.0; 16])).unwrap(); + let shape = serde_json::json!({"kind": "object", "hint": [{"x": 0.5, "y": 0.5}], "detail": vec![patch; 10_000]}); + let s: crate::MaskShape = serde_json::from_value(shape).unwrap(); + let crate::MaskShape::Object { detail, .. } = s else { panic!("an Object") }; + assert_eq!(detail.len(), MAX_DETAIL); + // and a whole develop document with them still loads + let mut d = crate::DevelopSettings::default(); + d.masks.push(crate::Mask { + id: 1, + components: vec![crate::MaskComponent { + name: None, + op: crate::MaskOp::Add, + invert: false, + shape: crate::MaskShape::Prompt { text: "x".into(), seg: None, detail: vec![SegMask::from_logits(4, &[1.0; 16]); 50], edge: 0.0 }, + }], + ..Default::default() + }); + let back: crate::DevelopSettings = serde_json::from_str(&serde_json::to_string(&d).unwrap()).unwrap(); + let crate::MaskShape::Prompt { detail, .. } = &back.masks[0].components[0].shape else { panic!("a Prompt") }; + assert_eq!(detail.len(), MAX_DETAIL); + } + #[test] fn damaged_data_is_none() { let mut m = SegMask::from_logits(8, &[1.0; 64]); diff --git a/crates/develop/src/settings.rs b/crates/develop/src/settings.rs index 736586fe..800f229b 100644 --- a/crates/develop/src/settings.rs +++ b/crates/develop/src/settings.rs @@ -782,7 +782,7 @@ pub enum MaskShape { 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")] + #[serde(default, skip_serializing_if = "Vec::is_empty", deserialize_with = "crate::segmask::de_detail")] detail: Vec, /// Edge −100..100: below 0 harder (a steeper transition), above 0 softer (feathered). #[serde(default, skip_serializing_if = "is_zero")] @@ -794,7 +794,7 @@ pub enum MaskShape { text: String, #[serde(default, skip_serializing_if = "Option::is_none")] seg: Option, - #[serde(default, skip_serializing_if = "Vec::is_empty")] + #[serde(default, skip_serializing_if = "Vec::is_empty", deserialize_with = "crate::segmask::de_detail")] detail: Vec, /// Edge −100..100, as for `Object`. #[serde(default, skip_serializing_if = "is_zero")] diff --git a/crates/engine/src/cmd/masks.rs b/crates/engine/src/cmd/masks.rs index da96f967..1615fb06 100644 --- a/crates/engine/src/cmd/masks.rs +++ b/crates/engine/src/cmd/masks.rs @@ -5,7 +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::segment::{Click, MAX_CLICKS}; use crate::{Result, Session}; fn shape_from(kind: &str, p: &Value, c: &str) -> Result { @@ -39,8 +39,10 @@ fn shape_from(kind: &str, p: &Value, c: &str) -> Result { .unwrap_or_default(); MaskShape::ColorRange { samples, refine: f64_or(p, "refine", 50.0) } } - "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 }, + "object" => MaskShape::Object { hint: points(p, "points"), exclude: points(p, "exclude"), seg: seg_param(p), detail: vec![], edge: 0.0 }, + "prompt" => { + MaskShape::Prompt { text: str_param(p, "text").unwrap_or_default().trim().to_string(), seg: seg_param(p), detail: vec![], edge: 0.0 } + } other => { return Err(bad( c, @@ -58,19 +60,50 @@ fn points(p: &Value, key: &str) -> Vec { .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<()> { +/// A segmentation passed in (`seg`: a stored or replayed AI mask; the model isn't needed). +fn seg_param(p: &Value) -> Option { + p.get("seg").and_then(|v| serde_json::from_value(v.clone()).ok()) +} + +/// What a new AI shape still needs once it is added. +enum Later { + Nothing, + /// Background mode: these Object clicks are segmented on the worker. + Clicks(Vec, Vec), + /// Background mode: the Describe selection is computed first; the mask is added when found. + Text(String), +} + +/// Compute the segmentation of a new AI shape (Object with clicks, Prompt) — or, in background +/// mode, say what to queue. An Object without clicks stays empty until the photo is clicked +/// (the photo is analyzed meanwhile); a Prompt that matches nothing is an error. Shapes given +/// with their `seg` need no model. +fn resolve_ai(s: &mut Session, shape: &mut MaskShape, c: &str) -> Result { let id = s.active().ok_or_else(|| bad(c, "no active photo"))?; + let background = s.segmenter.background; match shape { + MaskShape::Object { seg: Some(_), .. } | MaskShape::Prompt { seg: Some(_), .. } => {} + MaskShape::Object { hint, exclude, .. } if !hint.is_empty() && background => { + s.segmenter.model_dir().map_err(ai_err)?; + return Ok(Later::Clicks(std::mem::take(hint), std::mem::take(exclude))); + } 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)?; + // the clicks come later: refuse an Object mask that could never be computed (no + // model) rather than leave an empty one; in the app, analyze the photo meanwhile + s.segmenter.model_dir().map_err(ai_err)?; + if background { + s.segment_prepare(id).map_err(ai_err)?; + } + } + MaskShape::Prompt { text, .. } if background => { + s.segmenter.model_dir().map_err(ai_err)?; + if text.trim().is_empty() { + return Err(ai_err("describe what to select")); + } + return Ok(Later::Text(text.clone())); } MaskShape::Prompt { text, seg, .. } => { let found = s.segment_text(id, text).map_err(ai_err)?; @@ -78,7 +111,17 @@ fn resolve_ai(s: &mut Session, shape: &mut MaskShape, c: &str) -> Result<()> { } _ => {} } - Ok(()) + Ok(Later::Nothing) +} + +/// Queue what a new AI component (component `comp` of mask `mask`) still needs. +fn queue_later(s: &mut Session, later: Later, mask: u32, comp: usize, mut out: Value) -> Result { + if let Later::Clicks(hint, exclude) = later { + let id = s.active().ok_or_else(|| ai_err("no active photo"))?; + s.segment_clicks_later(id, mask, comp, hint, exclude).map_err(ai_err)?; + out["pending"] = json!(true); + } + Ok(out) } /// AI mask failures are shown as they are (not as "invalid parameters"). @@ -86,7 +129,7 @@ fn ai_err(msg: impl Into) -> crate::EngineError { crate::EngineError::Other(msg.into()) } -fn clicks_of(include: &[Point], exclude: &[Point]) -> Vec { +pub(crate) 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() } @@ -115,15 +158,20 @@ pub fn specs() -> Vec { "Create New Mask", [], None, - "{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?}", + "{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; object/prompt: seg? a stored segmentation, no model needed), name?} — AI kinds need the SAM 3 model; in the app a prompt returns {pending} and the mask appears when found", has_active, |s, p| { let kind = str_param(p, "kind").unwrap_or("radial").to_string(); let mut shape = shape_from(&kind, p, "mask.add")?; - resolve_ai(s, &mut shape, "mask.add")?; + let later = resolve_ai(s, &mut shape, "mask.add")?; let name = str_param(p, "name").map(str::to_string); + if let Later::Text(text) = &later { + let id = s.active().ok_or_else(|| bad("mask.add", "no active photo"))?; + s.segment_text_later(id, None, "add", name, text).map_err(ai_err)?; + return Ok(json!({"pending": true})); + } 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| { + let out = masks_edit(s, "mask.add", "Add Mask", |masks, active| { masks.push(Mask { id: next, name: name.unwrap_or_else(|| format!("Mask {next}")), @@ -132,7 +180,8 @@ pub fn specs() -> Vec { }); *active = Some(next); Ok(()) - }) + })?; + queue_later(s, later, next, 0, out) } ), cmd!( @@ -144,17 +193,29 @@ pub fn specs() -> Vec { has_active, |s, p| { let kind = str_param(p, "kind").unwrap_or("brush").to_string(); - 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; let mid = mask_id(p, active, "mask.addComponent")?; - masks_edit(s, "mask.addComponent", "Edit Mask", |masks, _| { + let id = s.active().ok_or_else(|| bad("mask.addComponent", "no active photo"))?; + if !s.develop_of(id).is_some_and(|d| d.masks.iter().any(|m| m.id == mid)) { + return Err(bad("mask.addComponent", format!("no mask {mid}"))); + } + let mut shape = shape_from(&kind, p, "mask.addComponent")?; + let later = resolve_ai(s, &mut shape, "mask.addComponent")?; + if let Later::Text(text) = &later { + let op = p.get("op").and_then(Value::as_str).unwrap_or("add"); + s.segment_text_later(id, Some(mid), op, None, text).map_err(ai_err)?; + return Ok(json!({"pending": true})); + } + let mut comp = 0; + let out = masks_edit(s, "mask.addComponent", "Edit Mask", |masks, _| { let i = find(masks, mid, "mask.addComponent")?; masks[i].components.push(MaskComponent { name: None, op, invert: bool_or(p, "invert", false), shape }); + comp = masks[i].components.len() - 1; Ok(()) - }) + })?; + queue_later(s, later, mid, comp, out) } ), cmd!( @@ -290,11 +351,23 @@ pub fn specs() -> Vec { return Err(bad(c, "not an Object component")); }; let (mut hint, mut ex, edge) = (hint.clone(), ex.clone(), *edge); + // clicks still on their way to the model count too + if let Some(pending) = s.segmenter.pending_clicks().filter(|q| (q.photo, q.mask, q.comp) == (id, mid, k)) { + (hint, ex) = (pending.hint.clone(), pending.exclude.clone()); + } + if hint.len() + ex.len() >= MAX_CLICKS { + return Err(ai_err(format!("an Object selection takes at most {MAX_CLICKS} clicks; start a new one"))); + } if exclude { ex.push(at) } else { hint.push(at) } + if s.segmenter.background && !hint.is_empty() { + let out = json!({"include": hint.len(), "exclude": ex.len(), "pending": true}); + s.segment_clicks_later(id, mid, k, hint, ex).map_err(ai_err)?; + return Ok(out); + } 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, _| { @@ -346,6 +419,63 @@ pub fn specs() -> Vec { Ok(json!({"busy": s.segmenter.busy()})) } ), + cmd!( + query "segment.model.status", + "AI Mask Model Status", + [], + None, + "{} — whether this build has AI masks, whether the SAM 3 model is installed (and where), loaded, busy, and the download's progress → {available, installed, dir, loaded, busy, analyzing, sizeBytes, license, licenseUrl, mirrors, download: {running, done, total, file, error, finished}}", + super::always, + |s, _| { + let g = &s.segmenter; + Ok(json!({ + "available": crate::segment::Segmenter::AVAILABLE, + "installed": g.installed(), + "dir": g.dir.as_ref().map(|d| d.display().to_string()), + "loaded": g.loaded(), + "busy": g.busy(), + "analyzing": g.analyzing(), + "sizeBytes": crate::segment::MODEL_BYTES, + "license": crate::segment::LICENSE_NAME, + "licenseUrl": crate::segment::LICENSE_URL, + "mirrors": g.mirrors().len(), + "download": g.download_status(), + })) + } + ), + cmd!( + query "segment.model.download", + "Download AI Mask Model", + [], + None, + "{acknowledged: true} — download the SAM 3 model (about 3.4 GB, Meta's SAM License, not LightCraft's) in the background, from the configured mirrors; only after the user agreed to it. Watch segment.model.status; segment.model.cancel stops it (it resumes later) → {started, installed, downloading}", + super::always, + |s, p| { + let c = "segment.model.download"; + if p.get("acknowledged").and_then(Value::as_bool) != Some(true) { + return Err(bad( + c, + format!( + "the SAM 3 model is a {:.1} GB download under Meta's {} ({}): pass `acknowledged: true` once the user has agreed to download it", + crate::segment::MODEL_BYTES as f64 / 1e9, + crate::segment::LICENSE_NAME, + crate::segment::LICENSE_URL + ), + )); + } + let started = s.segmenter.start_download().map_err(ai_err)?; + Ok(json!({"started": started, "installed": s.segmenter.installed(), "downloading": s.segmenter.download_status().running})) + } + ), + cmd!( + query "segment.model.cancel", + "Cancel AI Mask Model Download", + [], + None, + "{} — stop the SAM 3 download (what has arrived is kept, and a new download resumes from it) → {cancelled}", + super::always, + |s, _| Ok(json!({"cancelled": s.segmenter.cancel_download()})) + ), cmd!( "mask.refine", "Refine Mask Edges", diff --git a/crates/engine/src/lib.rs b/crates/engine/src/lib.rs index 5edcad67..ec39222a 100644 --- a/crates/engine/src/lib.rs +++ b/crates/engine/src/lib.rs @@ -727,6 +727,8 @@ mod tests_persist; #[cfg(test)] mod tests_prefs; #[cfg(test)] +mod tests_segment; +#[cfg(test)] mod tests_settings_files; #[cfg(test)] mod tests_spots; diff --git a/crates/engine/src/segment.rs b/crates/engine/src/segment.rs deleted file mode 100644 index 41aed09e..00000000 --- a/crates/engine/src/segment.rs +++ /dev/null @@ -1,588 +0,0 @@ -//! 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/segment/download.rs b/crates/engine/src/segment/download.rs new file mode 100644 index 00000000..875d652f --- /dev/null +++ b/crates/engine/src/segment/download.rs @@ -0,0 +1,146 @@ +//! The model download as a background task the session starts, watches and cancels without +//! ever blocking: progress lives in atomics, messages behind a mutex only ever `try_lock`ed +//! by the session (the download thread holds it for a few instructions at a time). + +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; + +use lightcraft_segment::fetch::{self, Options, Progress}; + +/// What the session sees of a download. +#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize)] +pub struct Status { + pub running: bool, + pub done: u64, + pub total: u64, + /// The file being fetched. + pub file: String, + /// The last download's error (cleared when a new one starts). + pub error: Option, + /// The last download finished and the files are in place. + pub finished: bool, +} + +#[derive(Default)] +struct Text { + file: String, + error: Option, +} + +#[derive(Default)] +struct Shared { + running: AtomicBool, + cancel: AtomicBool, + finished: AtomicBool, + done: AtomicU64, + total: AtomicU64, + text: Mutex, +} + +/// Clears `running` when the download thread ends, however it ends. +struct Running(Arc); + +impl Drop for Running { + fn drop(&mut self) { + self.0.running.store(false, Ordering::SeqCst); + } +} + +#[derive(Default)] +pub struct Downloader { + shared: Arc, + /// The last text seen (when the download thread held the lock at that moment). + last: Mutex, +} + +fn lock(m: &Mutex) -> std::sync::MutexGuard<'_, Text> { + m.lock().unwrap_or_else(std::sync::PoisonError::into_inner) +} + +impl Downloader { + pub fn status(&self) -> Status { + let s = &self.shared; + if let Ok(t) = s.text.try_lock() + && let Ok(mut last) = self.last.try_lock() + { + last.file.clone_from(&t.file); + last.error.clone_from(&t.error); + } + // (only the session's thread locks `last`) + let (file, error) = self.last.try_lock().map(|l| (l.file.clone(), l.error.clone())).unwrap_or_default(); + Status { + running: s.running.load(Ordering::SeqCst), + done: s.done.load(Ordering::SeqCst), + total: s.total.load(Ordering::SeqCst), + file, + error, + finished: s.finished.load(Ordering::SeqCst), + } + } + + pub fn running(&self) -> bool { + self.shared.running.load(Ordering::SeqCst) + } + + /// Ask a running download to stop (it keeps its partial files to resume later). + pub fn cancel(&self) -> bool { + let running = self.running(); + if running { + self.shared.cancel.store(true, Ordering::SeqCst); + } + running + } + + /// Start downloading `files` from `mirrors` into `dir` on a background thread. False when + /// one is already running. + pub fn start(&self, files: &'static [fetch::FileSpec], mirrors: Vec, dir: PathBuf, opts: Options) -> Result { + let s = self.shared.clone(); + if s.running.swap(true, Ordering::SeqCst) { + return Ok(false); + } + s.cancel.store(false, Ordering::SeqCst); + s.finished.store(false, Ordering::SeqCst); + s.done.store(0, Ordering::SeqCst); + s.total.store(files.iter().filter_map(|f| f.size).sum(), Ordering::SeqCst); + // nothing else runs while `running` was false + *lock(&s.text) = Text::default(); + *lock(&self.last) = Text::default(); + let guard = Running(s.clone()); + let spawned = std::thread::Builder::new().name("sam3-download".into()).spawn(move || { + let s = guard.0.clone(); + let mut last_file = String::new(); + let mut on_progress = |p: &Progress| { + s.done.store(p.done, Ordering::SeqCst); + s.total.store(p.total, Ordering::SeqCst); + if p.file != last_file { + last_file.clone_from(&p.file); + lock(&s.text).file.clone_from(&p.file); + } + }; + let r = + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| fetch::download(files, &mirrors, &dir, &opts, &s.cancel, &mut on_progress))); + let error = match r { + Ok(Ok(())) => { + s.finished.store(true, Ordering::SeqCst); + log::info!("SAM 3 model downloaded to {}", dir.display()); + None + } + Ok(Err(e)) => Some(e.to_string()), + Err(_) => Some("the download failed unexpectedly".to_string()), + }; + if let Some(e) = &error { + log::warn!("SAM 3 download: {e}"); + } + lock(&s.text).error = error; + drop(guard); + }); + match spawned { + Ok(_) => Ok(true), + Err(e) => { + // the closure (and its guard) was dropped: `running` is false again + Err(format!("could not start the download: {e}")) + } + } + } +} diff --git a/crates/engine/src/segment/mod.rs b/crates/engine/src/segment/mod.rs new file mode 100644 index 00000000..84397042 --- /dev/null +++ b/crates/engine/src/segment/mod.rs @@ -0,0 +1,715 @@ +//! AI masks: Object (clicks) and Describe (text) selections computed with SAM 3 +//! (`lightcraft-segment`, cargo feature `sam`; the desktop app enables it). +//! +//! **Nothing requires the model.** It is not part of LightCraft (SAM License); the user +//! downloads it when they first use an AI mask and agree to (`segment.model.download`), or +//! puts the files in the model folder themselves. Without it, AI mask requests fail with a +//! clear "not installed" error; renders, exports and every other feature never touch it +//! (stored segmentations are plain data in the develop settings). +//! +//! 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. +//! Everything that touches the model (loading it, rendering and encoding the photo, clicks, +//! descriptions, detail passes) runs on one worker thread ([`worker`]); the session talks to it +//! over channels and never waits on a lock. With [`Segmenter::background`] (the desktop app) +//! commands return at once and [`Session::segment_poll`] applies the results; without it (CLI, +//! MCP, tests) they wait for the worker. + +#[cfg(feature = "sam")] +mod download; +#[cfg(feature = "sam")] +mod worker; + +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; +/// Download size of the model (`model.safetensors`; the tokenizer files add ~2 MB). +pub const MODEL_BYTES: u64 = 3_439_938_512; +/// The model's licence (not LightCraft's): shown before downloading. +pub const LICENSE_NAME: &str = "SAM License (Meta)"; +pub const LICENSE_URL: &str = "https://github.com/facebookresearch/sam3/blob/main/LICENSE"; +/// How errors about a missing model start (the UI offers the download on it). +pub const NOT_INSTALLED: &str = "The SAM 3 model is not installed"; +/// Longest a waiting command (CLI, MCP) waits for the model (loading on a CPU can take minutes). +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +const WAIT: std::time::Duration = std::time::Duration::from_secs(15 * 60); +/// Most clicks one Object selection takes (the model's prompt grows with each). +pub const MAX_CLICKS: usize = 64; + +/// 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's result is for (how [`Session::segment_poll`] applies it). +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +pub(crate) enum Tag { + /// A command waiting for the result itself. + Wait, + /// Load the model and encode the photo ahead of the first click. + Prepare, + /// Object clicks on component `comp` of mask `mask`. + Clicks { photo: PhotoId, mask: u32, comp: usize, hint: Vec, exclude: Vec }, + /// A Describe selection: a new mask (`mask: None`) or a component combined by `op`. + NewPrompt { photo: PhotoId, mask: Option, op: String, name: Option, text: String }, + /// A zoomed-in pass for a component, computed from `shape` (without its detail). + Detail { photo: PhotoId, mask: u32, comp: usize, shape: MaskShape }, +} + +/// The model input: 8-bit RGB. +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +pub(crate) struct Input { + rgb: Vec, + w: usize, + h: usize, +} + +/// What a detail pass re-runs. +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +pub(crate) enum Prompt { + Clicks(Vec), + Text(String), +} + +/// Object clicks sent to the worker and not applied yet (background mode): further clicks add +/// to these, and the loupe shows them. +#[derive(Clone, Debug, PartialEq)] +pub struct PendingClicks { + pub photo: PhotoId, + pub mask: u32, + pub comp: usize, + pub hint: Vec, + pub exclude: Vec, +} + +/// What [`Session::segment_poll`] did. +#[derive(Debug, Default)] +pub struct Polled { + /// A mask changed. + pub changed: bool, + /// Errors and notices from background work, for the user. + pub messages: Vec, + /// A mask that got a new selection (worth a detail pass). + pub refine: Option, +} + +/// The model download (`segment.model.*`). +#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize)] +pub struct DownloadStatus { + pub running: bool, + pub done: u64, + pub total: u64, + pub file: String, + pub error: Option, + pub finished: bool, +} + +/// The model (on its worker thread), the download, and AI mask requests in flight. +pub struct Segmenter { + /// Where the `facebook/sam3` files live (set by the app; `None`: no AI masks). + pub dir: Option, + /// The user's list of download mirrors (one base URL per line), if any. + pub mirrors_file: Option, + /// Run requests in the background and apply them in [`Session::segment_poll`] (the + /// desktop app); otherwise commands wait for their result (CLI, MCP, tests). + pub background: bool, + pending: Option, + messages: Vec, + #[cfg(feature = "sam")] + worker: worker::Worker, + #[cfg(feature = "sam")] + download: download::Downloader, + #[cfg(feature = "sam")] + results: (std::sync::mpsc::Sender, std::sync::mpsc::Receiver), +} + +// (derivable only without the `sam` feature: the results channel has no Default) +#[allow(clippy::derivable_impls)] +impl Default for Segmenter { + fn default() -> Self { + Segmenter { + dir: None, + mirrors_file: None, + background: false, + pending: None, + messages: Vec::new(), + #[cfg(feature = "sam")] + worker: worker::Worker::default(), + #[cfg(feature = "sam")] + download: download::Downloader::default(), + #[cfg(feature = "sam")] + results: std::sync::mpsc::channel(), + } + } +} + +impl Segmenter { + /// Whether this build can compute AI masks at all. + pub const AVAILABLE: bool = cfg!(feature = "sam"); + + /// Whether the model's files are in place. + pub fn installed(&self) -> bool { + #[cfg(feature = "sam")] + if let Some(dir) = &self.dir { + return lightcraft_segment::is_model_dir(dir); + } + false + } + + /// Whether requests are queued or running (loading, analyzing the photo, a click…). + pub fn busy(&self) -> bool { + #[cfg(feature = "sam")] + return self.worker.pending() > 0; + #[cfg(not(feature = "sam"))] + false + } + + /// Whether the slow part runs now: loading the model or analyzing a photo. + pub fn analyzing(&self) -> bool { + #[cfg(feature = "sam")] + return self.worker.encoding(); + #[cfg(not(feature = "sam"))] + false + } + + /// Whether a zoomed-in detail pass is queued or running. + pub fn detail_busy(&self) -> bool { + #[cfg(feature = "sam")] + return self.worker.detail() > 0; + #[cfg(not(feature = "sam"))] + false + } + + /// Whether the model is in memory (it is unloaded after a while without use). + pub fn loaded(&self) -> bool { + #[cfg(feature = "sam")] + return self.worker.loaded(); + #[cfg(not(feature = "sam"))] + false + } + + /// Clicks waiting for their segmentation. + pub fn pending_clicks(&self) -> Option<&PendingClicks> { + self.pending.as_ref() + } + + /// The model download's state. + pub fn download_status(&self) -> DownloadStatus { + #[cfg(feature = "sam")] + { + let s = self.download.status(); + DownloadStatus { running: s.running, done: s.done, total: s.total, file: s.file, error: s.error, finished: s.finished } + } + #[cfg(not(feature = "sam"))] + DownloadStatus::default() + } + + /// Where the model is downloaded from, in order (the user's mirrors first). + pub fn mirrors(&self) -> Vec { + #[cfg(feature = "sam")] + { + let env = std::env::var(lightcraft_segment::fetch::MIRRORS_ENV).ok(); + lightcraft_segment::fetch::mirrors(env.as_deref(), self.mirrors_file.as_deref()) + } + #[cfg(not(feature = "sam"))] + Vec::new() + } + + /// The model folder, if the model is there; otherwise an error that says what to do. + 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")?; + if self.installed() { + return Ok(dir); + } + let d = self.download_status(); + if d.running { + let pct = d.done.saturating_mul(100).checked_div(d.total).unwrap_or(0); + return Err(format!("{NOT_INSTALLED} yet: it is downloading ({pct} %).")); + } + Err(format!( + "{NOT_INSTALLED}. Download it (about {:.1} GB, {LICENSE_NAME}) when LightCraft offers it, with `segment.model.download {{\"acknowledged\": true}}`, or put model.safetensors, vocab.json and merges.txt in {}.", + MODEL_BYTES as f64 / 1e9, + dir.display() + )) + } + + /// Start downloading the model (on a background thread; see `segment.model.download`). + /// `Ok(false)` when it is installed or downloading already. + pub fn start_download(&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")?; + if self.installed() { + return Ok(false); + } + #[cfg(feature = "sam")] + { + let mirrors = self.mirrors(); + if mirrors.is_empty() { + return Err(lightcraft_segment::fetch::DownloadError::NoMirrors.to_string()); + } + self.download.start(lightcraft_segment::fetch::SAM3_FILES, mirrors, dir, lightcraft_segment::fetch::Options::default()) + } + #[cfg(not(feature = "sam"))] + { + let _ = dir; + Ok(false) + } + } + + /// Stop a running download (its partial files stay, to resume). False when none runs. + pub fn cancel_download(&self) -> bool { + #[cfg(feature = "sam")] + return self.download.cancel(); + #[cfg(not(feature = "sam"))] + false + } +} + +/// Mix `v` into a 64-bit hash. +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +fn mix(h: u64, v: u64) -> u64 { + (h ^ v).wrapping_mul(0x0100_0000_01b3).rotate_left(29) +} + +/// 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 = if aspect.is_finite() { aspect.clamp(0.1, 10.0) } else { 1.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) +} + +/// `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 +} + +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() +} + +impl Session { + /// The key and settings of photo `id` as the model sees it: its look without masks, uncropped. + #[cfg_attr(not(feature = "sam"), allow(dead_code))] + 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)) + } + + /// Clicks (in normalized photo coordinates) for the model. + #[cfg(feature = "sam")] + fn model_clicks(clicks: &[Click]) -> Vec { + clicks.iter().map(|c| lightcraft_segment::Click { x: c.at.x as f32, y: c.at.y as f32, positive: c.include }).collect() + } + + /// A worker request about photo `id` (the render job for its model input goes along; the + /// worker runs it only when its cache doesn't hold the photo). `region` renders that part + /// of the photo, from a source sharp enough for it (detail passes). + #[cfg(feature = "sam")] + fn segment_job( + &mut self, + id: PhotoId, + kind: worker::Kind, + tag: Tag, + reply: std::sync::mpsc::Sender, + region: Option<[f64; 4]>, + ) -> Result { + let dir = self.segmenter.model_dir()?; + let (key, mut d) = self.segment_key(id).ok_or("no such photo")?; + let p = self.catalog.photo(id).ok_or("no such photo")?.clone(); + let missing = match &p.source { + lightcraft_catalog::Source::File { path } if !std::path::Path::new(path).exists() => Some(path.clone()), + _ => None, + }; + let render = match region { + None => self.preview_job(id, INPUT_EDGE, INPUT_EDGE, false, &d).ok_or("no such photo")?, + Some(region) => { + 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); + job.level = crate::media::SourceLevel::for_size((INPUT_EDGE as f64 / frac) as usize); + job.source = self.media.source_ref(&p, job.level); + job + } + }; + Ok(worker::Job { dir, key: if region.is_some() { 0 } else { key }, render: Some(render), missing, kind, tag, reply }) + } + + /// Send a request and wait for its result (CLI, MCP, tests; or when not in background mode). + #[cfg(feature = "sam")] + fn segment_wait(&mut self, id: PhotoId, kind: worker::Kind, region: Option<[f64; 4]>) -> Result, String> { + let (tx, rx) = std::sync::mpsc::channel(); + let job = self.segment_job(id, kind, Tag::Wait, tx, region)?; + self.segmenter.worker.submit(job)?; + match rx.recv_timeout(WAIT) { + Ok(o) => o.result, + Err(_) => Err("the AI model took too long; try again".into()), + } + } + + /// Send a request whose result [`Session::segment_poll`] applies. + #[cfg(feature = "sam")] + fn segment_queue(&mut self, id: PhotoId, kind: worker::Kind, tag: Tag, region: Option<[f64; 4]>) -> Result<(), String> { + let reply = self.segmenter.results.0.clone(); + let job = self.segment_job(id, kind, tag, reply, region)?; + self.segmenter.worker.submit(job) + } + + /// Start preparing photo `id` for AI masks (load the model, encode the image), so the first + /// click is fast. In background mode it returns at once. + pub fn segment_prepare(&mut self, id: PhotoId) -> Result<(), String> { + self.segmenter.model_dir()?; + #[cfg(feature = "sam")] + { + if self.segmenter.background { + return self.segment_queue(id, worker::Kind::Prepare, Tag::Prepare, None); + } + self.segment_wait(id, worker::Kind::Prepare, None)?; + } + #[cfg(not(feature = "sam"))] + let _ = id; + Ok(()) + } + + /// Segment the object `clicks` pick on photo `id`, waiting for the model. + pub fn segment_clicks(&mut self, id: PhotoId, clicks: &[Click]) -> Result { + self.segmenter.model_dir()?; + if !clicks.iter().any(|c| c.include) { + return Err("click the object to include it first".into()); + } + #[cfg(feature = "sam")] + { + self.segment_wait(id, worker::Kind::Clicks(Self::model_clicks(clicks)), None)?.ok_or_else(|| "nothing was selected".to_string()) + } + #[cfg(not(feature = "sam"))] + { + let _ = id; + Err("AI masks are not available in this build".into()) + } + } + + /// Segment everything `text` describes on photo `id`, waiting for the model; `Ok(None)` + /// when nothing matches. + pub fn segment_text(&mut self, id: PhotoId, text: &str) -> Result, String> { + self.segmenter.model_dir()?; + let text = text.trim(); + if text.is_empty() { + return Err("describe what to select".into()); + } + #[cfg(feature = "sam")] + { + self.segment_wait(id, worker::Kind::Text(text.to_string()), None) + } + #[cfg(not(feature = "sam"))] + { + let _ = id; + Err("AI masks are not available in this build".into()) + } + } + + /// Background mode: queue Object clicks for component `comp` of mask `mask`; the result + /// (clicks and segmentation, one undo step) is applied by [`Session::segment_poll`]. + pub(crate) fn segment_clicks_later(&mut self, id: PhotoId, mask: u32, comp: usize, hint: Vec, exclude: Vec) -> Result<(), String> { + self.segmenter.model_dir()?; + if hint.is_empty() { + return Err("click the object to include it first".into()); + } + #[cfg(feature = "sam")] + { + let clicks = Self::model_clicks(&clicks_of(&hint, &exclude)); + let tag = Tag::Clicks { photo: id, mask, comp, hint: hint.clone(), exclude: exclude.clone() }; + self.segment_queue(id, worker::Kind::Clicks(clicks), tag, None)?; + self.segmenter.pending = Some(PendingClicks { photo: id, mask, comp, hint, exclude }); + Ok(()) + } + #[cfg(not(feature = "sam"))] + { + let _ = (id, mask, comp, exclude); + Err("AI masks are not available in this build".into()) + } + } + + /// Background mode: queue a Describe selection; the mask (or component, combined by `op` + /// into mask `mask`) is created by [`Session::segment_poll`] when something matches. + pub(crate) fn segment_text_later(&mut self, id: PhotoId, mask: Option, op: &str, name: Option, text: &str) -> Result<(), String> { + self.segmenter.model_dir()?; + let text = text.trim(); + if text.is_empty() { + return Err("describe what to select".into()); + } + #[cfg(feature = "sam")] + { + let tag = Tag::NewPrompt { photo: id, mask, op: op.to_string(), name, text: text.to_string() }; + self.segment_queue(id, worker::Kind::Text(text.to_string()), tag, None) + } + #[cfg(not(feature = "sam"))] + { + let _ = (id, mask, op, name); + Err("AI masks are not available in this build".into()) + } + } + + /// 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. In background mode [`Session::segment_poll`] applies the result; + /// otherwise it is applied before this returns. `Ok(false)` when there is nothing to + /// refine (no selection yet, or it is most of the photo) or a pass is running already. + pub fn segment_detail(&mut self, id: PhotoId, mask: u32, comp: usize) -> Result { + 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 { + // (a document with absurdly many clicks isn't run through the model) + MaskShape::Object { hint, exclude, seg: Some(seg), .. } if !hint.is_empty() && hint.len() + exclude.len() <= MAX_CLICKS => { + (seg.clone(), Prompt::Clicks(clicks_of(hint, exclude))) + } + 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 kind = worker::Kind::Detail { prompt, region }; + let tag = Tag::Detail { photo: id, mask, comp, shape: strip_detail(shape) }; + if self.segmenter.background { + if self.segmenter.detail_busy() { + return Ok(false); + } + self.segment_queue(id, kind, tag, Some(region))?; + } else { + let patch = self.segment_wait(id, kind, Some(region))?; + let outcome = worker::Outcome { tag, result: Ok(patch), superseded: false }; + let mut polled = Polled::default(); + self.apply_outcome(outcome, &mut polled); + } + Ok(true) + } + #[cfg(not(feature = "sam"))] + { + let _ = (prompt, region, mask); + Ok(false) + } + } + + /// Apply finished background requests (to masks that haven't changed since) and collect + /// their messages. Cheap when nothing is pending: call it every frame. + pub fn segment_poll(&mut self) -> Polled { + #[cfg_attr(not(feature = "sam"), allow(unused_mut))] + let mut polled = Polled { messages: std::mem::take(&mut self.segmenter.messages), ..Default::default() }; + #[cfg(feature = "sam")] + { + // (before draining: a worker that is idle now has sent everything) + let idle = !self.segmenter.busy(); + let done: Vec = self.segmenter.results.1.try_iter().collect(); + for o in done { + self.apply_outcome(o, &mut polled); + } + // clicks whose request was lost (the worker restarted) aren't pending forever + if idle { + self.segmenter.pending = None; + } + } + polled + } + + #[cfg(feature = "sam")] + fn apply_outcome(&mut self, o: worker::Outcome, polled: &mut Polled) { + if o.superseded { + return; + } + match o.tag { + Tag::Wait => {} + Tag::Prepare => { + if let Err(e) = o.result { + polled.messages.push(e); + } + } + Tag::Clicks { photo, mask, comp, hint, exclude } => { + let latest = PendingClicks { photo, mask, comp, hint, exclude }; + if self.segmenter.pending.as_ref() != Some(&latest) { + return; // newer clicks are on their way + } + self.segmenter.pending = None; + let seg = match o.result { + Ok(Some(seg)) => seg, + Ok(None) => return, + Err(e) => return polled.messages.push(e), + }; + let Some(d) = self.develop_of(photo) else { return }; + let mut d = (*d).clone(); + let Some(c) = d.masks.iter_mut().find(|m| m.id == mask).and_then(|m| m.components.get_mut(comp)) else { return }; + let MaskShape::Object { edge, .. } = c.shape else { return }; + c.shape = MaskShape::Object { hint: latest.hint, exclude: latest.exclude, seg: Some(seg), detail: vec![], edge }; + match self.set_develop(photo, d, "Object Mask") { + Ok(()) => { + polled.changed = true; + polled.refine = Some(mask); + } + Err(e) => polled.messages.push(e.to_string()), + } + } + Tag::NewPrompt { photo, mask, op, name, text } => { + let seg = match o.result { + Ok(Some(seg)) => seg, + Ok(None) => return polled.messages.push(format!("Nothing matching “{text}” was found in this photo.")), + Err(e) => return polled.messages.push(e), + }; + if self.active() != Some(photo) { + return polled.messages.push(format!("The photo changed before “{text}” was found; try again.")); + } + let Ok(seg) = serde_json::to_value(seg) else { return }; + let mut p = serde_json::json!({"kind": "prompt", "text": text, "seg": seg}); + if let Some(n) = name { + p["name"] = serde_json::json!(n); + } + let r = match mask { + None => self.execute("mask.add", &p), + Some(m) => { + if !self.develop_of(photo).is_some_and(|d| d.masks.iter().any(|x| x.id == m)) { + return polled.messages.push("The mask was removed before the selection was found.".into()); + } + self.active_mask = Some(m); + p["op"] = serde_json::json!(op); + self.execute("mask.addComponent", &p) + } + }; + match r { + Ok(_) => { + polled.changed = true; + polled.refine = self.active_mask; + } + Err(e) => polled.messages.push(e.to_string()), + } + } + Tag::Detail { photo, mask, comp, shape } => { + let patch = match o.result { + Ok(Some(p)) => p, + Ok(None) => return, + Err(e) => return log::warn!("SAM 3 detail pass: {e}"), + }; + let Some(d) = self.develop_of(photo) else { return }; + let mut d = (*d).clone(); + let Some(c) = d.masks.iter_mut().find(|m| m.id == mask).and_then(|m| m.components.get_mut(comp)) else { return }; + if strip_detail(c.shape.clone()) != shape { + return; // clicked again meanwhile: a newer pass will follow + } + match &mut c.shape { + MaskShape::Object { detail, .. } | MaskShape::Prompt { detail, .. } => *detail = vec![patch], + _ => return, + } + match self.set_develop(photo, d, "AI Mask Detail") { + Ok(()) => polled.changed = true, + Err(e) => log::warn!("AI mask detail: {e}"), + } + } + } + } +} + +#[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}"); + // a hostile aspect ratio stays sane + assert!(detail_region(&seg_box(100, 40, 45, 50, 55), f64::NAN).is_some_and(|r| r.iter().all(|v| v.is_finite()))); + } + + #[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/segment/worker.rs b/crates/engine/src/segment/worker.rs new file mode 100644 index 00000000..9ded0fbe --- /dev/null +++ b/crates/engine/src/segment/worker.rs @@ -0,0 +1,426 @@ +//! The SAM 3 worker thread: it owns the model and the last encoded photo, and runs every +//! request (load, render the model input, encode, clicks, descriptions, detail passes) in +//! turn. The session only sends jobs and receives outcomes over channels, so no thread ever +//! waits on a lock the model holds, and the UI thread never runs the model. +//! +//! Every job runs under `catch_unwind`: a panic inside candle becomes an error outcome, drops +//! the (possibly inconsistent) model and leaves the worker ready for the next job. The busy +//! counters are decremented by drop guards, so they can't stay stuck. The model is unloaded +//! after [`IDLE_UNLOAD`] without requests (it holds several GB). + +use std::collections::VecDeque; +use std::panic::AssertUnwindSafe; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::mpsc::{self, RecvTimeoutError}; +use std::time::Duration; + +use lightcraft_develop::SegMask; +use lightcraft_segment::{Encoded, Sam3}; + +use super::{Input, Prompt, Tag}; +use crate::media::RenderJob; + +/// Unload the model after this long without a request. +pub const IDLE_UNLOAD: Duration = Duration::from_secs(10 * 60); + +/// What to compute. +pub enum Kind { + /// Load the model and encode the photo (so the first click is fast). + Prepare, + Clicks(Vec), + Text(String), + /// A zoomed-in pass over `region` (normalized) of the photo, rendered by the job's `render`. + Detail { + prompt: Prompt, + region: [f64; 4], + }, + /// Panics (tests of the worker's recovery). + #[cfg(test)] + Panic, +} + +/// One request. +pub struct Job { + /// The model folder. + pub dir: PathBuf, + /// Which photo and look the request is about (the encode cache's key). + pub key: u64, + /// Renders the model input when the cache doesn't hold `key` (always for a detail pass). + pub render: Option, + /// The photo's file, when it is missing on disk (for a clearer error if rendering fails). + pub missing: Option, + pub kind: Kind, + pub tag: Tag, + pub reply: mpsc::Sender, +} + +/// A finished request: the segmentation (`None`: nothing selected / nothing to report). +pub struct Outcome { + pub tag: Tag, + pub result: Result, String>, + /// A newer request for the same selection came in; this one was skipped. + pub superseded: bool, +} + +/// Counters the session reads without locking. +#[derive(Default)] +pub struct Shared { + /// Requests other than detail passes, queued or running. + pub pending: AtomicUsize, + /// Detail passes queued or running. + pub detail: AtomicUsize, + /// The model is in memory. + pub loaded: AtomicBool, + /// Loading the model or encoding a photo right now (the slow parts). + pub encoding: AtomicBool, +} + +/// Decrements a counter when dropped (also when unwinding). +struct Done<'a>(&'a AtomicUsize); + +impl Drop for Done<'_> { + fn drop(&mut self) { + let _ = self.0.try_update(Ordering::SeqCst, Ordering::SeqCst, |n| Some(n.saturating_sub(1))); + } +} + +/// Sets a flag while alive. +struct Flag<'a>(&'a AtomicBool); + +impl<'a> Flag<'a> { + fn raise(f: &'a AtomicBool) -> Self { + f.store(true, Ordering::SeqCst); + Flag(f) + } +} + +impl Drop for Flag<'_> { + fn drop(&mut self) { + self.0.store(false, Ordering::SeqCst); + } +} + +/// The session's handle on the worker thread (started on the first request, restarted if it +/// ever ends). +#[derive(Default)] +pub struct Worker { + tx: Option>, + thread: Option>, + pub shared: Arc, +} + +impl Worker { + fn counter(&self, job: &Job) -> &AtomicUsize { + if matches!(job.kind, Kind::Detail { .. }) { &self.shared.detail } else { &self.shared.pending } + } + + /// Whether the thread has ended (it never should); its counters are then meaningless. + fn dead(&self) -> bool { + self.thread.as_ref().is_some_and(|t| t.is_finished()) + } + + pub fn pending(&self) -> usize { + if self.dead() { 0 } else { self.shared.pending.load(Ordering::SeqCst) } + } + + pub fn detail(&self) -> usize { + if self.dead() { 0 } else { self.shared.detail.load(Ordering::SeqCst) } + } + + pub fn encoding(&self) -> bool { + !self.dead() && self.shared.encoding.load(Ordering::SeqCst) + } + + pub fn loaded(&self) -> bool { + !self.dead() && self.shared.loaded.load(Ordering::SeqCst) + } + + /// Queue `job`; returns at once. + pub fn submit(&mut self, job: Job) -> Result<(), String> { + if self.dead() { + // a thread that ended took its counters' meaning with it + self.tx = None; + self.thread = None; + self.shared = Arc::new(Shared::default()); + } + self.counter(&job).fetch_add(1, Ordering::SeqCst); + let job = match &self.tx { + Some(tx) => match tx.send(job) { + Ok(()) => return Ok(()), + Err(mpsc::SendError(job)) => job, + }, + None => job, + }; + // (re)start the thread + let (tx, rx) = mpsc::channel(); + let shared = self.shared.clone(); + match std::thread::Builder::new().name("sam3".into()).spawn(move || run(&rx, &shared, IDLE_UNLOAD)) { + Ok(t) => { + self.thread = Some(t); + self.tx = Some(tx.clone()); + tx.send(job).map_err(|_| "the AI mask worker stopped".to_string()) + } + Err(e) => { + let _ = self.counter(&job).try_update(Ordering::SeqCst, Ordering::SeqCst, |n| Some(n.saturating_sub(1))); + Err(format!("could not start the AI mask worker: {e}")) + } + } + } +} + +/// The model and the last encoded photo. +#[derive(Default)] +struct State { + model: Option<(PathBuf, Sam3)>, + cache: Option<(u64, Encoded)>, +} + +/// Whether `later` makes `job` pointless (it asks for the same selection again). +fn supersedes(later: &Job, job: &Job) -> bool { + match (&later.tag, &job.tag) { + (Tag::Prepare, Tag::Prepare) => later.key == job.key, + (Tag::Clicks { photo: a, mask: am, comp: ac, .. }, Tag::Clicks { photo: b, mask: bm, comp: bc, .. }) => (a, am, ac) == (b, bm, bc), + (Tag::Detail { photo: a, mask: am, comp: ac, .. }, Tag::Detail { photo: b, mask: bm, comp: bc, .. }) => (a, am, ac) == (b, bm, bc), + _ => false, + } +} + +fn run(rx: &mpsc::Receiver, shared: &Shared, idle: Duration) { + let mut state = State::default(); + let mut queue: VecDeque = VecDeque::new(); + loop { + if queue.is_empty() { + match rx.recv_timeout(idle) { + Ok(j) => queue.push_back(j), + Err(RecvTimeoutError::Timeout) => { + if state.model.is_some() { + state = State::default(); + shared.loaded.store(false, Ordering::SeqCst); + log::info!("SAM 3 unloaded after {} minutes without use", idle.as_secs() / 60); + } + continue; + } + Err(RecvTimeoutError::Disconnected) => return, + } + } + while let Ok(j) = rx.try_recv() { + queue.push_back(j); + } + let Some(mut job) = queue.pop_front() else { continue }; + let is_detail = matches!(job.kind, Kind::Detail { .. }); + let _done = Done(if is_detail { &shared.detail } else { &shared.pending }); + if queue.iter().any(|later| supersedes(later, &job)) { + let _ = job.reply.send(Outcome { tag: job.tag, result: Ok(None), superseded: true }); + continue; + } + let result = match std::panic::catch_unwind(AssertUnwindSafe(|| run_job(&mut state, shared, &mut job))) { + Ok(r) => r, + Err(p) => { + let why = p.downcast_ref::<&str>().map(|s| s.to_string()).or_else(|| p.downcast_ref::().cloned()).unwrap_or_default(); + log::error!("SAM 3 panicked: {why}"); + // don't trust what the model kept: load it again next time + state = State::default(); + Err("the AI model failed unexpectedly (it will be reloaded on the next try)".to_string()) + } + }; + if result.is_err() && !is_detail { + state.cache = None; + } + shared.loaded.store(state.model.is_some(), Ordering::SeqCst); + let _ = job.reply.send(Outcome { tag: job.tag, result, superseded: false }); + } +} + +/// The model for `dir`, loading it when needed. +fn model<'a>(state: &'a mut State, shared: &Shared, dir: &std::path::Path) -> Result<&'a mut Sam3, String> { + if state.model.as_ref().is_none_or(|(d, _)| d != dir) { + state.model = None; + state.cache = None; + let _busy = Flag::raise(&shared.encoding); + let t = web_time::Instant::now(); + let m = Sam3::load(dir).map_err(|e| e.to_string())?; + log::info!("SAM 3 loaded in {:?}", t.elapsed()); + state.model = Some((dir.to_path_buf(), m)); + shared.loaded.store(true, Ordering::SeqCst); + } + state.model.as_mut().map(|(_, m)| m).ok_or_else(|| "the model did not load".to_string()) +} + +/// Render the model input. +fn render(job: &mut Job) -> Result { + let r = job.render.take().ok_or("the photo is not available")?.run(); + let img = match r.rendered { + Ok(r) => r.image, + Err(e) => { + return Err(match &job.missing { + Some(path) => 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." + ), + None => 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 { rgb, w: img.width, h: img.height }) +} + +/// The model and the encoded photo `job.key` (rendering and encoding it when needed). +fn encoded<'a>(state: &'a mut State, shared: &Shared, job: &mut Job) -> Result<(&'a mut Sam3, &'a mut Encoded), String> { + model(state, shared, &job.dir)?; + if state.cache.as_ref().is_none_or(|(k, _)| *k != job.key) { + state.cache = None; + let input = render(job)?; + let _busy = Flag::raise(&shared.encoding); + let m = state.model.as_mut().map(|(_, m)| m).ok_or("the model did not load")?; + let t = web_time::Instant::now(); + let mut enc = m.encode(&input.rgb, input.w, input.h).map_err(|e| e.to_string())?; + m.prepare_clicks(&mut enc).map_err(|e| e.to_string())?; + log::info!("SAM 3 encoded {}×{} in {:?}", input.w, input.h, t.elapsed()); + state.cache = Some((job.key, enc)); + } + let State { model, cache } = state; + match (model.as_mut(), cache.as_mut()) { + (Some((_, m)), Some((_, e))) => Ok((m, e)), + _ => Err("the model is not ready".into()), + } +} + +/// Every phrase of `text` ("car, road"), merged (the per-pixel maximum); `None` when nothing +/// matches. +fn text_logits(model: &mut Sam3, enc: &mut Encoded, text: &str) -> Result>, String> { + 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) +} + +fn run_job(state: &mut State, shared: &Shared, job: &mut Job) -> Result, String> { + let side = lightcraft_segment::MASK_SIDE; + match std::mem::replace(&mut job.kind, Kind::Prepare) { + Kind::Prepare => { + encoded(state, shared, job)?; + Ok(None) + } + Kind::Clicks(clicks) => { + let (m, enc) = encoded(state, shared, job)?; + let p = m.segment_clicks(enc, &clicks).map_err(|e| e.to_string())?; + Ok(Some(SegMask::from_logits(side, &p.logits))) + } + Kind::Text(text) => { + let (m, enc) = encoded(state, shared, job)?; + Ok(text_logits(m, enc, &text)?.map(|l| SegMask::from_logits(side, &l))) + } + #[cfg(test)] + #[allow(clippy::panic)] + Kind::Panic => panic!("a test panic"), + Kind::Detail { prompt, region } => { + let input = render(job)?; + let m = model(state, shared, &job.dir)?; + let t = web_time::Instant::now(); + let mut enc = m.encode(&input.rgb, input.w, input.h).map_err(|e| e.to_string())?; + let (rw, rh) = (region[2] - region[0], region[3] - region[1]); + if !(rw > 0.0 && rh > 0.0) { + return Ok(None); + } + 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(m.segment_clicks(&mut enc, &inside).map_err(|e| e.to_string())?.logits) + } + Prompt::Text(text) => text_logits(m, &mut enc, &text)?, + }; + log::info!("SAM 3 detail pass ({}×{}) in {:?}", input.w, input.h, t.elapsed()); + Ok(logits.map(|l| SegMask::from_logits_in(side, &l, region))) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn job(tag: Tag, reply: &mpsc::Sender) -> Job { + // an empty folder: loading fails with an error (no model files) + let dir = std::env::temp_dir().join(format!("lc-sam3-worker-none-{}", std::process::id())); + Job { dir, key: 1, render: None, missing: None, kind: Kind::Prepare, tag, reply: reply.clone() } + } + + #[test] + fn errors_come_back_and_the_counters_return_to_zero() { + let mut w = Worker::default(); + let (tx, rx) = mpsc::channel(); + for _ in 0..3 { + w.submit(job(Tag::Wait, &tx)).unwrap(); + } + for _ in 0..3 { + let o = rx.recv_timeout(Duration::from_secs(30)).unwrap(); + assert!(o.result.unwrap_err().contains("not found")); + } + let t = web_time::Instant::now(); + while w.pending() > 0 && t.elapsed() < Duration::from_secs(5) { + std::thread::yield_now(); + } + assert_eq!(w.pending(), 0); + assert!(!w.encoding() && !w.loaded()); + } + + #[test] + fn a_panic_in_a_job_is_an_error_and_the_worker_carries_on() { + let shared = Shared::default(); + let (tx, rx) = mpsc::channel(); + let (jtx, jrx) = mpsc::channel(); + let mut j = job(Tag::Wait, &tx); + j.kind = Kind::Panic; + jtx.send(j).unwrap(); + jtx.send(job(Tag::Wait, &tx)).unwrap(); + shared.pending.store(2, Ordering::SeqCst); + drop(jtx); + run(&jrx, &shared, Duration::from_secs(60)); + let first = rx.recv().unwrap(); + assert!(first.result.unwrap_err().contains("failed unexpectedly")); + let second = rx.recv().unwrap(); + assert!(second.result.is_err(), "the next job still runs"); + assert_eq!(shared.pending.load(Ordering::SeqCst), 0, "counters reset"); + assert!(!shared.encoding.load(Ordering::SeqCst)); + } + + #[test] + fn a_newer_request_for_the_same_selection_supersedes_a_queued_one() { + let (tx, rx) = mpsc::channel(); + let (jtx, jrx) = mpsc::channel(); + let clicks = |n: usize| Tag::Clicks { + photo: lightcraft_catalog::PhotoId(1), + mask: 1, + comp: 0, + hint: vec![lightcraft_geom::Point::new(0.5, 0.5); n], + exclude: vec![], + }; + for n in 1..=3 { + jtx.send(job(clicks(n), &tx)).unwrap(); + } + drop(jtx); + let shared = Shared::default(); + shared.pending.store(3, Ordering::SeqCst); + run(&jrx, &shared, Duration::from_secs(60)); + let out: Vec = rx.try_iter().collect(); + assert_eq!(out.len(), 3); + assert!(out[0].superseded && out[1].superseded && !out[2].superseded); + assert_eq!(shared.pending.load(Ordering::SeqCst), 0); + } +} diff --git a/crates/engine/src/tests_segment.rs b/crates/engine/src/tests_segment.rs new file mode 100644 index 00000000..59aee18c --- /dev/null +++ b/crates/engine/src/tests_segment.rs @@ -0,0 +1,232 @@ +//! AI masks without the SAM 3 model (it is never required): every request is a clear error that +//! comes back at once, nothing waits on the network or the model, a damaged model is an error +//! (not a crash or a stuck "busy"), and the download command needs consent and fails cleanly. +//! No test touches the internet: downloads go to a local server. + +use std::time::{Duration, Instant}; + +use serde_json::json; + +use crate::Session; +use crate::segment::{NOT_INSTALLED, Segmenter}; + +fn tmp(name: &str) -> std::path::PathBuf { + let d = std::env::temp_dir().join(format!("lc-sam3-engine-{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&d); + d +} + +/// A folder that looks installed but holds no usable model. +fn damaged_model(name: &str) -> std::path::PathBuf { + let d = tmp(name); + std::fs::create_dir_all(&d).unwrap(); + std::fs::write(d.join("model.safetensors"), b"\x10\x00\x00\x00\x00\x00\x00\x00{not json at all}").unwrap(); + std::fs::write(d.join("vocab.json"), b"{}").unwrap(); + std::fs::write(d.join("merges.txt"), b"#version: 0.2\n").unwrap(); + d +} + +/// Poll until the worker is idle (or `limit` passes), collecting messages. +fn settle(s: &mut Session, limit: Duration) -> Vec { + let t = Instant::now(); + let mut messages = Vec::new(); + loop { + let p = s.segment_poll(); + messages.extend(p.messages); + if !s.segmenter.busy() && !s.segmenter.detail_busy() { + messages.extend(s.segment_poll().messages); + return messages; + } + assert!(t.elapsed() < limit, "the worker never finished"); + std::thread::sleep(Duration::from_millis(10)); + } +} + +#[test] +fn status_and_model_commands_without_the_model() { + let mut s = Session::with_demo(); + let st = s.execute("segment.model.status", &json!({})).unwrap(); + assert_eq!(st["available"], Segmenter::AVAILABLE); + assert_eq!(st["installed"], false); + assert_eq!(st["download"]["running"], false); + assert!(st["license"].as_str().unwrap().contains("SAM License")); + // downloading needs the user's consent, and names the size and the licence + let e = s.execute("segment.model.download", &json!({})).unwrap_err().to_string(); + assert!(e.contains("acknowledged") && e.contains("SAM License") && e.contains("GB"), "{e}"); + let e = s.execute("segment.model.download", &json!({"acknowledged": "yes"})).unwrap_err().to_string(); + assert!(e.contains("acknowledged"), "{e}"); + assert_eq!(s.execute("segment.model.cancel", &json!({})).unwrap()["cancelled"], false); + if !Segmenter::AVAILABLE { + assert!(s.execute("segment.model.download", &json!({"acknowledged": true})).is_err()); + return; + } + // no folder, then no download location: clear errors, nothing started + let e = s.execute("segment.model.download", &json!({"acknowledged": true})).unwrap_err().to_string(); + assert!(e.contains("no folder"), "{e}"); + let dir = tmp("nomirrors"); + s.segmenter.dir = Some(dir.clone()); + s.segmenter.mirrors_file = Some(dir.join("none.txt")); + if s.segmenter.mirrors().is_empty() { + let e = s.execute("segment.model.download", &json!({"acknowledged": true})).unwrap_err().to_string(); + assert!(e.contains("LIGHTCRAFT_SAM3_MIRRORS"), "{e}"); + } + assert!(!s.segmenter.download_status().running); +} + +#[cfg(feature = "sam")] +#[test] +fn a_failing_download_ends_with_an_error_and_never_blocks() { + use std::io::{BufRead, BufReader, Write}; + // a local mirror that has nothing + let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}/sam3", l.local_addr().unwrap()); + std::thread::spawn(move || { + for mut c in l.incoming().flatten() { + let mut r = BufReader::new(c.try_clone().unwrap()); + let mut line = String::new(); + while r.read_line(&mut line).unwrap_or(0) > 2 { + line.clear(); + } + let _ = c.write_all(b"HTTP/1.1 404 Not Found\r\nContent-Length: 13\r\nConnection: close\r\n\r\n

404

\n"); + } + }); + let dir = tmp("download404"); + let mirrors = dir.with_extension("mirrors.txt"); + std::fs::write(&mirrors, format!("# test mirror\n{base}\n")).unwrap(); + let mut s = Session::with_demo(); + s.segmenter.dir = Some(dir.clone()); + s.segmenter.mirrors_file = Some(mirrors.clone()); + let t = Instant::now(); + let r = s.execute("segment.model.download", &json!({"acknowledged": true})).unwrap(); + assert_eq!(r["started"], true); + assert!(t.elapsed() < Duration::from_secs(1), "starting a download returns at once"); + // a second request while it runs starts nothing + let again = s.execute("segment.model.download", &json!({"acknowledged": true})).unwrap(); + assert!(again["started"] == false || again["downloading"] == false); + let t = Instant::now(); + while s.segmenter.download_status().running { + assert!(t.elapsed() < Duration::from_secs(30), "the download never ended"); + // AI mask requests meanwhile: an error at once, never a wait + let q = Instant::now(); + let e = s.execute("mask.add", &json!({"kind": "prompt", "text": "sky"})).unwrap_err().to_string(); + assert!(e.contains(NOT_INSTALLED), "{e}"); + assert!(q.elapsed() < Duration::from_millis(500)); + std::thread::sleep(Duration::from_millis(20)); + } + let st = s.segmenter.download_status(); + assert!(!st.finished); + let e = st.error.unwrap(); + assert!(e.contains("not found"), "{e}"); + assert!(!s.segmenter.installed()); + // nothing was left behind under a real name, and no error page was saved + for f in ["model.safetensors", "vocab.json", "merges.txt", "vocab.json.part"] { + assert!(!dir.join(f).exists(), "{f}"); + } + let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::remove_file(&mirrors); +} + +/// The app (background mode) never waits: with the model missing every request fails at once; +/// with a damaged model the requests are queued, fail on the worker, and come back as messages. +#[test] +fn background_requests_never_wait_and_failures_come_back_as_messages() { + let mut s = Session::with_demo(); + s.segmenter.background = true; + let id = s.active().unwrap(); + s.segmenter.dir = Some(tmp("missing")); + for (cmd, p) in [("mask.add", json!({"kind": "object"})), ("mask.add", json!({"kind": "prompt", "text": "sky"})), ("segment.prepare", json!({}))] + { + let t = Instant::now(); + let e = s.execute(cmd, &p).unwrap_err().to_string(); + assert!(t.elapsed() < Duration::from_millis(500), "{cmd}"); + if Segmenter::AVAILABLE { + assert!(e.starts_with(NOT_INSTALLED), "{cmd}: {e}"); + } else { + assert!(e.contains("not available"), "{cmd}: {e}"); + } + } + assert!(s.develop_of(id).unwrap().masks.is_empty()); + if !Segmenter::AVAILABLE { + return; + } + // a model folder whose weights are damaged + let dir = damaged_model("damaged"); + s.segmenter.dir = Some(dir.clone()); + assert!(s.segmenter.installed()); + let t = Instant::now(); + s.execute("mask.add", &json!({"kind": "object"})).unwrap(); + let r = s.execute("mask.objectPoint", &json!({"x": 0.5, "y": 0.5})).unwrap(); + assert_eq!(r["pending"], true); + assert!(s.segmenter.pending_clicks().is_some_and(|p| p.hint.len() == 1)); + // a second click while the first is on its way builds on it + let r = s.execute("mask.objectPoint", &json!({"x": 0.6, "y": 0.5, "exclude": true})).unwrap(); + assert_eq!((r["include"].as_u64(), r["exclude"].as_u64()), (Some(1), Some(1))); + let r = s.execute("mask.add", &json!({"kind": "prompt", "text": "the sky"})).unwrap(); + assert_eq!(r["pending"], true); + assert!(t.elapsed() < Duration::from_secs(1), "nothing waited for the model: {:?}", t.elapsed()); + let messages = settle(&mut s, Duration::from_secs(60)); + assert!(!messages.is_empty() && messages.iter().all(|m| !m.is_empty()), "{messages:?}"); + assert!(messages.iter().any(|m| m.contains("model.safetensors") || m.contains("SAM 3")), "{messages:?}"); + assert!(s.segmenter.pending_clicks().is_none()); + // the Object mask is there, without a selection; no Describe mask was made from nothing + let d = s.develop_of(id).unwrap(); + assert_eq!(d.masks.len(), 1); + assert!(matches!(&d.masks[0].components[0].shape, lightcraft_develop::MaskShape::Object { seg: None, .. })); + assert!(!s.segmenter.busy() && !s.segmenter.loaded()); + let _ = std::fs::remove_dir_all(&dir); +} + +/// CLI / MCP (waiting mode): a damaged model is an error from the command, quickly. +#[test] +fn waiting_requests_with_a_damaged_model_are_errors() { + if !Segmenter::AVAILABLE { + return; + } + let mut s = Session::with_demo(); + let id = s.active().unwrap(); + let dir = damaged_model("damaged-wait"); + s.segmenter.dir = Some(dir.clone()); + let e = s.execute("mask.add", &json!({"kind": "prompt", "text": "sky"})).unwrap_err().to_string(); + assert!(e.contains("model.safetensors") || e.contains("SAM 3"), "{e}"); + let e = s.execute("mask.add", &json!({"kind": "object", "points": [[0.5, 0.5]]})).unwrap_err().to_string(); + assert!(!e.is_empty()); + assert!(s.develop_of(id).unwrap().masks.is_empty()); + assert!(!s.segmenter.busy()); + // a stored segmentation needs no model at all (replayed or pasted masks) + let seg = serde_json::to_value(lightcraft_develop::SegMask::from_logits(4, &[5.0; 16])).unwrap(); + s.segmenter.dir = None; + s.execute("mask.add", &json!({"kind": "prompt", "text": "sky", "seg": seg})).unwrap(); + assert_eq!(s.develop_of(id).unwrap().masks.len(), 1); + let _ = std::fs::remove_dir_all(&dir); +} + +/// Too many clicks on one Object selection are refused, not sent to the model. +#[test] +fn object_clicks_are_capped() { + if !Segmenter::AVAILABLE { + return; + } + let mut s = Session::with_demo(); + s.segmenter.background = true; + let dir = damaged_model("clicks"); + s.segmenter.dir = Some(dir.clone()); + let id = s.active().unwrap(); + let mut d = (*s.develop_of(id).unwrap()).clone(); + let many = vec![lightcraft_geom::Point::new(0.5, 0.5); crate::segment::MAX_CLICKS]; + d.masks.push(lightcraft_develop::Mask { + id: 1, + components: vec![lightcraft_develop::MaskComponent { + name: None, + op: lightcraft_develop::MaskOp::Add, + invert: false, + shape: lightcraft_develop::MaskShape::Object { hint: many, exclude: vec![], seg: None, detail: vec![], edge: 0.0 }, + }], + ..Default::default() + }); + s.set_develop(id, d, "test").unwrap(); + s.active_mask = Some(1); + let e = s.execute("mask.objectPoint", &json!({"x": 0.5, "y": 0.5})).unwrap_err().to_string(); + assert!(e.contains("at most"), "{e}"); + assert!(!s.segmenter.busy()); + let _ = std::fs::remove_dir_all(&dir); +} diff --git a/crates/pipeline/src/masks.rs b/crates/pipeline/src/masks.rs index 01985b5b..a71667bc 100644 --- a/crates/pipeline/src/masks.rs +++ b/crates/pipeline/src/masks.rs @@ -228,7 +228,8 @@ impl SegGrid { /// 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(); + // (capped: a hostile document can hold many) + let patches: Vec = detail.iter().take(lightcraft_develop::segmask::MAX_DETAIL).filter_map(SegGrid::new).collect(); if coarse.is_none() && patches.is_empty() { return; } @@ -466,6 +467,21 @@ mod tests { assert!(a.get(190, 50) < 0.01, "inside the patch, its unselected half"); } + #[test] + fn a_hostile_number_of_detail_patches_is_capped() { + // built in memory (a deserialized document is capped already): only the first + // MAX_DETAIL are decoded and sampled, so this needs 4 × 4 MB, not 20 000 × 4 MB + let side = 1024; + let l: Vec = vec![12.0; side * side]; + let patch = SegMask::from_logits_in(side, &l, [0.0, 0.0, 0.5, 1.0]); + let shape = MaskShape::Object { hint: vec![Point::new(0.2, 0.5)], exclude: vec![], seg: None, detail: vec![patch; 20_000], edge: 0.0 }; + let (w, h) = (64, 32); + let t = std::time::Instant::now(); + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + assert!(a.get(5, 16) > 0.99 && a.get(60, 16) < 0.01); + assert!(t.elapsed() < std::time::Duration::from_secs(20), "{:?}", t.elapsed()); + } + #[test] fn the_edge_setting_hardens_or_feathers_the_transition() { let side = 16; diff --git a/crates/ui-egui/src/control.rs b/crates/ui-egui/src/control.rs index b592557a..1b9d22bb 100644 --- a/crates/ui-egui/src/control.rs +++ b/crates/ui-egui/src/control.rs @@ -278,7 +278,9 @@ pub fn handle(app: &mut LightcraftApp, ctx: &egui::Context, req: &ControlRequest Some(d) => { let r = crate::panels::dialogs::confirm_dialog(app, &d); // the import review stays open on an error, as with its button - if r.is_err() && matches!(d, crate::state::Dialog::Import { .. }) { + if (r.is_err() && matches!(d, crate::state::Dialog::Import { .. } | crate::state::Dialog::SamModel { .. })) + || (r.is_ok() && crate::panels::dialogs::keeps_open(app, &d)) + { app.ui.dialog = Some(d); } wrap(r) diff --git a/crates/ui-egui/src/lib.rs b/crates/ui-egui/src/lib.rs index bdc2780f..49ea140f 100644 --- a/crates/ui-egui/src/lib.rs +++ b/crates/ui-egui/src/lib.rs @@ -203,7 +203,9 @@ pub struct LightcraftApp { } impl LightcraftApp { - pub fn new(session: Session, services: Services) -> Self { + pub fn new(mut session: Session, services: Services) -> Self { + // AI mask requests run on the model's worker; frames apply their results (never wait) + session.segmenter.background = true; Self { session, ui: UiState::default(), @@ -358,13 +360,42 @@ impl LightcraftApp { self.ui.toast = Some((text.into(), t + secs)); } - /// AI masks: start a zoomed-in detail pass once the clicking stops (`ui.detail_due`), and - /// apply finished passes. + /// AI masks: apply finished background requests (clicks, descriptions, detail passes) and + /// show their errors; start a zoomed-in detail pass once the clicking stops + /// (`ui.detail_due`); watch the model download. Never waits for the model. fn ai_mask_detail(&mut self, ctx: &egui::Context) { - if self.session.segment_poll() { + let polled = self.session.segment_poll(); + let now = ctx.input(|i| i.time); + if polled.changed { ctx.request_repaint(); } - let now = ctx.input(|i| i.time); + if let Some(mask) = polled.refine { + self.ui.detail_due = Some((now + 0.3, mask)); + } + if let Some(e) = polled.messages.into_iter().last() { + self.ai_error(ctx, e, None); + } + let seg = &self.session.segmenter; + if seg.busy() || seg.pending_clicks().is_some() { + ctx.request_repaint_after(std::time::Duration::from_millis(50)); + } + let download = seg.download_status(); + if download.running { + ctx.request_repaint_after(std::time::Duration::from_millis(250)); + } + // a download finishing while its dialog is closed: say so once + if self.ui.sam_downloading && !download.running { + self.ui.sam_downloading = false; + let open = matches!(self.ui.dialog, Some(state::Dialog::SamModel { .. })); + match (&download.error, open) { + (_, true) => {} + (Some(e), false) => self.toast_error(ctx, format!("The SAM 3 download failed: {e}")), + (None, false) if download.finished => { + self.toast_error(ctx, crate::i18n::tr("The SAM 3 model is installed: Object and Describe masks are ready.")) + } + (None, false) => {} + } + } 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; @@ -380,6 +411,24 @@ impl LightcraftApp { } } + /// An AI mask error: when the model isn't installed, the dialog that offers to download it + /// (`then`: the AI mask to start afterwards); otherwise a toast. + pub fn ai_error(&mut self, ctx: &egui::Context, e: impl Into, then: Option<(&str, &str)>) { + let e = e.into(); + if lightcraft_engine::segment::Segmenter::AVAILABLE && e.starts_with(lightcraft_engine::segment::NOT_INSTALLED) { + self.offer_sam_download(then); + } else { + self.toast_error(ctx, e); + } + } + + /// Open the dialog offering the SAM 3 download (never downloads by itself). + pub fn offer_sam_download(&mut self, then: Option<(&str, &str)>) { + if self.ui.dialog.is_none() || matches!(self.ui.dialog, Some(state::Dialog::SamModel { .. })) { + self.ui.dialog = Some(state::Dialog::SamModel { then: then.map(|(k, o)| (k.to_string(), o.to_string())), error: None }); + } + } + /// 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); diff --git a/crates/ui-egui/src/panels/detail.rs b/crates/ui-egui/src/panels/detail.rs index 67c01f47..3b3622be 100644 --- a/crates/ui-egui/src/panels/detail.rs +++ b/crates/ui-egui/src/panels/detail.rs @@ -698,9 +698,6 @@ 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 { @@ -708,13 +705,6 @@ 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)) @@ -1183,9 +1173,15 @@ 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 + // (clicks still on their way to the model included) + let pending = app.session.segmenter.pending_clicks().cloned(); for m in d.masks.iter().filter(|m| Some(m.id) == active) { - for c in &m.components { + for (k, c) in m.components.iter().enumerate() { if let MaskShape::Object { hint, exclude, .. } = &c.shape { + let (hint, exclude) = match pending.as_ref().filter(|q| (q.mask, q.comp) == (m.id, k)) { + Some(q) => (&q.hint, &q.exclude), + None => (hint, exclude), + }; 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); @@ -1212,7 +1208,7 @@ fn mask_overlay(app: &mut LightcraftApp, ui: &mut egui::Ui, resp: &egui::Respons app.ui.detail_due = Some((ui.input(|i| i.time) + 1.0, mid)); } } - Err(e) => app.toast_error(ui.ctx(), e), + Err(e) => app.ai_error(ui.ctx(), e, Some(("object", "new"))), } } return; diff --git a/crates/ui-egui/src/panels/dialogs.rs b/crates/ui-egui/src/panels/dialogs.rs index a9ecdcfe..67d495de 100644 --- a/crates/ui-egui/src/panels/dialogs.rs +++ b/crates/ui-egui/src/panels/dialogs.rs @@ -94,6 +94,7 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) { Dialog::Merge { opts } => opts.title(), Dialog::Settings { .. } => "Settings", Dialog::ConfirmDelete { .. } => "Delete Photos", + Dialog::SamModel { .. } => "Download the SAM 3 Model?", Dialog::About => "About LightCraft", Dialog::Shortcuts => "Keyboard Shortcuts", } @@ -700,6 +701,7 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) { Dialog::Merge { opts } => crate::merge::body(app, ui, opts), Dialog::Import { opts } => crate::import::body(app, ui, opts), Dialog::Settings { tab } => crate::panels::settings::body(app, ui, tab), + Dialog::SamModel { error, .. } => sam_model_body(app, ui, error.as_deref()), Dialog::ConfirmDelete { count } => { let what = if *count == 1 { crate::i18n::tr("this photo").to_string() } else { crate::i18n::tr_format!("these {count} photos", count = count) }; ui.label(crate::i18n::tr_format!("Move {what} to Recently Deleted?", what = what)); @@ -767,8 +769,19 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) { ui.add_space(4.0); ui.horizontal(|ui| { let informational = matches!(dlg, Dialog::About | Dialog::Shortcuts | Dialog::Settings { .. }); - if !informational && ui.button(crate::i18n::tr("Cancel")).clicked() { - close = true; + let sam = &app.session.segmenter; + let (sam_installed, sam_running, sam_failed) = (sam.installed(), sam.download_status().running, sam.download_status().error.is_some()); + let cancel = match &dlg { + Dialog::SamModel { .. } if sam_running || sam_installed => "Close", + Dialog::SamModel { .. } => "Not Now", + _ => "Cancel", + }; + if !informational { + let r = ui.button(crate::i18n::tr(cancel)); + crate::widgets::register(ui.ctx(), "button:dialogCancel", r.rect); + if r.clicked() { + close = true; + } } let add_label; let ok = match &dlg { @@ -780,10 +793,18 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) { } Dialog::Merge { .. } => "Merge", Dialog::ConfirmDelete { .. } => "Delete", + Dialog::SamModel { then: Some(_), .. } if sam_installed => "Continue", + Dialog::SamModel { .. } if sam_installed || sam_running => "", + Dialog::SamModel { error, .. } if error.is_some() || sam_failed => "Try Again", + Dialog::SamModel { .. } => "Download", _ if informational => "Close", _ => "OK", }; - if ui.button(crate::i18n::tr(ok)).clicked() { + let r = (!ok.is_empty()).then(|| ui.button(crate::i18n::tr(ok))); + if let Some(r) = &r { + crate::widgets::register(ui.ctx(), "button:dialogOk", r.rect); + } + if r.is_some_and(|r| r.clicked()) { if informational { close = true; } else { @@ -803,6 +824,17 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) { match confirm_dialog(app, &dlg) { // the import review stays open on an error (e.g. an unusable folder template) Err(e) if matches!(dlg, Dialog::Import { .. }) => app.toast(ctx, e), + // the SAM 3 dialog stays open to show the download (or why it can't start) + Err(e) if matches!(dlg, Dialog::SamModel { .. }) => { + if let Dialog::SamModel { error, .. } = &mut dlg { + *error = Some(e); + } + } + Ok(_) if keeps_open(app, &dlg) => { + if let Dialog::SamModel { error, .. } = &mut dlg { + *error = None; + } + } _ => close = true, } } @@ -823,8 +855,86 @@ pub fn fmt_gap(v: f64) -> String { } } +/// Whether a dialog stays open after its action succeeded (the SAM 3 dialog while the model +/// downloads). +pub fn keeps_open(app: &LightcraftApp, dlg: &Dialog) -> bool { + matches!(dlg, Dialog::SamModel { .. }) && !app.session.segmenter.installed() +} + +/// The SAM 3 dialog: what the model is, its size and licence, and the download's progress. +fn sam_model_body(app: &mut LightcraftApp, ui: &mut egui::Ui, error: Option<&str>) { + use lightcraft_engine::segment::{LICENSE_NAME, LICENSE_URL, MODEL_BYTES}; + let t = Tokens::get(ui.ctx()); + let seg = &app.session.segmenter; + let d = seg.download_status(); + if seg.installed() { + ui.label(crate::i18n::tr("The SAM 3 model is installed: Object and Describe masks are ready.")); + return; + } + let gb = |b: u64| b as f64 / 1e9; + ui.label(crate::i18n::tr( + "Object and Describe masks use SAM 3, Meta's segmentation model. It isn't part of LightCraft, and everything else works without it.", + )); + let dir = seg.dir.as_ref().map(|d| d.display().to_string()).unwrap_or_default(); + ui.label(format!("{} {:.1} GB · {dir}", crate::i18n::tr("Download it once (about):"), gb(MODEL_BYTES))); + ui.label( + egui::RichText::new(format!( + "{} {LICENSE_NAME} — {}", + crate::i18n::tr("Licence:"), + crate::i18n::tr("Meta's terms, not LightCraft's. Downloading it means accepting them.") + )) + .color(t.text_label), + ); + let r = ui.link(crate::i18n::tr("Read the SAM License")).on_hover_text(LICENSE_URL); + crate::widgets::register(ui.ctx(), "link:samLicense", r.rect); + if r.clicked() { + let _ = crate::links::open(app, LICENSE_URL); + } + let seg = &app.session.segmenter; + if d.running { + ui.add_space(4.0); + let frac = if d.total > 0 { d.done as f64 / d.total as f64 } else { 0.0 }; + let text = format!("{:.2} / {:.2} GB · {}", gb(d.done), gb(d.total), d.file); + let r = ui.add(egui::ProgressBar::new(frac as f32).text(text)); + crate::widgets::register(ui.ctx(), "progress:samDownload", r.rect); + let r = ui.button(crate::i18n::tr("Cancel Download")); + crate::widgets::register(ui.ctx(), "button:samCancel", r.rect); + if r.clicked() { + app.session.segmenter.cancel_download(); + } + ui.label(egui::RichText::new(crate::i18n::tr("You can close this: the download continues, and resumes if interrupted.")).color(t.text_dim)); + ui.ctx().request_repaint_after(std::time::Duration::from_millis(250)); + return; + } + if seg.mirrors().is_empty() { + ui.label( + egui::RichText::new(crate::i18n::tr( + "This build has no download location for the model yet: put the files in the folder above yourself (see docs/ai-masks.md).", + )) + .color(t.text_dim), + ); + } + if let Some(e) = error.map(str::to_string).or(d.error) { + ui.label(egui::RichText::new(format!("{} {e}", crate::i18n::tr("The download didn't work:"))).color(egui::Color32::from_rgb(230, 90, 80))); + } +} + pub fn confirm_dialog(app: &mut LightcraftApp, dlg: &Dialog) -> Result { match dlg { + Dialog::SamModel { then, .. } => { + if app.session.segmenter.installed() { + // installed: start what the user was doing + if let Some((kind, op)) = then { + crate::panels::masking::begin_ai(app, kind, op)?; + } + return Ok(serde_json::Value::Null); + } + let r = app.run("segment.model.download", json!({"acknowledged": true})); + if r.is_ok() { + app.ui.sam_downloading = true; + } + r + } Dialog::NewAlbum { name, folder } => app.run("album.create", json!({"name": name, "folder": folder, "addSelected": !folder})), Dialog::RenameAlbum { id, name } => app.run("album.rename", json!({"id": id, "name": name})), Dialog::TextPrompt { value, command, params, key, .. } => { diff --git a/crates/ui-egui/src/panels/masking.rs b/crates/ui-egui/src/panels/masking.rs index 86fa7aad..2e150a60 100644 --- a/crates/ui-egui/src/panels/masking.rs +++ b/crates/ui-egui/src/panels/masking.rs @@ -106,7 +106,7 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { 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())), + "prompt" => start_describe(app, "new"), "brush" | "linear" | "radial" => { app.ui.tool = kind.to_string(); if *kind != "brush" { @@ -124,9 +124,12 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { } }); describe_field(app, ui, true); - let status = if app.session.segmenter.busy() { + let seg = &app.session.segmenter; + let status = if seg.analyzing() { Some("Analyzing the photo for AI masks…") - } else if app.session.segmenter.detail_busy() || app.ui.detail_due.is_some() { + } else if seg.busy() { + Some("Selecting…") + } else if seg.detail_busy() || app.ui.detail_due.is_some() { Some("Refining the mask's detail…") } else { None @@ -136,6 +139,21 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { ui.label(egui::RichText::new(crate::i18n::tr(status)).color(t.text_dim)); ui.ctx().request_repaint_after(std::time::Duration::from_millis(250)); } + // the model download, while its dialog is closed + let download = seg.download_status(); + if download.running && app.ui.dialog.is_none() { + ui.add_space(4.0); + ui.horizontal(|ui| { + let pct = if download.total > 0 { download.done as f64 / download.total as f64 } else { 0.0 }; + let r = ui.add( + egui::ProgressBar::new(pct as f32).desired_width(ui.available_width() - 70.0).text(format!("SAM 3: {} %", (pct * 100.0).floor())), + ); + register(ui.ctx(), "maskSamProgress", r.rect); + if text_button(ui, "maskSamDetails", crate::i18n::tr("Details"), false).clicked() { + app.offer_sam_download(None); + } + }); + } }); divider(ui); // mask list @@ -144,17 +162,7 @@ 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, - "Always show the selected mask (O); hovering a mask shows it", - ) - .clicked() - { + if icon_button(ui, "maskOverlay", Icon::Eye, vec2(24.0, 24.0), on, true, "Show overlay (O)").clicked() { app.ui.mask_overlay = !on; } }); @@ -163,7 +171,6 @@ 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) { @@ -187,10 +194,6 @@ 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, @@ -297,9 +300,6 @@ 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())); } @@ -635,21 +635,62 @@ fn component_row_menu(app: &mut LightcraftApp, ui: &mut egui::Ui, mask: u32, k: } } +/// Whether the SAM 3 model is missing in a build that could use it: then the download is +/// offered (with `then` to start afterwards) instead of starting an AI mask. +fn needs_model(app: &mut LightcraftApp, kind: &str, op: &str) -> bool { + let seg = &app.session.segmenter; + let missing = lightcraft_engine::segment::Segmenter::AVAILABLE && seg.dir.is_some() && !seg.installed(); + if missing { + app.offer_sam_download(Some((kind, op))); + } + missing +} + /// 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) { +/// combined by `op`; the photo is analyzed meanwhile (in the background). +pub(crate) fn start_object(app: &mut LightcraftApp, ctx: &egui::Context, op: &str) { + if needs_model(app, "object", op) { + return; + } 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), + Err(e) => app.ai_error(ctx, e, Some(("object", op))), } } +/// Start an AI mask of `kind` (object|prompt) combined by `op` (new|add|subtract|intersect), +/// after the model was installed. +pub(crate) fn begin_ai(app: &mut LightcraftApp, kind: &str, op: &str) -> Result { + match kind { + "object" => { + let r = if op == "new" { + app.run("mask.add", json!({"kind": "object"})) + } else { + app.run("mask.addComponent", json!({"op": op, "kind": "object"})) + }?; + app.ui.tool = "object".into(); + Ok(r) + } + _ => { + app.ui.describe = Some((op.to_string(), String::new())); + Ok(serde_json::Value::Null) + } + } +} + +/// Open the Describe field (a new mask, or a component combined by `op`). +pub(crate) fn start_describe(app: &mut LightcraftApp, op: &str) { + if needs_model(app, "prompt", op) { + return; + } + app.ui.describe = Some((op.to_string(), String::new())); +} + /// 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 }; @@ -691,13 +732,9 @@ fn describe_field(app: &mut LightcraftApp, ui: &mut egui::Ui, new_mask: bool) { 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), + // the mask appears when the model has found it (a detail pass follows) + Ok(_) => app.ui.describe = None, + Err(e) => app.ai_error(ui.ctx(), e, Some(("prompt", op.as_str()))), } } } @@ -711,7 +748,7 @@ fn component_menu(app: &mut LightcraftApp, ui: &mut egui::Ui, op: &str) { 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())); + start_describe(app, op); } for (kind, label) in [ ("brush", "Brush"), diff --git a/crates/ui-egui/src/state.rs b/crates/ui-egui/src/state.rs index 5b8b21d2..bdbbcf29 100644 --- a/crates/ui-egui/src/state.rs +++ b/crates/ui-egui/src/state.rs @@ -291,9 +291,9 @@ pub struct UiState { /// (`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). + /// A SAM 3 download was started from the app (to report its end once). #[serde(skip)] - pub hover_mask: Option, + pub sam_downloading: bool, /// 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)] @@ -517,6 +517,15 @@ pub enum Dialog { Settings { tab: String, }, + /// Object and Describe masks need the SAM 3 model, which isn't installed: offer to download + /// it (size, licence, progress). `then`: the AI mask to start once it is there (`kind` + /// object|prompt, `op` new|add|subtract|intersect). + SamModel { + then: Option<(String, String)>, + /// Why the download couldn't start (shown in the dialog). + #[serde(default)] + error: Option, + }, /// Confirm moving photos to Recently Deleted. ConfirmDelete { count: usize, @@ -553,7 +562,7 @@ impl Default for UiState { histogram: true, soft_proof: false, proof: lightcraft_engine::pipeline::Proof { dest_warning: false, ..Default::default() }, - mask_overlay: false, + mask_overlay: true, mask_overlay_mode: "color".into(), mask_overlay_color: [230, 30, 40], mask_overlay_opacity: 50.0, @@ -573,8 +582,8 @@ impl Default for UiState { focus_search: false, renaming_mask: None, describe: None, - hover_mask: None, detail_due: None, + sam_downloading: false, 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 5be82af4..c3aac95c 100644 --- a/crates/ui-egui/src/tests_masking.rs +++ b/crates/ui-egui/src/tests_masking.rs @@ -47,20 +47,14 @@ 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)); - // no overlay until a mask is hovered in the Masks list; then that mask, in red + // the loupe asks the renderer for the selected mask's overlay 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; @@ -626,3 +620,53 @@ fn soft_proofing_flags_out_of_gamut_colours_and_makes_proof_copies() { assert!(!h.app.ui.soft_proof); assert_ne!(h.app.session.catalog.stack_of(copy.id).unwrap().collapsed, stack.collapsed, "S toggled the stack"); } + +/// Object / Describe without the SAM 3 model: the download is offered (never started without a +/// yes), the dialog says why it can't start when no mirror is configured, nothing freezes and no +/// empty mask is left behind. +#[test] +fn ai_masks_without_the_model_offer_the_download() { + use crate::state::Dialog; + use lightcraft_engine::segment::Segmenter; + let mut h = detail("panel.masking"); + let dir = std::env::temp_dir().join(format!("lc-ui-no-sam3-{}", std::process::id())); + h.app.session.segmenter.dir = Some(dir.clone()); + h.app.session.segmenter.mirrors_file = Some(dir.join("none.txt")); + let t = std::time::Instant::now(); + let r = h.request("ui.clickWidget", json!({"id": "maskNew:object"}), T); + assert_eq!(r["ok"], true, "{r}"); + assert!(develop(&h).masks.is_empty()); + if !Segmenter::AVAILABLE { + // a build without AI masks says so (a toast), no dialog + assert_eq!(h.app.ui.dialog, None); + return; + } + assert_eq!(h.app.ui.dialog, Some(Dialog::SamModel { then: Some(("object".into(), "new".into())), error: None })); + assert!(!h.app.session.segmenter.download_status().running, "nothing downloads without a yes"); + if h.app.session.segmenter.mirrors().is_empty() { + // Download: no location configured in this build → the reason, in the dialog + let r = h.request("ui.dialog.confirm", json!({}), T); + assert_eq!(r["ok"], false, "{r}"); + assert!(matches!(h.app.ui.dialog, Some(Dialog::SamModel { .. })), "stays open"); + let r = h.request("ui.clickWidget", json!({"id": "button:dialogOk"}), T); + assert_eq!(r["ok"], true, "{r}"); + let Some(Dialog::SamModel { error: Some(e), .. }) = &h.app.ui.dialog else { panic!("{:?}", h.app.ui.dialog) }; + assert!(e.contains("LIGHTCRAFT_SAM3_MIRRORS"), "{e}"); + } + // (a frame with the message laid out, so the buttons are where they are drawn) + h.step(); + h.step(); + let r = h.request("ui.clickWidget", json!({"id": "button:dialogCancel"}), T); + assert_eq!(r["ok"], true, "{r}"); + assert_eq!(h.app.ui.dialog, None); + // Describe asks too + let r = h.request("ui.clickWidget", json!({"id": "maskNew:prompt"}), T); + assert_eq!(r["ok"], true, "{r}"); + assert_eq!(h.app.ui.dialog, Some(Dialog::SamModel { then: Some(("prompt".into(), "new".into())), error: None })); + assert_eq!(h.app.ui.describe, None); + // a click command from an agent gets the not-installed error at once, and the app offers the model + let e = h.app.run("mask.add", json!({"kind": "prompt", "text": "sky"})).unwrap_err(); + assert!(e.starts_with(lightcraft_engine::segment::NOT_INSTALLED), "{e}"); + assert!(develop(&h).masks.is_empty()); + assert!(t.elapsed() < SETTLE, "{:?}", t.elapsed()); +} diff --git a/docs/ai-masks.md b/docs/ai-masks.md index dc3e6726..26d4f717 100644 --- a/docs/ai-masks.md +++ b/docs/ai-masks.md @@ -4,6 +4,10 @@ LightCraft selects objects with **SAM 3** (Segment Anything with Concepts, Meta 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). +**The model is optional.** It is not part of LightCraft, and nothing requires it: without it, +Object and Describe offer to download it, every other feature (and every existing mask, see +below) works, and no render, export or command ever waits for it. + ## Using them In the Masking panel (M): @@ -19,31 +23,87 @@ In the Masking panel (M): 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. + in the background); the result is kept as a high-resolution patch (at most 4 per selection), + 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. +click takes ~35 ms and each description ~0.6 s. Everything the model does runs on its own worker +thread: the window stays responsive, a click shows its dot at once and the selection follows +when it's computed ("Selecting…" / "Analyzing the photo…" under the mask tiles). The model is +unloaded after 10 minutes without use (it takes several GB of memory) and loaded again when +needed. 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 +## Getting 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: +### In the app (recommended) -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: +The first time you pick Object or Describe without the model, LightCraft asks: + +> **Download the SAM 3 model?** Object and Describe masks use SAM 3, Meta's segmentation model. +> It isn't part of LightCraft, and everything else works without it. Download it once (about +> 3.4 GB) · *folder*. Licence: SAM License (Meta) — Meta's terms, not LightCraft's. Downloading +> it means accepting them. [Read the SAM License] — **Not Now** / **Download** + +Nothing is downloaded until you choose **Download**. The download runs in the background (close +the dialog and keep working; the Masking panel shows its progress), can be cancelled, resumes +where it stopped if interrupted (also after quitting), and tries the next download location when +one fails or stalls. Each file is written as `.part` and moved into the model folder only +when complete and verified — `model.safetensors` must have the exact size and SHA-256 of Meta's +official checkpoint; a file that doesn't match is deleted, never used. If no location works, the +dialog says why and LightCraft carries on without the model. + +Agents and scripts use the same commands (control channel, `lightcraft-cli mcp --connect`): +`segment.model.status`, `segment.model.download {"acknowledged": true}` (refused without the +explicit `acknowledged`: pass it only after the user agreed to the download and the licence), +`segment.model.cancel`. + +### Download locations + +The model is fetched from an ordered list of mirrors: base URLs where `/model.safetensors`, +`/vocab.json` and `/merges.txt` live. The list is, in order: + +1. `LIGHTCRAFT_SAM3_MIRRORS` (environment variable; URLs separated by commas or spaces); +2. the file `models/sam3-mirrors.txt` in LightCraft's settings folder (one URL per line, `#` + comments) — `~/Library/Application Support/LightCraft/` (macOS), `%APPDATA%\LightCraft\` + (Windows), `~/.config/lightcraft/` (Linux); +3. the built-in list, `DEFAULT_MIRRORS` in `crates/segment/src/fetch/mod.rs`. + +> **Maintainers:** the built-in list is **empty** until LightCraft's own CDN locations exist +> (see the `TODO(maintainer)` there): add them in order of preference, host the three files +> unchanged, and pin `vocab.json` / `merges.txt` (size + SHA-256) in `SAM3_FILES` at the same +> time. Until then the in-app download needs a user-configured mirror, and the dialog says so. + +Hugging Face's `facebook/sam3` can't be a default: it is **gated** (each person must accept the +licence there, wait for approval and download with their own token). Someone with access can +use it as their mirror: + +```sh +LIGHTCRAFT_SAM3_MIRRORS=https://huggingface.co/facebook/sam3/resolve/main \ +LIGHTCRAFT_SAM3_TOKEN=hf_... lightcraft +``` + +(`LIGHTCRAFT_SAM3_TOKEN` is sent as a bearer token only over https, and only to the mirror's own +host — not to the storage host it redirects to.) + +The download is plain HTTP/1.1 over TLS in pure Rust (rustls with the RustCrypto provider and +the Mozilla root certificates; no OpenSSL, `ring` or `aws-lc`, nothing compiled from C), with a +15 s connect timeout and a 30 s stall timeout; it follows redirects but doesn't use HTTP proxies. + +### By hand + +Put `model.safetensors`, `vocab.json` and `merges.txt` from `facebook/sam3` in the model folder: +`models/sam3/` in the settings folder above, or the folder `LIGHTCRAFT_SAM3_DIR` names. Developers +can use the installer scripts, which download from Hugging Face with your token and verify the +weights: ```sh HF_TOKEN=hf_... tools/install-sam3.sh # macOS, Linux @@ -54,15 +114,79 @@ tools/install-sam3.sh --check # verify an installation $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. +## Licence of the model -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. +LightCraft's code is MIT OR Apache-2.0; `crates/segment` (a port of Apache-2.0 code) is +Apache-2.0. **The SAM 3 weights are neither**: they are Meta's, released under the **SAM +License** (), a custom licence that +is not an OSI open-source licence. It grants a non-exclusive, worldwide, non-transferable, +royalty-free licence to use, reproduce, distribute, copy, modify and create derivative works of +the "SAM Materials", subject to its terms, including: anyone redistributing them must provide a +copy of the licence with them, and derivative works are distributed under its terms; uses +subject to ITAR, military or warfare purposes, nuclear industries, espionage, guns or illegal +weapons, and reverse engineering are prohibited, and trade controls (sanctions, export rules) +apply; the materials come "as is", without warranty; Meta may terminate the licence on a breach +(then the materials must be deleted); California law governs. Read the licence itself — this +summary is not legal advice. + +Therefore: + +- **The weights are never part of this repository, its releases or installers**, and must not be + committed. +- They reach a user's computer only when that user asks for them, after the dialog names the + licence (or when they install them by hand). +- Whoever hosts a mirror redistributes the SAM Materials and must follow the licence (ship it + next to the files). + +## Commands (control channel, MCP) + +| Command | Params | | +|---|---|---| +| `mask.add` | `{kind: "object", points?: [[x,y],…], exclude?: [[x,y],…], seg?}` | an Object mask (empty until clicked) | +| `mask.add` | `{kind: "prompt", text, seg?}` | 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 (≤ 64 clicks) | +| `mask.refineDetail` | `{id?, component?}` | the zoomed-in detail pass → `{started}` | +| `segment.prepare` | `{}` | load the model and analyze the active photo → `{busy}` | +| `segment.model.status` | `{}` | `{available, installed, dir, loaded, busy, analyzing, sizeBytes, license, licenseUrl, mirrors, download}` | +| `segment.model.download` | `{acknowledged: true}` | start the download in the background → `{started}` | +| `segment.model.cancel` | `{}` | stop it (it resumes next time) → `{cancelled}` | + +Coordinates are normalized to the uncropped, oriented photo, like every mask shape. `seg` passes +a stored segmentation (as `mask.list` / the develop settings hold it): the mask is made without +the model. Without the model, the AI kinds fail with an error starting "The SAM 3 model is not +installed" that says how to get it. + +In the desktop app the commands don't wait for the model: `mask.objectPoint` returns +`{pending: true}` and the selection is applied when it's computed (one undo step per settled +selection); `mask.add {kind: "prompt"}` returns `{pending: true}` and the mask appears when +something is found (or a message says nothing was). `lightcraft-cli` and the headless MCP server +wait and return the result (when built with the `sam` feature; the default CLI has no AI masks). + +## Implementation notes + +- `crates/segment` ports the Hugging Face `transformers` implementation (Apache-2.0; see + `NOTICE` — the ported files say so and that they were modified): 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`. `fetch/` is the downloader (original code). +- 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`), only the tensors a path needs, and every + tensor's byte range is checked against the file's length when it is opened (a truncated or + damaged file is an error, never a crash). +- The engine (`crates/engine/src/segment/`) owns one worker thread with the model and the last + encoded photo; the session sends it jobs and receives results over channels, never holding a + lock the model holds. Each job runs under `catch_unwind`: a panic inside candle becomes an + error and the model is reloaded on the next request; busy counters are reset by drop guards. +- Hostile documents: a selection keeps at most 4 detail patches (when read and when rendered), + each at most 1024² logits; more than 64 clicks are never sent to the model. +- 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. **Developers:** the accuracy test against the Python reference needs `torch` and `transformers` (`tools/requirements-sam3-reference.txt`): @@ -72,32 +196,3 @@ python3 -m venv .venv-sam3 && .venv-sam3/bin/pip install -r tools/requirements-s .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 c15d1c09..82729971 100644 --- a/docs/parity.md +++ b/docs/parity.md @@ -69,9 +69,10 @@ 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): - 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. + Object and Describe masks run SAM 3 in pure Rust (`crates/segment`); the weights (SAM License) are never bundled: + the app offers a consented, verified download, but **LightCraft's CDN mirrors are not configured yet** (the + built-in list is empty, so today users need their own mirror or a manual install). Subject / Sky / People could + use the same model with fixed prompts; 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. @@ -329,7 +330,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 | 🟡 | `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-OBJECTS | Object selection | P2 | 🟡 | `cmd:mask.add` (`object`, `prompt`), `cmd:mask.objectPoint`, `cmd:mask.refineDetail`, `cmd:segment.prepare`, `cmd:segment.model.status`, `cmd:segment.model.download`, `cmd:segment.model.cancel`, `crates/segment`, `crates/engine/src/segment/mod.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; + / − under the selected mask; comma lists (`car, road`); a zoomed-in detail pass for 5–10× finer edges on small objects; per-selection Edge (hard ↔ soft). The model runs on its own worker thread (the UI never waits; panics become errors), is unloaded after 10 min idle, and is optional: the segmentation is stored with the mask (288² logits), so renders and exports never need it. Weights are never bundled (SAM License): without them the app offers a consented background download (mirrors, resume, timeouts, SHA-256) — **but no default download location is configured yet** (users set `LIGHTCRAFT_SAM3_MIRRORS` or install by hand; see docs/ai-masks.md). Not verified against Lightroom's Select Object; no brush/box object mode; the first click on a photo waits for its analysis (~4 s on an M4 Pro, much longer on CPU) | | 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/docs/whats-new.md b/docs/whats-new.md index 12984b9a..00480132 100644 --- a/docs/whats-new.md +++ b/docs/whats-new.md @@ -131,6 +131,12 @@ lost everything at quit; a temporary session shows a banner the whole time and never writes to your library. ### Editing +- AI masks with SAM 3 (Object and Describe in the Masking panel): click an object to select it (⌥-click leaves a + part out), or type what to select ("sky", "the red car", "car, road"); both combine with other masks, have an + Edge setting, and get a sharper zoomed-in pass in the background. The model runs inside LightCraft in pure Rust + and never freezes the window. It is optional and not part of LightCraft (Meta's SAM License, about 3.4 GB): the + first time you use an AI mask, LightCraft asks before downloading it, shows the progress, can cancel and resume, + and checks the file before using it. Masks keep their selection, so they render and export without the model. - Auto Sync: edits apply to every selected photo. Auto B&W mix. Automatic versions. - Colour-range masks: click the photo to sample. Luminance ranges: range bar, smoothness, luminance map. - ⌘-drag to straighten, ⇧G Guided Upright, a grid while transforming. diff --git a/tools/install-sam3.ps1 b/tools/install-sam3.ps1 index bd43fd98..0036c561 100644 --- a/tools/install-sam3.ps1 +++ b/tools/install-sam3.ps1 @@ -1,5 +1,6 @@ # 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). +# (see docs/ai-masks.md; tools/install-sam3.sh is the macOS/Linux version). The desktop app +# offers to download the model itself when an AI mask first needs it; this is for developers. # # 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 @@ -34,7 +35,26 @@ 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) + # into a .part file, moved into place only when complete (and, for the weights, verified): + # a failed or damaged download never ends up under the real name + $final = Join-Path $Dir $f + $part = "$final.part" + Remove-Item -Force -ErrorAction SilentlyContinue $part + try { + Invoke-WebRequest -Uri "$Base/$f" -Headers $Headers -OutFile $part + } catch { + Remove-Item -Force -ErrorAction SilentlyContinue $part + throw "download of $Repo/$f failed: $($_.Exception.Message)" + } + if ($f -eq "model.safetensors") { + Write-Host "checking model.safetensors (SHA-256 of 3.4 GB)..." + $got = (Get-FileHash $part -Algorithm SHA256).Hash.ToLower() + if ($got -ne $Sha256) { + Remove-Item -Force $part + throw "the download is damaged or not the official checkpoint (got $got); deleted it" + } + } + Move-Item -Force $part $final } 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 index 86554c39..4ec9d0a9 100755 --- a/tools/install-sam3.sh +++ b/tools/install-sam3.sh @@ -1,6 +1,9 @@ #!/bin/sh # Install the SAM 3 model for LightCraft's Object and Describe masks (see docs/ai-masks.md). # +# The desktop app offers to download the model itself when an AI mask first needs it; this script +# is for developers (reference tests, benchmarks) and for installing from Hugging Face directly. +# # 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: @@ -75,20 +78,29 @@ fi mkdir -p "$DIR" BASE="https://huggingface.co/$REPO/resolve/main" -fetch() { # fetch FILE DEST [resume] +# fetch FILE DEST [resume]: HTTP errors never reach DEST (--fail writes no error page), and a +# resume the server refuses (416) starts over instead of counting as done. +fetch() { 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=$? + code=$(curl -fsSL ${3:+-C -} -H "$auth" -w '%{http_code}' -o "$2" "$BASE/$1") || true case "$code" in - 200|206|416) return 0 ;; + 200|206) return 0 ;; + 416) + echo "the partial download of $1 doesn't match the server's; starting over" >&2 + rm -f "$2" + code=$(curl -fsSL -H "$auth" -w '%{http_code}' -o "$2" "$BASE/$1") || true + [ "$code" = 200 ] && return 0 + echo "download of $REPO/$1 failed ($code)" >&2; return 1 ;; 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 ;; + *) echo "download of $REPO/$1 failed (${code:-no response})" >&2; return 1 ;; esac } for f in $FILES; do echo "fetching $f" - fetch "$f" "$DIR/$f" + # small files: whole, into place only when complete + fetch "$f" "$DIR/$f.part" && mv -f "$DIR/$f.part" "$DIR/$f" || { rm -f "$DIR/$f.part"; exit 1; } done if [ -s "$DIR/model.safetensors" ] && [ "$(sha256 "$DIR/model.safetensors")" = "$SHA256" ]; then @@ -96,13 +108,21 @@ if [ -s "$DIR/model.safetensors" ] && [ "$(sha256 "$DIR/model.safetensors")" = " 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 + # a wrong model.safetensors is never resumed from: it goes + rm -f "$DIR/model.safetensors" + if [ -f "$part" ] && [ "$(wc -c < "$part" | tr -d ' ')" -ge "$SIZE" ]; then rm -f "$part"; fi + fetch model.safetensors "$part" resume || exit 1 have=$(wc -c < "$part" | tr -d ' ') - if [ "$have" != "$SIZE" ]; then + if [ "$have" -lt "$SIZE" ]; then echo "incomplete download ($have of $SIZE bytes); run the script again to resume" >&2 exit 1 fi + echo "checking model.safetensors (SHA-256 of 3.4 GB, takes a few seconds)…" + if [ "$have" != "$SIZE" ] || [ "$(sha256 "$part")" != "$SHA256" ]; then + rm -f "$part" + echo "the download is damaged or not the official checkpoint; deleted it, run the script again" >&2 + exit 1 + fi mv -f "$part" "$DIR/model.safetensors" fi From d5be20143159afc2ea19bf9444860d56283835b9 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 7 Oct 2026 10:06:46 -0400 Subject: [PATCH 4/4] M12: SAM 3 download dialog: no Download button without a location, clearer wording Looked at in the running app: without a configured mirror the dialog offered Download / Try Again that could only fail; it now shows the manual-install hint and Close (an agent confirming still gets the reason). "A one-time download of about 3.4 GB, saved in "; a cancelled download is reported as stopped (resumable), not failed. Co-Authored-By: Claude Opus 5.5 --- crates/ui-egui/src/lib.rs | 3 +++ crates/ui-egui/src/panels/dialogs.rs | 8 +++++--- crates/ui-egui/src/tests_masking.rs | 11 ++++++++--- 3 files changed, 16 insertions(+), 6 deletions(-) diff --git a/crates/ui-egui/src/lib.rs b/crates/ui-egui/src/lib.rs index 49ea140f..a4addbf3 100644 --- a/crates/ui-egui/src/lib.rs +++ b/crates/ui-egui/src/lib.rs @@ -389,6 +389,9 @@ impl LightcraftApp { let open = matches!(self.ui.dialog, Some(state::Dialog::SamModel { .. })); match (&download.error, open) { (_, true) => {} + (Some(e), false) if e.contains("cancelled") => { + self.toast(ctx, crate::i18n::tr("SAM 3 download stopped: it resumes where it left off next time.")) + } (Some(e), false) => self.toast_error(ctx, format!("The SAM 3 download failed: {e}")), (None, false) if download.finished => { self.toast_error(ctx, crate::i18n::tr("The SAM 3 model is installed: Object and Describe masks are ready.")) diff --git a/crates/ui-egui/src/panels/dialogs.rs b/crates/ui-egui/src/panels/dialogs.rs index 67d495de..bc23383b 100644 --- a/crates/ui-egui/src/panels/dialogs.rs +++ b/crates/ui-egui/src/panels/dialogs.rs @@ -771,8 +771,10 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) { let informational = matches!(dlg, Dialog::About | Dialog::Shortcuts | Dialog::Settings { .. }); let sam = &app.session.segmenter; let (sam_installed, sam_running, sam_failed) = (sam.installed(), sam.download_status().running, sam.download_status().error.is_some()); + // no download location in this build: nothing to offer but the manual install + let sam_nowhere = !sam_installed && !sam_running && sam.mirrors().is_empty(); let cancel = match &dlg { - Dialog::SamModel { .. } if sam_running || sam_installed => "Close", + Dialog::SamModel { .. } if sam_running || sam_installed || sam_nowhere => "Close", Dialog::SamModel { .. } => "Not Now", _ => "Cancel", }; @@ -794,7 +796,7 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) { Dialog::Merge { .. } => "Merge", Dialog::ConfirmDelete { .. } => "Delete", Dialog::SamModel { then: Some(_), .. } if sam_installed => "Continue", - Dialog::SamModel { .. } if sam_installed || sam_running => "", + Dialog::SamModel { .. } if sam_installed || sam_running || sam_nowhere => "", Dialog::SamModel { error, .. } if error.is_some() || sam_failed => "Try Again", Dialog::SamModel { .. } => "Download", _ if informational => "Close", @@ -876,7 +878,7 @@ fn sam_model_body(app: &mut LightcraftApp, ui: &mut egui::Ui, error: Option<&str "Object and Describe masks use SAM 3, Meta's segmentation model. It isn't part of LightCraft, and everything else works without it.", )); let dir = seg.dir.as_ref().map(|d| d.display().to_string()).unwrap_or_default(); - ui.label(format!("{} {:.1} GB · {dir}", crate::i18n::tr("Download it once (about):"), gb(MODEL_BYTES))); + ui.label(format!("{} {:.1} GB, {} {dir}", crate::i18n::tr("A one-time download of about"), gb(MODEL_BYTES), crate::i18n::tr("saved in"))); ui.label( egui::RichText::new(format!( "{} {LICENSE_NAME} — {}", diff --git a/crates/ui-egui/src/tests_masking.rs b/crates/ui-egui/src/tests_masking.rs index c3aac95c..2abc141f 100644 --- a/crates/ui-egui/src/tests_masking.rs +++ b/crates/ui-egui/src/tests_masking.rs @@ -644,14 +644,19 @@ fn ai_masks_without_the_model_offer_the_download() { assert_eq!(h.app.ui.dialog, Some(Dialog::SamModel { then: Some(("object".into(), "new".into())), error: None })); assert!(!h.app.session.segmenter.download_status().running, "nothing downloads without a yes"); if h.app.session.segmenter.mirrors().is_empty() { - // Download: no location configured in this build → the reason, in the dialog + // no location configured in this build: no Download button (only Close), and an agent + // confirming anyway gets the reason; the dialog stays + h.step(); + let r = h.request("ui.clickWidget", json!({"id": "button:dialogOk"}), T); + assert_eq!(r["ok"], false, "{r}"); let r = h.request("ui.dialog.confirm", json!({}), T); assert_eq!(r["ok"], false, "{r}"); + assert!(r["error"].as_str().unwrap_or_default().contains("LIGHTCRAFT_SAM3_MIRRORS"), "{r}"); assert!(matches!(h.app.ui.dialog, Some(Dialog::SamModel { .. })), "stays open"); + } else { + // with a mirror: Download starts it in the background (here it fails: nothing listens) let r = h.request("ui.clickWidget", json!({"id": "button:dialogOk"}), T); assert_eq!(r["ok"], true, "{r}"); - let Some(Dialog::SamModel { error: Some(e), .. }) = &h.app.ui.dialog else { panic!("{:?}", h.app.ui.dialog) }; - assert!(e.contains("LIGHTCRAFT_SAM3_MIRRORS"), "{e}"); } // (a frame with the message laid out, so the buttons are where they are drawn) h.step();