shifathosn
/
lightcraft
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

M12: SAM 3 Object and Describe masks (pure Rust, Metal) 

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) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01C8BkS7JDqa1EMcjW6spDGA
This commit is contained in:
Guinther Kovalski 2026-10-06 20:15:04 -03:00
parent 5cee0c08d0
commit 6468bfc06b
40 changed files with 4328 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

712
Cargo.lock generated
View File

@ -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"

5
Cargo.toml
View File

@ -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"

6
NOTICE
View File

@ -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

2
apps/lightcraft/Cargo.toml
View File

@ -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 }

35
apps/lightcraft/src/main.rs
View File

@ -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: <settings folder>/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<StderrLog> = 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<u16> = 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 <config>/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;

1
crates/develop/Cargo.toml
View File

@ -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 }

2
crates/develop/src/lib.rs
View File

@ -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;

131
crates/develop/src/segmask.rs Normal file
View File

@ -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<u8> = 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<Vec<f32>> {
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<Vec<u8>> {
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<u8> = (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<f32> = (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());
}
}

25
crates/develop/src/settings.rs
View File

@ -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<Point>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
exclude: Vec<Point>,
#[serde(default, skip_serializing_if = "Option::is_none")]
seg: Option<crate::SegMask>,
/// 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<crate::SegMask>,
/// 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<crate::SegMask>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
detail: Vec<crate::SegMask>,
/// Edge −100..100, as for `Object`.
#[serde(default, skip_serializing_if = "is_zero")]
edge: f64,
},
People {
person: u32,

5
crates/engine/Cargo.toml
View File

@ -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 }

141
crates/engine/src/cmd/masks.rs
View File

@ -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<MaskShape> {
@ -38,10 +39,57 @@ fn shape_from(kind: &str, p: &Value, c: &str) -> Result<MaskShape> {
.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<Point> {
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<String>) -> crate::EngineError {
crate::EngineError::Other(msg.into())
}
fn clicks_of(include: &[Point], exclude: &[Point]) -> Vec<Click> {
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<Mask>, &mut Option<u32>) -> Result<()>) -> Result<Value> {
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<CommandSpec> {
"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<CommandSpec> {
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<CommandSpec> {
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",

4
crates/engine/src/lib.rs
View File

@ -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<u32>,
/// 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<usize>,
/// 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(),

588
crates/engine/src/segment.rs Normal file
View File

@ -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<u8>,
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<Sam3>,
/// The last encoded image and its key.
pub cache: Option<(u64, Encoded)>,
}
#[derive(Clone, Default)]
pub struct Shared {
pub inner: Arc<Mutex<Inner>>,
pub busy: Arc<AtomicBool>,
/// A zoomed-in detail pass is running.
pub detail_busy: Arc<AtomicBool>,
/// Finished detail passes, waiting for [`crate::Session::segment_poll`].
pub results: Arc<Mutex<Vec<super::DetailResult>>>,
}
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<PathBuf>,
#[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<PathBuf, String> {
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<SegMask>,
}
/// What a detail pass re-runs.
#[cfg_attr(not(feature = "sam"), allow(dead_code))]
enum Prompt {
Clicks(Vec<Click>),
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<Input, String> {
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<SegMask, String> {
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<lightcraft_segment::Click> =
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<Option<SegMask>, 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<Vec<f32>> = 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<Input, String> {
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<bool, String> {
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<DetailResult> =
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<Either, String> {
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<Option<SegMask>, 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<lightcraft_segment::Click> = 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<Vec<f32>> = 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<f32> =
(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());
}
}

34
crates/engine/src/tests.rs
View File

@ -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());
}

141
crates/pipeline/src/masks.rs
View File

@ -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<f32>,
side: usize,
r: [f64; 4],
}
impl SegGrid {
fn new(seg: &SegMask) -> Option<SegGrid> {
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<SegGrid> = 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<f32> = (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<f32> = (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<f32> = (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);

24
crates/segment/Cargo.toml Normal file
View File

@ -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"] }

25
crates/segment/examples/sam3_bench.rs Normal file
View File

@ -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 -- <model dir>`.
fn main() -> Result<(), Box<dyn std::error::Error>> {
let mut args = std::env::args().skip(1);
let dir = std::path::PathBuf::from(args.next().ok_or("usage: sam3_bench <model dir>")?);
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<u8> = (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(())
}

77
crates/segment/src/clip.rs Normal file
View File

@ -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<Layer>,
final_ln: LayerNorm,
device: Device,
}
impl TextEncoder {
/// `vb` at `detector_model.text_encoder.text_model`.
pub fn new(vb: VarBuilder) -> Result<Self> {
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<Tensor> {
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)?)
}
}

304
crates/segment/src/detector.rs Normal file
View File

@ -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<Tensor>,
/// 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<f32>,
/// 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<EncoderLayer>,
decoder: Vec<DecoderLayer>,
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<Conv2d>,
pixel_norms: Vec<GroupNorm>,
instance_proj: Conv2d,
mask_embedder: Vec<Linear>,
}
impl Detector {
/// `vb` at `detector_model`.
pub fn new(vb: VarBuilder) -> Result<Self> {
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::<Result<_>>()?,
pixel_norms: (0..3).map(|i| Ok(candle_nn::group_norm(8, HIDDEN, 1e-5, p.pp(format!("norms.{i}")))?)).collect::<Result<_>>()?,
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::<Result<_>>()?,
})
}
/// The detector's feature pyramid for the backbone output `[1, 1024, 72, 72]`.
pub fn features(&self, backbone: &Tensor) -> Result<Features> {
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<Tensor> {
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<Tensor> {
let d = (d * 8.0)?;
Ok(d.sign()?.mul(&((d.abs()? + 1.0)?.log()? / (8f64.ln()))?)?)
};
let pick = |a: usize, b: usize| -> Result<Tensor> { 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<Tensor> {
let feats = HIDDEN / 2;
let dim_t: Vec<f32> = (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<f32> = (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<Tensor> {
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<Output> {
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<f32> = valid.iter().map(|v| if *v { 1.0 } else { 0.0 }).collect();
let count = w.iter().sum::<f32>().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<Prediction> {
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<Vec<f32>> = 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 })
}

252
crates/segment/src/lib.rs Normal file
View File

@ -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<T> = std::result::Result<T, Error>;
/// 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<weights::Weights>,
backbone: vit::Vit,
tracker_neck: Option<neck::Neck>,
tracker: Option<tracker::Tracker>,
detector: Option<detector::Detector>,
tokenizer: Option<tokenizer::Tokenizer>,
}
#[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<tracker::Features>,
detector: Option<detector::Features>,
}
/// 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<f32>,
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> {
Self::load_on(dir, best_device())
}
pub fn load_on(dir: &Path, device: Device) -> Result<Self> {
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<weights::Weights>, 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<Encoded> {
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<Encoded> {
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<Prediction> {
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<Option<Prediction>> {
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<detector::Output> {
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()?)
}

126
crates/segment/src/mask.rs Normal file
View File

@ -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<f32>)> {
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<f32> = (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<Vec<f32>> {
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<f32> {
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<u8> = (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<f32> = (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));
}
}

73
crates/segment/src/neck.rs Normal file
View File

@ -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<Level>,
}
impl Neck {
/// `vb` at `…neck` (detector) or `tracker_neck`.
pub fn new(vb: VarBuilder) -> Result<Self> {
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<Vec<Tensor>> {
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)
}
}

238
crates/segment/src/nn.rs Normal file
View File

@ -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<LayerNorm> {
Ok(candle_nn::layer_norm(d, eps, vb)?)
}
pub fn linear(i: usize, o: usize, vb: VarBuilder) -> Result<Linear> {
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<Tensor> {
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> {
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<Self> {
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<Tensor> {
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<Tensor> {
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<Tensor> {
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<Tensor> {
let v: Vec<f32> = 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<Self> {
Ok(Self { fc1: linear(hidden, inter, vb.pp("fc1"))?, fc2: linear(inter, hidden, vb.pp("fc2"))?, gelu })
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
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<Linear>,
}
impl DecoderMlp {
pub fn new(input: usize, hidden: usize, output: usize, n: usize, vb: VarBuilder) -> Result<Self> {
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<Tensor> {
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<Linear>,
proj_out: Linear,
sigmoid: bool,
}
impl FeedForward {
pub fn new(input: usize, hidden: usize, output: usize, n: usize, sigmoid: bool, vb: VarBuilder) -> Result<Self> {
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<Tensor> {
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<Tensor> {
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<Tensor> {
let scale = 2.0 * std::f64::consts::PI;
let dim_t: Vec<f64> = (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<Tensor> {
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<Tensor> {
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<Vec<f32>> {
Ok(t.to_dtype(DType::F32)?.flatten_all()?.to_vec1::<f32>()?)
}

173
crates/segment/src/tokenizer.rs Normal file
View File

@ -0,0 +1,173 @@
//! The CLIP byte-pair-encoding tokenizer (lower-cased, `</w>` 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<String, u32>,
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<Self> {
let read = |name: &str| std::fs::read_to_string(dir.join(name)).map_err(|e| Error::Model(format!("{name}: {e}")));
let vocab: HashMap<String, u32> = 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<String, u32>, 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<u32>, Vec<bool>) {
let mut ids = vec![BOS];
for word in pre_tokenize(&text.split_whitespace().collect::<Vec<_>>().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<String> {
let mut syms: Vec<String> = word.chars().map(String::from).collect();
if let Some(last) = syms.last_mut() {
last.push_str("</w>");
}
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<String> {
let chars: Vec<char> = 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<String, u32> =
[("c", 1), ("a", 2), ("t</w>", 3), ("ca", 4), ("cat</w>", 5)].into_iter().map(|(k, v)| (k.to_string(), v)).collect();
let t = Tokenizer::new(vocab, "#version: 0.2\nc a\nca t</w>\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);
}
}

249
crates/segment/src/tracker.rs Normal file
View File

@ -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<f32>,
pub iou: f32,
pub object_score: f32,
}
/// Random-Fourier positional encoding of points normalized to 0..1 (`Sam3TrackerPositionalEmbedding`).
struct Fourier {
pe: Vec<f32>, // [2, 128]
}
impl Fourier {
fn encode(&self, x: f32, y: f32) -> Vec<f32> {
let half = self.pe.len() / 2;
let (cx, cy) = (2.0 * x - 1.0, 2.0 * y - 1.0);
let proj: Vec<f32> = (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<Tensor> {
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<TwoWayBlock>,
final_t2i: Attention,
ln_final: LayerNorm,
up1: ConvTranspose2d,
up_ln: LayerNorm,
up2: ConvTranspose2d,
hyper: Vec<FeedForward>,
iou_head: FeedForward,
obj_head: FeedForward,
}
fn read_vec(vb: &VarBuilder, shape: (usize, usize), name: &str) -> Result<Vec<f32>> {
to_vec(&vb.get(shape, name)?)
}
impl Tracker {
/// `vb` at `tracker_model`.
pub fn new(vb: VarBuilder) -> Result<Self> {
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::<Result<_>>()?,
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<Features> {
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<Tensor> {
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(&not_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<Prediction> {
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::<Result<Vec<_>>>()?;
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<Prediction> {
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::<f32>()?;
let union = single.gt(-0.05)?.to_dtype(DType::F32)?.sum_all()?.to_scalar::<f32>()?;
let stability = if union > 0.0 { inter / union } else { 1.0 };
if stability >= 0.98 { pick(0) } else { pick(best_multi()) }
}
}

166
crates/segment/src/vit.rs Normal file
View File

@ -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<f64> = (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<Tensor> {
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<Tensor> {
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<Tensor> {
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<Tensor> {
let (b, n, _) = x.dims3()?;
let split = |t: Tensor| -> Result<Tensor> { 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<Layer>,
rope_window: (Tensor, Tensor),
rope_global: (Tensor, Tensor),
}
impl Vit {
/// `vb` at `detector_model.vision_encoder.backbone`.
pub fn new(vb: VarBuilder) -> Result<Self> {
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<Tensor> {
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()?)
}
}

108
crates/segment/src/weights.rs Normal file
View File

@ -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<usize>,
start: u64,
end: u64,
}
/// An open safetensors file: its header, and the file to read tensor bytes from.
pub struct Weights {
file: Mutex<File>,
entries: HashMap<String, Entry>,
data_start: u64,
}
impl Weights {
pub fn open(path: &Path) -> Result<Arc<Self>> {
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<String, serde_json::Value> =
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<usize> = 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<u64> =
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<Tensor> {
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<f32> = 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<Weights>);
impl SimpleBackend for Backend {
fn get(&self, s: Shape, name: &str, _h: candle_nn::Init, dtype: DType, dev: &Device) -> candle_core::Result<Tensor> {
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<Tensor> {
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)
}
}

97
crates/segment/tests/reference.rs Normal file
View File

@ -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=<model dir> LIGHTCRAFT_SAM3_REF=<ref.safetensors> 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<String, Tensor>)> {
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<f32> {
t.to_device(&Device::Cpu).unwrap().flatten_all().unwrap().to_vec1::<f32>().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<i64> = r[&format!("tok{i}")].flatten_all().unwrap().to_vec1().unwrap();
assert_eq!(ids.iter().map(|v| *v as i64).collect::<Vec<_>>(), 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<f32> = 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);
}
}

30
crates/ui-egui/src/lib.rs
View File

@ -260,6 +260,7 @@ impl LightcraftApp {
}
let r = self.session.execute(id, &params).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<String>) {
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));

45
crates/ui-egui/src/panels/detail.rs
View File

@ -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;

158
crates/ui-egui/src/panels/masking.rs
View File

@ -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"),

16
crates/ui-egui/src/state.rs
View File

@ -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<u32>,
/// 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,

12
crates/ui-egui/src/tests_masking.rs
View File

@ -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;

103
docs/ai-masks.md Normal file
View File

@ -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 <https://huggingface.co/facebook/sam3>, 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=<model 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.

7
docs/parity.md
View File

@ -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 |

40
tools/install-sam3.ps1 Normal file
View File

@ -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."

110
tools/install-sam3.sh Executable file
View File

@ -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."

8
tools/requirements-sam3-reference.txt Normal file
View File

@ -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

71
tools/sam3_reference.py Normal file
View File

@ -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=<model 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)

1
xtask/src/layers.rs
View File

@ -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)),