diff --git a/.gitignore b/.gitignore index 823f9276..51531a34 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,6 @@ plan/ .DS_Store .claude/ /craft-fonts +# model weights are never committed (SAM 3: docs/ai-masks.md) +*.safetensors +*.safetensors.part diff --git a/Cargo.lock b/Cargo.lock index aa91c3d9..3f086161 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -146,6 +146,41 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "aead" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d122413f284cf2d62fb1b7db97e02edb8cda96d769b16e443a4f6195e35662b0" +dependencies = [ + "crypto-common", + "generic-array", +] + +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures", +] + +[[package]] +name = "aes-gcm" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "831010a0f742e1209b3bcea8fab6a8e149051ba6099432c8cb2cc117dec3ead1" +dependencies = [ + "aead", + "aes", + "cipher", + "ctr", + "ghash", + "subtle", +] + [[package]] name = "ahash" version = "0.8.12" @@ -543,6 +578,18 @@ dependencies = [ "windows-link", ] +[[package]] +name = "base16ct" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c7f02d4ea65f2c1853089ffd8d2787bdbc63de2f0d29dedbcf8ccdfa0ccd4cf" + +[[package]] +name = "base64ct" +version = "1.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" + [[package]] name = "bit-set" version = "0.8.0" @@ -594,6 +641,21 @@ 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 = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "block2" version = "0.5.1" @@ -678,6 +740,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 +809,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 0.9.5", + "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" @@ -774,6 +912,41 @@ dependencies = [ "libc", ] +[[package]] +name = "chacha20" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3613f74bd2eac03dad61bd53dbe620703d4371614fe0bc3b9f04dd36fe4e818" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures", +] + +[[package]] +name = "chacha20poly1305" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10cd79432192d1c0f4e1a0fef9527696cc039165d729fb41b3f4f4f354c2dc35" +dependencies = [ + "aead", + "chacha20", + "cipher", + "poly1305", + "zeroize", +] + +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common", + "inout", + "zeroize", +] + [[package]] name = "clipboard-win" version = "5.4.1" @@ -828,6 +1001,12 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "const-oid" +version = "0.9.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" + [[package]] name = "core-foundation" version = "0.9.4" @@ -868,6 +1047,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.5.2" @@ -917,12 +1105,81 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +[[package]] +name = "crypto-bigint" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0dc92fb57ca44df6db8059111ab3af99a63d5d0f8375d9972e319a379c6bab76" +dependencies = [ + "generic-array", + "rand_core 0.6.4", + "subtle", + "zeroize", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "ctr" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835" +dependencies = [ + "cipher", +] + [[package]] name = "cursor-icon" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f27ae1dd37df86211c42e150270f82743308803d90a6f6e6651cd730d5e1732f" +[[package]] +name = "curve25519-dalek" +version = "4.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97fb8b7c4503de7d6ae7b42ab72a5a59857b4c937ec27a3d4539dba95b5ab2be" +dependencies = [ + "cfg-if", + "cpufeatures", + "curve25519-dalek-derive", + "digest", + "fiat-crypto", + "rustc_version", + "subtle", + "zeroize", +] + +[[package]] +name = "curve25519-dalek-derive" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "der" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" +dependencies = [ + "const-oid", + "pem-rfc7468", + "zeroize", +] + [[package]] name = "dhat" version = "0.3.3" @@ -939,6 +1196,18 @@ dependencies = [ "thousands", ] +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "const-oid", + "crypto-common", + "subtle", +] + [[package]] name = "dispatch" version = "0.2.0" @@ -987,6 +1256,36 @@ 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 = "ecdsa" +version = "0.16.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee27f32b5c5292967d2d4a9d7f1e0b0aed2c15daded5a60300e4abb9d8020bca" +dependencies = [ + "der", + "digest", + "elliptic-curve", + "rfc6979", + "signature", + "spki", +] + [[package]] name = "ecolor" version = "0.36.2" @@ -997,6 +1296,30 @@ dependencies = [ "emath", ] +[[package]] +name = "ed25519" +version = "2.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53" +dependencies = [ + "pkcs8", + "signature", +] + +[[package]] +name = "ed25519-dalek" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9" +dependencies = [ + "curve25519-dalek", + "ed25519", + "serde", + "sha2", + "subtle", + "zeroize", +] + [[package]] name = "eframe" version = "0.36.2" @@ -1114,6 +1437,27 @@ version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" +[[package]] +name = "elliptic-curve" +version = "0.13.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5e6043086bf7973472e0c7dff2142ea0b680d30e18d9cc40f267efbf222bd47" +dependencies = [ + "base16ct", + "crypto-bigint", + "digest", + "ff", + "generic-array", + "group", + "hkdf", + "pem-rfc7468", + "pkcs8", + "rand_core 0.6.4", + "sec1", + "subtle", + "zeroize", +] + [[package]] name = "emath" version = "0.36.2" @@ -1129,6 +1473,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" @@ -1279,6 +1635,22 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b97b65636e5b9ef369943878ac74335ba1c55c1cb6adbf1e2c293c624248d693" +[[package]] +name = "ff" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0b50bfb653653f9ca9095b427bed08ab8d75a137839d9ad64eb11810d5b6393" +dependencies = [ + "rand_core 0.6.4", + "subtle", +] + +[[package]] +name = "fiat-crypto" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d" + [[package]] name = "find-msvc-tools" version = "0.1.14" @@ -1296,6 +1668,18 @@ dependencies = [ "zlib-rs", ] +[[package]] +name = "float8" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "719a903cc23e4a89e87962c2a80fdb45cdaad0983a89bd150bb57b4c8571a7d5" +dependencies = [ + "half", + "num-traits", + "rand 0.9.5", + "rand_distr", +] + [[package]] name = "fnv" version = "1.0.7" @@ -1399,6 +1783,255 @@ 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 = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", + "zeroize", +] + [[package]] name = "gethostname" version = "1.1.0" @@ -1409,6 +2042,17 @@ dependencies = [ "windows-link", ] +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + [[package]] name = "getrandom" version = "0.3.4" @@ -1432,6 +2076,16 @@ dependencies = [ "r-efi 6.0.0", ] +[[package]] +name = "ghash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0d8a4362ccb29cb0b265253fb0a2728f592895ee6854fd9bc13f2ffda266ff1" +dependencies = [ + "opaque-debug", + "polyval", +] + [[package]] name = "gif" version = "0.14.2" @@ -1567,6 +2221,17 @@ dependencies = [ "windows", ] +[[package]] +name = "group" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0f9ef7462f7c099f518d754361858f86d8a07af53ba9af0fe635bbccb151a63" +dependencies = [ + "ff", + "rand_core 0.6.4", + "subtle", +] + [[package]] name = "guillotiere" version = "0.7.0" @@ -1582,9 +2247,12 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "num-traits", + "rand 0.9.5", + "rand_distr", "zerocopy", ] @@ -1609,6 +2277,8 @@ dependencies = [ "allocator-api2", "equivalent", "foldhash", + "serde", + "serde_core", ] [[package]] @@ -1620,6 +2290,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" @@ -1632,6 +2308,24 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "hkdf" +version = "0.12.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b5f8eb2ad728638ea2c7d47a21db23b7b58a72ed6a38256b8a1849f15fbbdf7" +dependencies = [ + "hmac", +] + +[[package]] +name = "hmac" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" +dependencies = [ + "digest", +] + [[package]] name = "image" version = "0.25.10" @@ -1677,6 +2371,15 @@ dependencies = [ "hashbrown 0.17.1", ] +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "generic-array", +] + [[package]] name = "interpolate_name" version = "0.2.4" @@ -2026,6 +2729,9 @@ name = "lazy_static" version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" +dependencies = [ + "spin", +] [[package]] name = "libc" @@ -2159,6 +2865,7 @@ version = "0.2.1" dependencies = [ "lightcraft-color", "lightcraft-geom", + "miniz_oxide 0.8.9", "proptest", "serde", "serde_json", @@ -2183,6 +2890,7 @@ dependencies = [ "lightcraft-raster", "lightcraft-raw", "lightcraft-scenes", + "lightcraft-segment", "log", "miniz_oxide 0.8.9", "rayon", @@ -2325,6 +3033,22 @@ dependencies = [ "serde", ] +[[package]] +name = "lightcraft-segment" +version = "0.2.1" +dependencies = [ + "candle-core", + "candle-nn", + "log", + "rustls", + "rustls-rustcrypto", + "serde", + "serde_json", + "sha2", + "thiserror 2.0.21", + "webpki-roots", +] + [[package]] name = "lightcraft-sysmem" version = "0.2.1" @@ -2428,6 +3152,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" @@ -2451,6 +3184,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" dependencies = [ "libc", + "stable_deref_trait", ] [[package]] @@ -2462,6 +3196,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" @@ -2631,6 +3380,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" @@ -2641,6 +3404,32 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-bigint-dig" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e661dda6640fad38e827a6d4a310ff4763082116fe217f279885c97f511bb0b7" +dependencies = [ + "lazy_static", + "libm", + "num-integer", + "num-iter", + "num-traits", + "rand 0.8.8", + "smallvec", + "zeroize", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "bytemuck", + "num-traits", +] + [[package]] name = "num-derive" version = "0.4.2" @@ -2661,6 +3450,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" @@ -2682,6 +3481,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" @@ -2704,6 +3513,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" @@ -2870,6 +3688,7 @@ checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272" dependencies = [ "bitflags 2.13.2", "block2 0.6.2", + "libc", "objc2 0.6.4", "objc2-core-foundation", ] @@ -2917,7 +3736,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", ] @@ -3030,6 +3851,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "opaque-debug" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" + [[package]] name = "orbclient" version = "0.3.55" @@ -3068,6 +3895,30 @@ dependencies = [ "ttf-parser", ] +[[package]] +name = "p256" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9863ad85fa8f4460f9c48cb909d38a0d689dba1f6f6988a5e3e0d31071bcd4b" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2", +] + +[[package]] +name = "p384" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe42f1670a52a47d448f14b6a5c61dd78fce51856e68edaa38f7ae3a46b8d6b6" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2", +] + [[package]] name = "parking" version = "2.2.1" @@ -3109,6 +3960,15 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35fb2e5f958ec131621fdd531e9fc186ed768cbe395337403ae56c17a74c68ec" +[[package]] +name = "pem-rfc7468" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88b39c9bfcfc231068454382784bb460aae594343fb030d46e9f50a645418412" +dependencies = [ + "base64ct", +] + [[package]] name = "peniko" version = "0.6.1" @@ -3208,6 +4068,38 @@ dependencies = [ "futures-io", ] +[[package]] +name = "pkcs1" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8ffb9f10fa047879315e6625af03c164b16962a5368d724ed16323b68ace47f" +dependencies = [ + "der", + "pkcs8", + "spki", +] + +[[package]] +name = "pkcs5" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e847e2c91a18bfa887dd028ec33f2fe6f25db77db3619024764914affe8b69a6" +dependencies = [ + "der", + "spki", +] + +[[package]] +name = "pkcs8" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" +dependencies = [ + "der", + "pkcs5", + "spki", +] + [[package]] name = "pkg-config" version = "0.3.34" @@ -3259,6 +4151,17 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bc6355899e1c9462875b6757c79f3caa011a1fdae12bbb1a2e72dd1f234f8336" +[[package]] +name = "poly1305" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8159bd90725d2df49889a078b54f4f79e87f1f8a8444194cdca81d38f5393abf" +dependencies = [ + "cpufeatures", + "opaque-debug", + "universal-hash", +] + [[package]] name = "polycool" version = "0.4.0" @@ -3268,6 +4171,18 @@ dependencies = [ "arrayvec", ] +[[package]] +name = "polyval" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1fe60d06143b2430aa532c94cfe9e29783047f06c0d7fd359a9a51b729fa25" +dependencies = [ + "cfg-if", + "cpufeatures", + "opaque-debug", + "universal-hash", +] + [[package]] name = "portable-atomic" version = "1.15.0" @@ -3298,6 +4213,15 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8cf8e6a8aa66ce33f63993ffc4ea4271eb5b0530a9002db8455ea6050c77bfa" +[[package]] +name = "primeorder" +version = "0.13.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "353e1ca18966c16d9deb1c69278edbc5f194139612772bd9537af60ac231e1e6" +dependencies = [ + "elliptic-curve", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -3345,8 +4269,8 @@ dependencies = [ "bit-vec 0.8.0", "bitflags 2.13.2", "num-traits", - "rand", - "rand_chacha", + "rand 0.9.5", + "rand_chacha 0.9.0", "rand_xorshift", "regex-syntax", "rusty-fork", @@ -3354,6 +4278,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" @@ -3411,14 +4372,34 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "rand" +version = "0.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e058c7de0b26af77780c769414d6257830bb240f3c38477dbc2c16e5f54d6d4c" +dependencies = [ + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + [[package]] name = "rand" version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" dependencies = [ - "rand_chacha", - "rand_core", + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", ] [[package]] @@ -3428,7 +4409,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", ] [[package]] @@ -3440,13 +4430,23 @@ 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 0.9.5", +] + [[package]] name = "rand_xorshift" version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a" dependencies = [ - "rand_core", + "rand_core 0.9.5", ] [[package]] @@ -3482,8 +4482,8 @@ dependencies = [ "num-traits", "paste", "profiling", - "rand", - "rand_chacha", + "rand 0.9.5", + "rand_chacha 0.9.0", "simd_helpers", "thiserror 2.0.21", "v_frame", @@ -3505,6 +4505,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" @@ -3554,6 +4563,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" @@ -3593,6 +4608,16 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19b30a45b0cd0bcca8037f3d0dc3421eaf95327a17cad11964fb8179b4fc4832" +[[package]] +name = "rfc6979" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dd2a808d456c4a54e300a23e9f5a67e122c3024119acbfd73e3bf664491cb2" +dependencies = [ + "hmac", + "subtle", +] + [[package]] name = "rfd" version = "0.17.2" @@ -3626,6 +4651,41 @@ version = "0.8.53" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47b34b781b31e5d73e9fbc8689c70551fd1ade9a19e3e28cfec8580a79290cc4" +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "rsa" +version = "0.9.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8573f03f5883dcaebdfcf4725caa1ecb9c15b2ef50c43a07b816e06799bb12d" +dependencies = [ + "const-oid", + "digest", + "num-bigint-dig", + "num-integer", + "num-traits", + "pkcs1", + "pkcs8", + "rand_core 0.6.4", + "sha2", + "signature", + "spki", + "subtle", + "zeroize", +] + [[package]] name = "rustc-demangle" version = "0.1.28" @@ -3679,6 +4739,80 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "rustls" +version = "0.23.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" +dependencies = [ + "once_cell", + "rustls-pki-types", + "rustls-webpki 0.103.15", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-pki-types" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96" +dependencies = [ + "zeroize", +] + +[[package]] +name = "rustls-rustcrypto" +version = "0.0.2-alpha" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f12052947763ab8515f753315357599e9b0b4dab3b8ba15f30f725fe6d025557" +dependencies = [ + "aead", + "aes-gcm", + "chacha20poly1305", + "crypto-common", + "der", + "digest", + "ecdsa", + "ed25519-dalek", + "hmac", + "p256", + "p384", + "paste", + "pkcs8", + "rand_core 0.6.4", + "rsa", + "rustls", + "rustls-pki-types", + "rustls-webpki 0.102.8", + "sec1", + "sha2", + "signature", + "x25519-dalek", +] + +[[package]] +name = "rustls-webpki" +version = "0.102.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64ca1bc8749bd4cf37b5ce386cc146580777b4e8572c7b97baf22c83f444bee9" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + [[package]] name = "rustversion" version = "1.0.23" @@ -3697,6 +4831,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" @@ -3718,6 +4873,20 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "sec1" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3e97a565f76233a6003f9f5c54be1d9c5bdfa3eccfb189469f11ec4901c47dc" +dependencies = [ + "base16ct", + "der", + "generic-array", + "pkcs8", + "subtle", + "zeroize", +] + [[package]] name = "self_cell" version = "1.3.0" @@ -3730,6 +4899,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" @@ -3793,6 +4968,17 @@ dependencies = [ "serde_core", ] +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "shlex" version = "2.0.1" @@ -3809,6 +4995,16 @@ dependencies = [ "libc", ] +[[package]] +name = "signature" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" +dependencies = [ + "digest", + "rand_core 0.6.4", +] + [[package]] name = "simd-adler32" version = "0.3.10" @@ -3949,6 +5145,12 @@ dependencies = [ "serde", ] +[[package]] +name = "spin" +version = "0.9.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3763264f6b73151db08c50ff20d7d8a0b8796e021cdea7ceedad07b80155fa0e" + [[package]] name = "spirv" version = "0.4.0+sdk-1.4.341.0" @@ -3958,6 +5160,16 @@ dependencies = [ "bitflags 2.13.2", ] +[[package]] +name = "spki" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d" +dependencies = [ + "base64ct", + "der", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -3970,6 +5182,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + [[package]] name = "syn" version = "2.0.119" @@ -3992,6 +5210,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" @@ -4181,6 +5435,18 @@ 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 = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + [[package]] name = "uds_windows" version = "1.2.1" @@ -4192,6 +5458,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" @@ -4222,6 +5523,22 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" +[[package]] +name = "universal-hash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea" +dependencies = [ + "crypto-common", + "subtle", +] + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + [[package]] name = "uuid" version = "1.26.1" @@ -4296,6 +5613,12 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + [[package]] name = "wasip2" version = "1.0.4+wasi-0.2.12" @@ -4516,6 +5839,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki-roots" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "weezl" version = "0.1.12" @@ -5097,6 +6429,17 @@ version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ea6fc2961e4ef194dcbfe56bb845534d0dc8098940c7e5c012a258bfec6701bd" +[[package]] +name = "x25519-dalek" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7e468321c81fb07fa7f4c636c3972b9100f0346e5b6a9f2bd0603a52f7ed277" +dependencies = [ + "curve25519-dalek", + "rand_core 0.6.4", + "zeroize", +] + [[package]] name = "xcursor" version = "0.3.11" @@ -5143,6 +6486,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" @@ -5269,6 +6659,59 @@ 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 = "zeroize" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13084392c5e4bc371903e2935a5eaeed24905a7511356b883835e18a78f6879" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c50655cbb0fe3fc43170059e702f1ce5e19b84cec58dc87b037a09935c2f328" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zip" +version = "7.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c42e33efc22a0650c311c2ef19115ce232583abbe80850bc8b66509ebef02de0" +dependencies = [ + "crc32fast", + "indexmap", + "memchr", + "typed-path", +] + [[package]] name = "zlib-rs" version = "0.6.8" diff --git a/Cargo.toml b/Cargo.toml index b1882f9b..d59f8f45 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,6 +29,7 @@ lightcraft-merge = { path = "crates/merge" } lightcraft-engine = { path = "crates/engine" } lightcraft-ui-egui = { path = "crates/ui-egui" } lightcraft-mcp = { path = "crates/mcp" } +lightcraft-segment = { path = "crates/segment" } # foundation thiserror = "2" @@ -62,6 +63,10 @@ jxl-oxide = { version = "0.12", default-features = false } ravif = { version = "0.13", default-features = false, features = ["threading"] } # no "asm": it needs nasm (non-Rust) # gpu (crates/gpu; native backends only, no GL: everything stays pure Rust) +# SAM 3 inference (lightcraft-segment). 0.9.2: later releases pull `tokenizers` with the +# Oniguruma C library into candle-core, and the product is pure Rust. +candle-core = { version = "=0.9.2", default-features = false } +candle-nn = { version = "=0.9.2", default-features = false } wgpu = { version = "30", default-features = false, features = ["std", "parking_lot", "metal", "vulkan", "dx12", "wgsl"] } pollster = "1" diff --git a/NOTICE b/NOTICE index e6f95a8a..f3f6499f 100644 --- a/NOTICE +++ b/NOTICE @@ -15,6 +15,18 @@ 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 (lightcraft-segment) ports the SAM 3 model code of Hugging Face Transformers +(https://github.com/huggingface/transformers, src/transformers/models/sam3 and sam3_tracker), +Copyright The HuggingFace Team and Meta Platforms, Inc., licensed under the Apache License 2.0. +The port was modified by the LightCraft contributors (2026): translated from Python/PyTorch to +Rust on candle and restructured; each ported file says so. Because of this, crates/segment is +licensed under the Apache License 2.0 only (LICENSE-APACHE), not under the MIT option. + +The SAM 3 weights are not part of LightCraft and are never distributed with it. A user downloads +the facebook/sam3 checkpoint only when they choose to (in the app, or by hand), under Meta's SAM +License (https://github.com/facebookresearch/sam3/blob/main/LICENSE), a separate, non-OSI +licence. See docs/ai-masks.md. + Builds made with craft-fonts (the optional CRAFT_FONTS_DIR build input; all official releases) also embed the fonts listed in https://github.com/storytold/craft-fonts/blob/main/ATTRIBUTION.md (the Japanese fonts BIZ UDPGothic, BIZ UDMincho and Shippori Mincho; the web build embeds only BIZ diff --git a/ROADMAP.md b/ROADMAP.md index fdfbec0c..f2dcb9bc 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -39,7 +39,7 @@ duplicate Local entries #22, black GPU exports on an Intel iGPU #78). | Nikon / Sony / older-Canon raw shooter | ~65% | Camera colour fidelity and coverage (ARW and NEF preview estimates are only a starting point) | | Canon CR3 / Fujifilm / Olympus shooter | ~35% | Their raws open as embedded previews only (CR3: the full-size JPEG with full metadata) | | Lightroom Classic power user | ~45% | Print, Book, Map, publish, tethering | -| Relies on AI (masks, denoise) | ~25% | No segmentation or AI denoise models | +| Relies on AI (masks, denoise) | ~25% | Object / Describe masks via optional SAM 3 (download not yet hosted); no AI denoise | ## Where we're going @@ -57,7 +57,9 @@ Priorities, in order. Each points at tracker rows in [`docs/parity.md`](docs/par (references stay in the local `plan/`), then tune against the numbers. 5. **Lens profiles of our own** (LR-EDIT-OPTICS-PROFILE, P1). 6. **AI model strategy** (maintainer decision): which permissively licensed models (or our own training) for - segmentation masks and denoise; then pure-Rust inference. Unblocks M12 and Enhance. + segmentation masks and denoise; then pure-Rust inference. Unblocks M12 and Enhance. Object / Describe masks now + run Meta's SAM 3 in pure Rust (optional, user-consented download under the non-OSI SAM License; CDN mirrors still + to be set up, docs/ai-masks.md); Subject / Sky / People and Enhance are still open. 7. **Then:** HDR (Q), the Classic output modules (Print first, then Map view, Book, Slideshow), video (R), localisation and accessibility. diff --git a/apps/lightcraft/Cargo.toml b/apps/lightcraft/Cargo.toml index 46acbb73..a29299da 100644 --- a/apps/lightcraft/Cargo.toml +++ b/apps/lightcraft/Cargo.toml @@ -10,7 +10,7 @@ default-run = "lightcraft" workspace = true [dependencies] -lightcraft-engine = { workspace = true } +lightcraft-engine = { workspace = true, features = ["sam"] } lightcraft-sysmem = { workspace = true } lightcraft-ui-egui = { workspace = true } lightcraft-codecs = { workspace = true } diff --git a/apps/lightcraft/src/main.rs b/apps/lightcraft/src/main.rs index bb5714ff..e46fd2ee 100644 --- a/apps/lightcraft/src/main.rs +++ b/apps/lightcraft/src/main.rs @@ -406,10 +406,42 @@ OPTIONS: ENVIRONMENT: LIGHTCRAFT_GPU_BACKEND=dx12|vulkan|metal|auto|off graphics backend (default: DX12 on Windows, Metal on macOS, Vulkan on Linux; off = render on the CPU); else WGPU_BACKEND. LIGHTCRAFT_GPU=0: CPU rendering. + LIGHTCRAFT_SAM3_DIR=DIR the SAM 3 model for Object / Describe masks (default: /models/sam3; + optional: LightCraft offers to download it when first needed) + LIGHTCRAFT_SAM3_MIRRORS=URL,… where to download the SAM 3 model from (base URLs, tried in order) "; +/// Warnings and errors (failed commands, AI mask analysis) on stderr; `LIGHTCRAFT_LOG=info` +/// (or `debug`) shows more. +struct StderrLog(log::LevelFilter); + +impl log::Log for StderrLog { + fn enabled(&self, m: &log::Metadata) -> bool { + m.level() <= self.0 && (m.level() <= log::Level::Warn || m.target().starts_with("lightcraft")) + } + fn log(&self, r: &log::Record) { + if self.enabled(r.metadata()) { + eprintln!("lightcraft: {} {}: {}", r.level(), r.target(), r.args()); + } + } + fn flush(&self) {} +} + +fn install_log() { + let level = match std::env::var("LIGHTCRAFT_LOG").unwrap_or_default().as_str() { + "debug" => log::LevelFilter::Debug, + "info" => log::LevelFilter::Info, + _ => log::LevelFilter::Warn, + }; + static LOGGER: std::sync::OnceLock = std::sync::OnceLock::new(); + if log::set_logger(LOGGER.get_or_init(|| StderrLog(level))).is_ok() { + log::set_max_level(level); + } +} + fn main() -> eframe::Result { lightcraft_engine::guard::install_hook(std::env::temp_dir().join("lightcraft-panics.log")); + install_log(); alloc_release::install(); let mut control_port: Option = std::env::var("LIGHTCRAFT_CONTROL_PORT").ok().and_then(|p| p.parse().ok()); let mut files = Vec::new(); @@ -472,7 +504,13 @@ fn main() -> eframe::Result { "LightCraft", options, Box::new(move |cc| { - let (session, problem) = open_session(in_memory, library_dir, seed_demo && files.is_empty()); + let (mut session, problem) = open_session(in_memory, library_dir, seed_demo && files.is_empty()); + // AI masks: the SAM 3 checkpoint (facebook/sam3) in /models/sam3, or LIGHTCRAFT_SAM3_DIR + // (never required: without it, AI masks offer to download it; see docs/ai-masks.md) + session.segmenter.dir = + std::env::var_os("LIGHTCRAFT_SAM3_DIR").map(std::path::PathBuf::from).or_else(|| config_dir().map(|d| d.join("models").join("sam3"))); + // the user's own download locations, one base URL per line (LIGHTCRAFT_SAM3_MIRRORS too) + session.segmenter.mirrors_file = config_dir().map(|d| d.join("models").join("sam3-mirrors.txt")); let mut app = LightcraftApp::new(session, services()); if let Some(ui) = prefs { app.ui = ui; diff --git a/crates/develop/Cargo.toml b/crates/develop/Cargo.toml index c42aaaa3..6d606148 100644 --- a/crates/develop/Cargo.toml +++ b/crates/develop/Cargo.toml @@ -13,6 +13,7 @@ lightcraft-geom = { workspace = true } lightcraft-color = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +miniz_oxide = { workspace = true } [dev-dependencies] proptest = { workspace = true } diff --git a/crates/develop/src/lib.rs b/crates/develop/src/lib.rs index 72660afd..ff4daedd 100644 --- a/crates/develop/src/lib.rs +++ b/crates/develop/src/lib.rs @@ -8,10 +8,12 @@ pub mod controls; pub mod presets; +pub mod segmask; pub mod settings; pub use controls::{CONTROLS, ControlSpec, Section, Track}; pub use presets::{Preset, SettingsGroup, apply_partial, extract_groups}; +pub use segmask::SegMask; pub use settings::*; use serde_json::Value; diff --git a/crates/develop/src/segmask.rs b/crates/develop/src/segmask.rs new file mode 100644 index 00000000..172d53fc --- /dev/null +++ b/crates/develop/src/segmask.rs @@ -0,0 +1,183 @@ +//! Model segmentations stored with a mask ([`SegMask`]), so rendering never needs the model. + +use serde::{Deserialize, Serialize}; + +/// Logits are stored in steps of 1/`SCALE`, clamped to ±127/`SCALE` (≈ ±15.9: a sigmoid of +/// 1 − 1.2e-7, i.e. fully in or out). +const SCALE: f32 = 8.0; +/// Largest grid a stored segmentation may claim (hostile input caps the allocation). +pub const MAX_SIDE: usize = 1024; + +/// A segmentation computed by a model (SAM 3 in `lightcraft-segment`): `side × side` logits +/// over the uncropped, oriented image stretched to a square (row `y`, column `x` covers +/// normalized `((x + 0.5) / side, (y + 0.5) / side)`), quantized to `i8`, deflated and base64 +/// encoded. Coverage at a point is the sigmoid of the bilinearly sampled logit, so edges stay +/// smooth at any output size. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct SegMask { + pub side: u32, + pub data: String, + /// The part of the image the grid covers, normalized `[x0, y0, x1, y1]` (`None` = all of + /// it). A zoomed-in pass over one object covers just its surroundings, at a higher + /// resolution than a pass over the whole photo. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub rect: Option<[f64; 4]>, +} + +impl SegMask { + /// From `side × side` logits over the whole image. + pub fn from_logits(side: usize, logits: &[f32]) -> SegMask { + let q: Vec = logits.iter().take(side * side).map(|l| ((l * SCALE).round().clamp(-127.0, 127.0) as i8) as u8).collect(); + let z = miniz_oxide::deflate::compress_to_vec(&q, 8); + SegMask { side: side as u32, data: base64_encode(&z), rect: None } + } + + /// From `side × side` logits over the part `rect` of the image. + pub fn from_logits_in(side: usize, logits: &[f32], rect: [f64; 4]) -> SegMask { + SegMask { rect: Some(rect), ..SegMask::from_logits(side, logits) } + } + + /// The covered part, normalized (the whole image without a `rect`); `None` when damaged. + pub fn bounds(&self) -> Option<[f64; 4]> { + let r = self.rect.unwrap_or([0.0, 0.0, 1.0, 1.0]); + (r.iter().all(|v| v.is_finite()) && r[2] > r[0] && r[3] > r[1]).then_some(r) + } + + /// The logits, or `None` when the data is damaged or claims an absurd size. + pub fn logits(&self) -> Option> { + let side = self.side as usize; + if side == 0 || side > MAX_SIDE { + return None; + } + let z = base64_decode(&self.data)?; + let q = miniz_oxide::inflate::decompress_to_vec_with_limit(&z, side * side).ok()?; + (q.len() == side * side).then(|| q.iter().map(|b| f32::from(*b as i8) / SCALE).collect()) + } +} + +/// Most zoomed-in detail patches one AI mask component keeps (the app makes one; each costs a +/// decode per render, so a hostile document can't make rendering explode). +pub const MAX_DETAIL: usize = 4; + +/// Deserialize a component's detail patches, keeping at most [`MAX_DETAIL`] (the rest are +/// skipped without being decoded). +pub(crate) fn de_detail<'de, D: serde::Deserializer<'de>>(d: D) -> Result, D::Error> { + struct Capped; + impl<'de> serde::de::Visitor<'de> for Capped { + type Value = Vec; + fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + f.write_str("a list of segmentations") + } + fn visit_seq>(self, mut seq: A) -> Result, A::Error> { + let mut out = Vec::new(); + while out.len() < MAX_DETAIL { + match seq.next_element::()? { + Some(m) => out.push(m), + None => return Ok(out), + } + } + while seq.next_element::()?.is_some() {} + Ok(out) + } + } + d.deserialize_seq(Capped) +} + +const B64: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; + +fn base64_encode(bytes: &[u8]) -> String { + let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4); + for chunk in bytes.chunks(3) { + let b = [chunk.first().copied().unwrap_or(0), chunk.get(1).copied().unwrap_or(0), chunk.get(2).copied().unwrap_or(0)]; + let n = (u32::from(b[0]) << 16) | (u32::from(b[1]) << 8) | u32::from(b[2]); + for i in 0..4 { + if i <= chunk.len() { + out.push(char::from(B64[((n >> (18 - 6 * i)) & 63) as usize])); + } else { + out.push('='); + } + } + } + out +} + +fn base64_decode(s: &str) -> Option> { + let mut out = Vec::with_capacity(s.len() / 4 * 3); + let (mut acc, mut bits) = (0u32, 0u32); + for c in s.bytes().filter(|c| !c.is_ascii_whitespace()) { + if c == b'=' { + break; + } + let v = B64.iter().position(|b| *b == c)? as u32; + acc = (acc << 6) | v; + bits += 6; + if bits >= 8 { + bits -= 8; + out.push((acc >> bits) as u8); + acc &= (1 << bits) - 1; + } + } + Some(out) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn base64_round_trips() { + for n in 0..10 { + let v: Vec = (0..n).map(|i| (i * 37 + 11) as u8).collect(); + assert_eq!(base64_decode(&base64_encode(&v)).unwrap(), v); + } + assert_eq!(base64_encode(b"Man"), "TWFu"); + assert_eq!(base64_encode(b"Ma"), "TWE="); + assert!(base64_decode("TW!u").is_none()); + } + + #[test] + fn logits_round_trip_quantized() { + let side = 16; + let l: Vec = (0..side * side).map(|i| (i as f32 - 128.0) / 4.0).collect(); + let m = SegMask::from_logits(side, &l); + let back = m.logits().unwrap(); + for (a, b) in l.iter().zip(&back) { + assert!((a.clamp(-127.0 / SCALE, 127.0 / SCALE) - b).abs() <= 0.5 / SCALE + 1e-6); + } + } + + #[test] + fn a_hostile_document_keeps_at_most_max_detail_patches() { + let patch = serde_json::to_value(SegMask::from_logits(4, &[1.0; 16])).unwrap(); + let shape = serde_json::json!({"kind": "object", "hint": [{"x": 0.5, "y": 0.5}], "detail": vec![patch; 10_000]}); + let s: crate::MaskShape = serde_json::from_value(shape).unwrap(); + let crate::MaskShape::Object { detail, .. } = s else { panic!("an Object") }; + assert_eq!(detail.len(), MAX_DETAIL); + // and a whole develop document with them still loads + let mut d = crate::DevelopSettings::default(); + d.masks.push(crate::Mask { + id: 1, + components: vec![crate::MaskComponent { + name: None, + op: crate::MaskOp::Add, + invert: false, + shape: crate::MaskShape::Prompt { text: "x".into(), seg: None, detail: vec![SegMask::from_logits(4, &[1.0; 16]); 50], edge: 0.0 }, + }], + ..Default::default() + }); + let back: crate::DevelopSettings = serde_json::from_str(&serde_json::to_string(&d).unwrap()).unwrap(); + let crate::MaskShape::Prompt { detail, .. } = &back.masks[0].components[0].shape else { panic!("a Prompt") }; + assert_eq!(detail.len(), MAX_DETAIL); + } + + #[test] + fn damaged_data_is_none() { + let mut m = SegMask::from_logits(8, &[1.0; 64]); + m.side = 9; + assert!(m.logits().is_none(), "size mismatch"); + m.side = 100_000; + assert!(m.logits().is_none(), "absurd size"); + assert!(SegMask { side: 8, data: "not base64 !".into(), rect: None }.logits().is_none()); + assert!(SegMask { rect: Some([0.5, 0.0, 0.2, 1.0]), ..SegMask::from_logits(2, &[0.0; 4]) }.bounds().is_none()); + } +} diff --git a/crates/develop/src/settings.rs b/crates/develop/src/settings.rs index eeb6ff04..800f229b 100644 --- a/crates/develop/src/settings.rs +++ b/crates/develop/src/settings.rs @@ -772,8 +772,33 @@ pub enum MaskShape { Subject, Sky, Background, + /// One object picked by clicks (SAM 3 point prompts): `hint` holds the clicks that include, + /// `exclude` the ones that exclude; `seg` the segmentation computed from them. Object { hint: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + exclude: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + seg: Option, + /// Zoomed-in passes over parts of the image (higher resolution than `seg`), used + /// inside their rectangles. + #[serde(default, skip_serializing_if = "Vec::is_empty", deserialize_with = "crate::segmask::de_detail")] + detail: Vec, + /// Edge −100..100: below 0 harder (a steeper transition), above 0 softer (feathered). + #[serde(default, skip_serializing_if = "is_zero")] + edge: f64, + }, + /// Everything a description names ("sky", "the red car": SAM 3 concept prompts); `seg` is + /// the segmentation computed for it. + Prompt { + text: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + seg: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty", deserialize_with = "crate::segmask::de_detail")] + detail: Vec, + /// Edge −100..100, as for `Object`. + #[serde(default, skip_serializing_if = "is_zero")] + edge: f64, }, People { person: u32, diff --git a/crates/engine/Cargo.toml b/crates/engine/Cargo.toml index 5c8e3e46..c2d4e8a8 100644 --- a/crates/engine/Cargo.toml +++ b/crates/engine/Cargo.toml @@ -30,6 +30,11 @@ serde_json = { workspace = true } thiserror = { workspace = true } log = { workspace = true } ab_glyph = { workspace = true } +lightcraft-segment = { workspace = true, optional = true } + +[features] +# AI masks with SAM 3 (Object / Describe). Off for the web and CLI builds. +sam = ["dep:lightcraft-segment"] [dev-dependencies] libc = { workspace = true } diff --git a/crates/engine/src/cmd/masks.rs b/crates/engine/src/cmd/masks.rs index ec3ffac0..1615fb06 100644 --- a/crates/engine/src/cmd/masks.rs +++ b/crates/engine/src/cmd/masks.rs @@ -5,6 +5,7 @@ use lightcraft_geom::Point; use serde_json::{Value, json}; use super::{CommandSpec, bad, bool_or, cmd, f64_or, has_active, ok, point, str_param}; +use crate::segment::{Click, MAX_CLICKS}; use crate::{Result, Session}; fn shape_from(kind: &str, p: &Value, c: &str) -> Result { @@ -38,10 +39,100 @@ fn shape_from(kind: &str, p: &Value, c: &str) -> Result { .unwrap_or_default(); MaskShape::ColorRange { samples, refine: f64_or(p, "refine", 50.0) } } - other => return Err(bad(c, format!("unknown mask kind `{other}` (brush|linear|radial|sky|subject|background|luminanceRange|colorRange)"))), + "object" => MaskShape::Object { hint: points(p, "points"), exclude: points(p, "exclude"), seg: seg_param(p), detail: vec![], edge: 0.0 }, + "prompt" => { + MaskShape::Prompt { text: str_param(p, "text").unwrap_or_default().trim().to_string(), seg: seg_param(p), detail: vec![], edge: 0.0 } + } + other => { + return Err(bad( + c, + format!("unknown mask kind `{other}` (brush|linear|radial|sky|subject|background|luminanceRange|colorRange|object|prompt)"), + )); + } }) } +/// `[[x, y], …]` normalized points under `key`. +fn points(p: &Value, key: &str) -> Vec { + p.get(key) + .and_then(Value::as_array) + .map(|a| a.iter().filter_map(|q| Some(Point::new(q.get(0)?.as_f64()?, q.get(1)?.as_f64()?))).collect()) + .unwrap_or_default() +} + +/// A segmentation passed in (`seg`: a stored or replayed AI mask; the model isn't needed). +fn seg_param(p: &Value) -> Option { + p.get("seg").and_then(|v| serde_json::from_value(v.clone()).ok()) +} + +/// What a new AI shape still needs once it is added. +enum Later { + Nothing, + /// Background mode: these Object clicks are segmented on the worker. + Clicks(Vec, Vec), + /// Background mode: the Describe selection is computed first; the mask is added when found. + Text(String), +} + +/// Compute the segmentation of a new AI shape (Object with clicks, Prompt) — or, in background +/// mode, say what to queue. An Object without clicks stays empty until the photo is clicked +/// (the photo is analyzed meanwhile); a Prompt that matches nothing is an error. Shapes given +/// with their `seg` need no model. +fn resolve_ai(s: &mut Session, shape: &mut MaskShape, c: &str) -> Result { + let id = s.active().ok_or_else(|| bad(c, "no active photo"))?; + let background = s.segmenter.background; + match shape { + MaskShape::Object { seg: Some(_), .. } | MaskShape::Prompt { seg: Some(_), .. } => {} + MaskShape::Object { hint, exclude, .. } if !hint.is_empty() && background => { + s.segmenter.model_dir().map_err(ai_err)?; + return Ok(Later::Clicks(std::mem::take(hint), std::mem::take(exclude))); + } + MaskShape::Object { hint, exclude, seg, .. } if !hint.is_empty() => { + *seg = Some(s.segment_clicks(id, &clicks_of(hint, exclude)).map_err(ai_err)?); + } + MaskShape::Object { .. } => { + // the clicks come later: refuse an Object mask that could never be computed (no + // model) rather than leave an empty one; in the app, analyze the photo meanwhile + s.segmenter.model_dir().map_err(ai_err)?; + if background { + s.segment_prepare(id).map_err(ai_err)?; + } + } + MaskShape::Prompt { text, .. } if background => { + s.segmenter.model_dir().map_err(ai_err)?; + if text.trim().is_empty() { + return Err(ai_err("describe what to select")); + } + return Ok(Later::Text(text.clone())); + } + MaskShape::Prompt { text, seg, .. } => { + let found = s.segment_text(id, text).map_err(ai_err)?; + *seg = Some(found.ok_or_else(|| ai_err(format!("nothing matching “{text}” was found in this photo")))?); + } + _ => {} + } + Ok(Later::Nothing) +} + +/// Queue what a new AI component (component `comp` of mask `mask`) still needs. +fn queue_later(s: &mut Session, later: Later, mask: u32, comp: usize, mut out: Value) -> Result { + if let Later::Clicks(hint, exclude) = later { + let id = s.active().ok_or_else(|| ai_err("no active photo"))?; + s.segment_clicks_later(id, mask, comp, hint, exclude).map_err(ai_err)?; + out["pending"] = json!(true); + } + Ok(out) +} + +/// AI mask failures are shown as they are (not as "invalid parameters"). +fn ai_err(msg: impl Into) -> crate::EngineError { + crate::EngineError::Other(msg.into()) +} + +pub(crate) fn clicks_of(include: &[Point], exclude: &[Point]) -> Vec { + include.iter().map(|p| Click { at: *p, include: true }).chain(exclude.iter().map(|p| Click { at: *p, include: false })).collect() +} + fn masks_edit(s: &mut Session, c: &str, label: &str, f: impl FnOnce(&mut Vec, &mut Option) -> Result<()>) -> Result { let id = s.active().ok_or_else(|| bad(c, "no active photo"))?; let mut d = (*s.develop_of(id).unwrap_or_default()).clone(); @@ -67,14 +158,20 @@ pub fn specs() -> Vec { "Create New Mask", [], None, - "{kind: brush|linear|radial|sky|subject|background|luminanceRange|colorRange, ...shape params (start/end, center/rx/ry/angle/feather, lo/hi…), name?}", + "{kind: brush|linear|radial|sky|subject|background|luminanceRange|colorRange|object|prompt, ...shape params (start/end, center/rx/ry/angle/feather, lo/hi…; object: points/exclude [[x,y],…]; prompt: text; object/prompt: seg? a stored segmentation, no model needed), name?} — AI kinds need the SAM 3 model; in the app a prompt returns {pending} and the mask appears when found", has_active, |s, p| { let kind = str_param(p, "kind").unwrap_or("radial").to_string(); - let shape = shape_from(&kind, p, "mask.add")?; + let mut shape = shape_from(&kind, p, "mask.add")?; + let later = resolve_ai(s, &mut shape, "mask.add")?; let name = str_param(p, "name").map(str::to_string); + if let Later::Text(text) = &later { + let id = s.active().ok_or_else(|| bad("mask.add", "no active photo"))?; + s.segment_text_later(id, None, "add", name, text).map_err(ai_err)?; + return Ok(json!({"pending": true})); + } let next = s.active().and_then(|id| s.develop_of(id)).map(|d| d.next_mask_id()).unwrap_or(1); - masks_edit(s, "mask.add", "Add Mask", |masks, active| { + let out = masks_edit(s, "mask.add", "Add Mask", |masks, active| { masks.push(Mask { id: next, name: name.unwrap_or_else(|| format!("Mask {next}")), @@ -83,7 +180,8 @@ pub fn specs() -> Vec { }); *active = Some(next); Ok(()) - }) + })?; + queue_later(s, later, next, 0, out) } ), cmd!( @@ -95,16 +193,29 @@ pub fn specs() -> Vec { has_active, |s, p| { let kind = str_param(p, "kind").unwrap_or("brush").to_string(); - let shape = shape_from(&kind, p, "mask.addComponent")?; let op: MaskOp = serde_json::from_value(p.get("op").cloned().unwrap_or(json!("add"))).map_err(|e| bad("mask.addComponent", e.to_string()))?; let active = s.active_mask; let mid = mask_id(p, active, "mask.addComponent")?; - masks_edit(s, "mask.addComponent", "Edit Mask", |masks, _| { + let id = s.active().ok_or_else(|| bad("mask.addComponent", "no active photo"))?; + if !s.develop_of(id).is_some_and(|d| d.masks.iter().any(|m| m.id == mid)) { + return Err(bad("mask.addComponent", format!("no mask {mid}"))); + } + let mut shape = shape_from(&kind, p, "mask.addComponent")?; + let later = resolve_ai(s, &mut shape, "mask.addComponent")?; + if let Later::Text(text) = &later { + let op = p.get("op").and_then(Value::as_str).unwrap_or("add"); + s.segment_text_later(id, Some(mid), op, None, text).map_err(ai_err)?; + return Ok(json!({"pending": true})); + } + let mut comp = 0; + let out = masks_edit(s, "mask.addComponent", "Edit Mask", |masks, _| { let i = find(masks, mid, "mask.addComponent")?; masks[i].components.push(MaskComponent { name: None, op, invert: bool_or(p, "invert", false), shape }); + comp = masks[i].components.len() - 1; Ok(()) - }) + })?; + queue_later(s, later, mid, comp, out) } ), cmd!( @@ -213,6 +324,158 @@ pub fn specs() -> Vec { Ok(json!({"samples": out})) } ), + cmd!( + "mask.objectPoint", + "Add Object Click", + [], + None, + "{x, y: normalized image point, exclude?: bool (⌥-click: leave this part out), id?: maskId} — a click on the selected mask's Object selection (its last Object component); SAM 3 re-segments the object → {include, exclude}", + has_active, + |s, p| { + let c = "mask.objectPoint"; + let at = Point::new(super::f64_req(p, "x", c)?, super::f64_req(p, "y", c)?); + if !(0.0..=1.0).contains(&at.x) || !(0.0..=1.0).contains(&at.y) { + return Err(bad(c, "the point is outside the photo")); + } + let exclude = bool_or(p, "exclude", false); + let id = s.active().ok_or_else(|| bad(c, "no active photo"))?; + let mid = mask_id(p, s.active_mask, c)?; + let d = s.develop_of(id).unwrap_or_default(); + let m = d.masks.iter().find(|m| m.id == mid).ok_or_else(|| bad(c, format!("no mask {mid}")))?; + let k = m + .components + .iter() + .rposition(|x| matches!(x.shape, MaskShape::Object { .. })) + .ok_or_else(|| bad(c, "the mask has no Object selection (add one first)"))?; + let Some(MaskShape::Object { hint, exclude: ex, edge, .. }) = m.components.get(k).map(|x| &x.shape) else { + return Err(bad(c, "not an Object component")); + }; + let (mut hint, mut ex, edge) = (hint.clone(), ex.clone(), *edge); + // clicks still on their way to the model count too + if let Some(pending) = s.segmenter.pending_clicks().filter(|q| (q.photo, q.mask, q.comp) == (id, mid, k)) { + (hint, ex) = (pending.hint.clone(), pending.exclude.clone()); + } + if hint.len() + ex.len() >= MAX_CLICKS { + return Err(ai_err(format!("an Object selection takes at most {MAX_CLICKS} clicks; start a new one"))); + } + if exclude { + ex.push(at) + } else { + hint.push(at) + } + if s.segmenter.background && !hint.is_empty() { + let out = json!({"include": hint.len(), "exclude": ex.len(), "pending": true}); + s.segment_clicks_later(id, mid, k, hint, ex).map_err(ai_err)?; + return Ok(out); + } + let seg = if hint.is_empty() { None } else { Some(s.segment_clicks(id, &clicks_of(&hint, &ex)).map_err(ai_err)?) }; + let out = json!({"include": hint.len(), "exclude": ex.len()}); + masks_edit(s, c, "Object Mask", |masks, _| { + 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!( + query "segment.model.status", + "AI Mask Model Status", + [], + None, + "{} — whether this build has AI masks, whether the SAM 3 model is installed (and where), loaded, busy, and the download's progress → {available, installed, dir, loaded, busy, analyzing, sizeBytes, license, licenseUrl, mirrors, download: {running, done, total, file, error, finished}}", + super::always, + |s, _| { + let g = &s.segmenter; + Ok(json!({ + "available": crate::segment::Segmenter::AVAILABLE, + "installed": g.installed(), + "dir": g.dir.as_ref().map(|d| d.display().to_string()), + "loaded": g.loaded(), + "busy": g.busy(), + "analyzing": g.analyzing(), + "sizeBytes": crate::segment::MODEL_BYTES, + "license": crate::segment::LICENSE_NAME, + "licenseUrl": crate::segment::LICENSE_URL, + "mirrors": g.mirrors().len(), + "download": g.download_status(), + })) + } + ), + cmd!( + query "segment.model.download", + "Download AI Mask Model", + [], + None, + "{acknowledged: true} — download the SAM 3 model (about 3.4 GB, Meta's SAM License, not LightCraft's) in the background, from the configured mirrors; only after the user agreed to it. Watch segment.model.status; segment.model.cancel stops it (it resumes later) → {started, installed, downloading}", + super::always, + |s, p| { + let c = "segment.model.download"; + if p.get("acknowledged").and_then(Value::as_bool) != Some(true) { + return Err(bad( + c, + format!( + "the SAM 3 model is a {:.1} GB download under Meta's {} ({}): pass `acknowledged: true` once the user has agreed to download it", + crate::segment::MODEL_BYTES as f64 / 1e9, + crate::segment::LICENSE_NAME, + crate::segment::LICENSE_URL + ), + )); + } + let started = s.segmenter.start_download().map_err(ai_err)?; + Ok(json!({"started": started, "installed": s.segmenter.installed(), "downloading": s.segmenter.download_status().running})) + } + ), + cmd!( + query "segment.model.cancel", + "Cancel AI Mask Model Download", + [], + None, + "{} — stop the SAM 3 download (what has arrived is kept, and a new download resumes from it) → {cancelled}", + super::always, + |s, _| Ok(json!({"cancelled": s.segmenter.cancel_download()})) + ), cmd!( "mask.refine", "Refine Mask Edges", diff --git a/crates/engine/src/lib.rs b/crates/engine/src/lib.rs index 05a9cc90..ec39222a 100644 --- a/crates/engine/src/lib.rs +++ b/crates/engine/src/lib.rs @@ -33,6 +33,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; @@ -154,6 +155,8 @@ pub struct Session { depth: u32, /// Selected mask (Masking panel), by mask id. pub active_mask: Option, + /// AI masks (SAM 3): the model and the last photo prepared for it. + pub segmenter: segment::Segmenter, /// Selected spot (Remove panel), by index into the active photo's spots. pub active_spot: Option, /// The persistent library this session writes to (`None` = in-memory only). @@ -242,6 +245,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(), @@ -723,6 +727,8 @@ mod tests_persist; #[cfg(test)] mod tests_prefs; #[cfg(test)] +mod tests_segment; +#[cfg(test)] mod tests_settings_files; #[cfg(test)] mod tests_spots; diff --git a/crates/engine/src/segment/download.rs b/crates/engine/src/segment/download.rs new file mode 100644 index 00000000..875d652f --- /dev/null +++ b/crates/engine/src/segment/download.rs @@ -0,0 +1,146 @@ +//! The model download as a background task the session starts, watches and cancels without +//! ever blocking: progress lives in atomics, messages behind a mutex only ever `try_lock`ed +//! by the session (the download thread holds it for a few instructions at a time). + +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; + +use lightcraft_segment::fetch::{self, Options, Progress}; + +/// What the session sees of a download. +#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize)] +pub struct Status { + pub running: bool, + pub done: u64, + pub total: u64, + /// The file being fetched. + pub file: String, + /// The last download's error (cleared when a new one starts). + pub error: Option, + /// The last download finished and the files are in place. + pub finished: bool, +} + +#[derive(Default)] +struct Text { + file: String, + error: Option, +} + +#[derive(Default)] +struct Shared { + running: AtomicBool, + cancel: AtomicBool, + finished: AtomicBool, + done: AtomicU64, + total: AtomicU64, + text: Mutex, +} + +/// Clears `running` when the download thread ends, however it ends. +struct Running(Arc); + +impl Drop for Running { + fn drop(&mut self) { + self.0.running.store(false, Ordering::SeqCst); + } +} + +#[derive(Default)] +pub struct Downloader { + shared: Arc, + /// The last text seen (when the download thread held the lock at that moment). + last: Mutex, +} + +fn lock(m: &Mutex) -> std::sync::MutexGuard<'_, Text> { + m.lock().unwrap_or_else(std::sync::PoisonError::into_inner) +} + +impl Downloader { + pub fn status(&self) -> Status { + let s = &self.shared; + if let Ok(t) = s.text.try_lock() + && let Ok(mut last) = self.last.try_lock() + { + last.file.clone_from(&t.file); + last.error.clone_from(&t.error); + } + // (only the session's thread locks `last`) + let (file, error) = self.last.try_lock().map(|l| (l.file.clone(), l.error.clone())).unwrap_or_default(); + Status { + running: s.running.load(Ordering::SeqCst), + done: s.done.load(Ordering::SeqCst), + total: s.total.load(Ordering::SeqCst), + file, + error, + finished: s.finished.load(Ordering::SeqCst), + } + } + + pub fn running(&self) -> bool { + self.shared.running.load(Ordering::SeqCst) + } + + /// Ask a running download to stop (it keeps its partial files to resume later). + pub fn cancel(&self) -> bool { + let running = self.running(); + if running { + self.shared.cancel.store(true, Ordering::SeqCst); + } + running + } + + /// Start downloading `files` from `mirrors` into `dir` on a background thread. False when + /// one is already running. + pub fn start(&self, files: &'static [fetch::FileSpec], mirrors: Vec, dir: PathBuf, opts: Options) -> Result { + let s = self.shared.clone(); + if s.running.swap(true, Ordering::SeqCst) { + return Ok(false); + } + s.cancel.store(false, Ordering::SeqCst); + s.finished.store(false, Ordering::SeqCst); + s.done.store(0, Ordering::SeqCst); + s.total.store(files.iter().filter_map(|f| f.size).sum(), Ordering::SeqCst); + // nothing else runs while `running` was false + *lock(&s.text) = Text::default(); + *lock(&self.last) = Text::default(); + let guard = Running(s.clone()); + let spawned = std::thread::Builder::new().name("sam3-download".into()).spawn(move || { + let s = guard.0.clone(); + let mut last_file = String::new(); + let mut on_progress = |p: &Progress| { + s.done.store(p.done, Ordering::SeqCst); + s.total.store(p.total, Ordering::SeqCst); + if p.file != last_file { + last_file.clone_from(&p.file); + lock(&s.text).file.clone_from(&p.file); + } + }; + let r = + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| fetch::download(files, &mirrors, &dir, &opts, &s.cancel, &mut on_progress))); + let error = match r { + Ok(Ok(())) => { + s.finished.store(true, Ordering::SeqCst); + log::info!("SAM 3 model downloaded to {}", dir.display()); + None + } + Ok(Err(e)) => Some(e.to_string()), + Err(_) => Some("the download failed unexpectedly".to_string()), + }; + if let Some(e) = &error { + log::warn!("SAM 3 download: {e}"); + } + lock(&s.text).error = error; + drop(guard); + }); + match spawned { + Ok(_) => Ok(true), + Err(e) => { + // the closure (and its guard) was dropped: `running` is false again + Err(format!("could not start the download: {e}")) + } + } + } +} diff --git a/crates/engine/src/segment/mod.rs b/crates/engine/src/segment/mod.rs new file mode 100644 index 00000000..84397042 --- /dev/null +++ b/crates/engine/src/segment/mod.rs @@ -0,0 +1,715 @@ +//! AI masks: Object (clicks) and Describe (text) selections computed with SAM 3 +//! (`lightcraft-segment`, cargo feature `sam`; the desktop app enables it). +//! +//! **Nothing requires the model.** It is not part of LightCraft (SAM License); the user +//! downloads it when they first use an AI mask and agree to (`segment.model.download`), or +//! puts the files in the model folder themselves. Without it, AI mask requests fail with a +//! clear "not installed" error; renders, exports and every other feature never touch it +//! (stored segmentations are plain data in the develop settings). +//! +//! The model sees the photo as developed, uncropped and without masks, at 1008 px — the frame +//! mask coordinates live in — so a stored segmentation lines up with the photo at any size. +//! Everything that touches the model (loading it, rendering and encoding the photo, clicks, +//! descriptions, detail passes) runs on one worker thread ([`worker`]); the session talks to it +//! over channels and never waits on a lock. With [`Segmenter::background`] (the desktop app) +//! commands return at once and [`Session::segment_poll`] applies the results; without it (CLI, +//! MCP, tests) they wait for the worker. + +#[cfg(feature = "sam")] +mod download; +#[cfg(feature = "sam")] +mod worker; + +use std::path::PathBuf; + +use lightcraft_catalog::PhotoId; +use lightcraft_develop::{MaskShape, SegMask}; +use lightcraft_geom::Point; + +use crate::Session; + +/// Long edge of the image the model sees. +pub const INPUT_EDGE: usize = 1008; +/// Download size of the model (`model.safetensors`; the tokenizer files add ~2 MB). +pub const MODEL_BYTES: u64 = 3_439_938_512; +/// The model's licence (not LightCraft's): shown before downloading. +pub const LICENSE_NAME: &str = "SAM License (Meta)"; +pub const LICENSE_URL: &str = "https://github.com/facebookresearch/sam3/blob/main/LICENSE"; +/// How errors about a missing model start (the UI offers the download on it). +pub const NOT_INSTALLED: &str = "The SAM 3 model is not installed"; +/// Longest a waiting command (CLI, MCP) waits for the model (loading on a CPU can take minutes). +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +const WAIT: std::time::Duration = std::time::Duration::from_secs(15 * 60); +/// Most clicks one Object selection takes (the model's prompt grows with each). +pub const MAX_CLICKS: usize = 64; + +/// One click of an Object selection (normalized image coordinates). +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Click { + pub at: Point, + pub include: bool, +} + +/// What a request's result is for (how [`Session::segment_poll`] applies it). +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +pub(crate) enum Tag { + /// A command waiting for the result itself. + Wait, + /// Load the model and encode the photo ahead of the first click. + Prepare, + /// Object clicks on component `comp` of mask `mask`. + Clicks { photo: PhotoId, mask: u32, comp: usize, hint: Vec, exclude: Vec }, + /// A Describe selection: a new mask (`mask: None`) or a component combined by `op`. + NewPrompt { photo: PhotoId, mask: Option, op: String, name: Option, text: String }, + /// A zoomed-in pass for a component, computed from `shape` (without its detail). + Detail { photo: PhotoId, mask: u32, comp: usize, shape: MaskShape }, +} + +/// The model input: 8-bit RGB. +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +pub(crate) struct Input { + rgb: Vec, + w: usize, + h: usize, +} + +/// What a detail pass re-runs. +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +pub(crate) enum Prompt { + Clicks(Vec), + Text(String), +} + +/// Object clicks sent to the worker and not applied yet (background mode): further clicks add +/// to these, and the loupe shows them. +#[derive(Clone, Debug, PartialEq)] +pub struct PendingClicks { + pub photo: PhotoId, + pub mask: u32, + pub comp: usize, + pub hint: Vec, + pub exclude: Vec, +} + +/// What [`Session::segment_poll`] did. +#[derive(Debug, Default)] +pub struct Polled { + /// A mask changed. + pub changed: bool, + /// Errors and notices from background work, for the user. + pub messages: Vec, + /// A mask that got a new selection (worth a detail pass). + pub refine: Option, +} + +/// The model download (`segment.model.*`). +#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize)] +pub struct DownloadStatus { + pub running: bool, + pub done: u64, + pub total: u64, + pub file: String, + pub error: Option, + pub finished: bool, +} + +/// The model (on its worker thread), the download, and AI mask requests in flight. +pub struct Segmenter { + /// Where the `facebook/sam3` files live (set by the app; `None`: no AI masks). + pub dir: Option, + /// The user's list of download mirrors (one base URL per line), if any. + pub mirrors_file: Option, + /// Run requests in the background and apply them in [`Session::segment_poll`] (the + /// desktop app); otherwise commands wait for their result (CLI, MCP, tests). + pub background: bool, + pending: Option, + messages: Vec, + #[cfg(feature = "sam")] + worker: worker::Worker, + #[cfg(feature = "sam")] + download: download::Downloader, + #[cfg(feature = "sam")] + results: (std::sync::mpsc::Sender, std::sync::mpsc::Receiver), +} + +// (derivable only without the `sam` feature: the results channel has no Default) +#[allow(clippy::derivable_impls)] +impl Default for Segmenter { + fn default() -> Self { + Segmenter { + dir: None, + mirrors_file: None, + background: false, + pending: None, + messages: Vec::new(), + #[cfg(feature = "sam")] + worker: worker::Worker::default(), + #[cfg(feature = "sam")] + download: download::Downloader::default(), + #[cfg(feature = "sam")] + results: std::sync::mpsc::channel(), + } + } +} + +impl Segmenter { + /// Whether this build can compute AI masks at all. + pub const AVAILABLE: bool = cfg!(feature = "sam"); + + /// Whether the model's files are in place. + pub fn installed(&self) -> bool { + #[cfg(feature = "sam")] + if let Some(dir) = &self.dir { + return lightcraft_segment::is_model_dir(dir); + } + false + } + + /// Whether requests are queued or running (loading, analyzing the photo, a click…). + pub fn busy(&self) -> bool { + #[cfg(feature = "sam")] + return self.worker.pending() > 0; + #[cfg(not(feature = "sam"))] + false + } + + /// Whether the slow part runs now: loading the model or analyzing a photo. + pub fn analyzing(&self) -> bool { + #[cfg(feature = "sam")] + return self.worker.encoding(); + #[cfg(not(feature = "sam"))] + false + } + + /// Whether a zoomed-in detail pass is queued or running. + pub fn detail_busy(&self) -> bool { + #[cfg(feature = "sam")] + return self.worker.detail() > 0; + #[cfg(not(feature = "sam"))] + false + } + + /// Whether the model is in memory (it is unloaded after a while without use). + pub fn loaded(&self) -> bool { + #[cfg(feature = "sam")] + return self.worker.loaded(); + #[cfg(not(feature = "sam"))] + false + } + + /// Clicks waiting for their segmentation. + pub fn pending_clicks(&self) -> Option<&PendingClicks> { + self.pending.as_ref() + } + + /// The model download's state. + pub fn download_status(&self) -> DownloadStatus { + #[cfg(feature = "sam")] + { + let s = self.download.status(); + DownloadStatus { running: s.running, done: s.done, total: s.total, file: s.file, error: s.error, finished: s.finished } + } + #[cfg(not(feature = "sam"))] + DownloadStatus::default() + } + + /// Where the model is downloaded from, in order (the user's mirrors first). + pub fn mirrors(&self) -> Vec { + #[cfg(feature = "sam")] + { + let env = std::env::var(lightcraft_segment::fetch::MIRRORS_ENV).ok(); + lightcraft_segment::fetch::mirrors(env.as_deref(), self.mirrors_file.as_deref()) + } + #[cfg(not(feature = "sam"))] + Vec::new() + } + + /// The model folder, if the model is there; otherwise an error that says what to do. + pub fn model_dir(&self) -> Result { + if !Self::AVAILABLE { + return Err("AI masks are not available in this build".into()); + } + let dir = self.dir.clone().ok_or("no folder is set for the SAM 3 model")?; + if self.installed() { + return Ok(dir); + } + let d = self.download_status(); + if d.running { + let pct = d.done.saturating_mul(100).checked_div(d.total).unwrap_or(0); + return Err(format!("{NOT_INSTALLED} yet: it is downloading ({pct} %).")); + } + Err(format!( + "{NOT_INSTALLED}. Download it (about {:.1} GB, {LICENSE_NAME}) when LightCraft offers it, with `segment.model.download {{\"acknowledged\": true}}`, or put model.safetensors, vocab.json and merges.txt in {}.", + MODEL_BYTES as f64 / 1e9, + dir.display() + )) + } + + /// Start downloading the model (on a background thread; see `segment.model.download`). + /// `Ok(false)` when it is installed or downloading already. + pub fn start_download(&self) -> Result { + if !Self::AVAILABLE { + return Err("AI masks are not available in this build".into()); + } + let dir = self.dir.clone().ok_or("no folder is set for the SAM 3 model")?; + if self.installed() { + return Ok(false); + } + #[cfg(feature = "sam")] + { + let mirrors = self.mirrors(); + if mirrors.is_empty() { + return Err(lightcraft_segment::fetch::DownloadError::NoMirrors.to_string()); + } + self.download.start(lightcraft_segment::fetch::SAM3_FILES, mirrors, dir, lightcraft_segment::fetch::Options::default()) + } + #[cfg(not(feature = "sam"))] + { + let _ = dir; + Ok(false) + } + } + + /// Stop a running download (its partial files stay, to resume). False when none runs. + pub fn cancel_download(&self) -> bool { + #[cfg(feature = "sam")] + return self.download.cancel(); + #[cfg(not(feature = "sam"))] + false + } +} + +/// Mix `v` into a 64-bit hash. +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +fn mix(h: u64, v: u64) -> u64 { + (h ^ v).wrapping_mul(0x0100_0000_01b3).rotate_left(29) +} + +/// Objects covering more than this much of the photo gain little from a zoomed pass. +const DETAIL_MAX_AREA: f64 = 0.4; + +/// The part of the photo to zoom into for `seg`'s selection: its bounding box with a margin of +/// a quarter of its size, at least 12 % of the long edge, at most 2:1, inside the photo +/// (`aspect` = width / height). `None` when nothing is selected or it is most of the photo. +pub fn detail_region(seg: &SegMask, aspect: f64) -> Option<[f64; 4]> { + let logits = seg.logits()?; + let r = seg.bounds()?; + let side = seg.side as usize; + let (mut x0, mut y0, mut x1, mut y1) = (usize::MAX, usize::MAX, 0, 0); + for (i, l) in logits.iter().enumerate() { + if *l > 0.0 { + let (x, y) = (i % side, i / side); + (x0, y0, x1, y1) = (x0.min(x), y0.min(y), x1.max(x + 1), y1.max(y + 1)); + } + } + if x0 == usize::MAX { + return None; + } + let (rw, rh) = (r[2] - r[0], r[3] - r[1]); + let nx = |c: usize| r[0] + c as f64 / side as f64 * rw; + let ny = |c: usize| r[1] + c as f64 / side as f64 * rh; + let (bx0, by0, bx1, by1) = (nx(x0), ny(y0), nx(x1), ny(y1)); + if (bx1 - bx0) * (by1 - by0) > DETAIL_MAX_AREA { + return None; + } + // in units of the photo's height: x × aspect + let a = if aspect.is_finite() { aspect.clamp(0.1, 10.0) } else { 1.0 }; + let (cx, cy) = ((bx0 + bx1) / 2.0 * a, (by0 + by1) / 2.0); + let size = ((bx1 - bx0) * a).max(by1 - by0); + let m = 0.25 * size; + let mut w = ((bx1 - bx0) * a + 2.0 * m).max(0.12 * a.max(1.0)); + let mut h = ((by1 - by0) + 2.0 * m).max(0.12 * a.max(1.0)); + w = w.max(h / 2.0).min(a); + h = h.max(w / 2.0).min(1.0); + let x = (cx - w / 2.0).clamp(0.0, (a - w).max(0.0)); + let y = (cy - h / 2.0).clamp(0.0, (1.0 - h).max(0.0)); + let out = [x / a, y, (x + w) / a, y + h]; + // a region nearly the whole photo isn't worth a second pass + ((out[2] - out[0]) * (out[3] - out[1]) < 0.8).then_some(out) +} + +/// `shape` without its detail patches (what a detail pass was computed from) +#[cfg_attr(not(feature = "sam"), allow(dead_code))] +fn strip_detail(mut shape: MaskShape) -> MaskShape { + // (nor its edge setting: moving the slider meanwhile keeps the pass) + if let MaskShape::Object { detail, edge, .. } | MaskShape::Prompt { detail, edge, .. } = &mut shape { + detail.clear(); + *edge = 0.0; + } + shape +} + +fn clicks_of(include: &[Point], exclude: &[Point]) -> Vec { + include.iter().map(|p| Click { at: *p, include: true }).chain(exclude.iter().map(|p| Click { at: *p, include: false })).collect() +} + +impl Session { + /// The key and settings of photo `id` as the model sees it: its look without masks, uncropped. + #[cfg_attr(not(feature = "sam"), allow(dead_code))] + fn segment_key(&self, id: PhotoId) -> Option<(u64, lightcraft_develop::DevelopSettings)> { + let p = self.catalog.photo(id)?; + let mut d = (*p.develop).clone(); + d.masks.clear(); + let content = crate::media::content_key(p); + let mut h = mix(0xcbf2_9ce4_8422_2325, d.hash64()); + h = mix(h, id.0); + for b in content.bytes() { + h = mix(h, u64::from(b)); + } + Some((h, d)) + } + + /// Clicks (in normalized photo coordinates) for the model. + #[cfg(feature = "sam")] + fn model_clicks(clicks: &[Click]) -> Vec { + clicks.iter().map(|c| lightcraft_segment::Click { x: c.at.x as f32, y: c.at.y as f32, positive: c.include }).collect() + } + + /// A worker request about photo `id` (the render job for its model input goes along; the + /// worker runs it only when its cache doesn't hold the photo). `region` renders that part + /// of the photo, from a source sharp enough for it (detail passes). + #[cfg(feature = "sam")] + fn segment_job( + &mut self, + id: PhotoId, + kind: worker::Kind, + tag: Tag, + reply: std::sync::mpsc::Sender, + region: Option<[f64; 4]>, + ) -> Result { + let dir = self.segmenter.model_dir()?; + let (key, mut d) = self.segment_key(id).ok_or("no such photo")?; + let p = self.catalog.photo(id).ok_or("no such photo")?.clone(); + let missing = match &p.source { + lightcraft_catalog::Source::File { path } if !std::path::Path::new(path).exists() => Some(path.clone()), + _ => None, + }; + let render = match region { + None => self.preview_job(id, INPUT_EDGE, INPUT_EDGE, false, &d).ok_or("no such photo")?, + Some(region) => { + d.crop.geometry = + lightcraft_geom::CropGeometry { rect: lightcraft_geom::Rect::new(region[0], region[1], region[2], region[3]), angle: 0.0 }; + d.crop.flip_h = false; + d.crop.flip_v = false; + d.geometry.constrain_crop = false; + let mut job = self.preview_job(id, INPUT_EDGE, INPUT_EDGE, true, &d).ok_or("no such photo")?; + let frac = (region[2] - region[0]).max(region[3] - region[1]).max(0.01); + job.level = crate::media::SourceLevel::for_size((INPUT_EDGE as f64 / frac) as usize); + job.source = self.media.source_ref(&p, job.level); + job + } + }; + Ok(worker::Job { dir, key: if region.is_some() { 0 } else { key }, render: Some(render), missing, kind, tag, reply }) + } + + /// Send a request and wait for its result (CLI, MCP, tests; or when not in background mode). + #[cfg(feature = "sam")] + fn segment_wait(&mut self, id: PhotoId, kind: worker::Kind, region: Option<[f64; 4]>) -> Result, String> { + let (tx, rx) = std::sync::mpsc::channel(); + let job = self.segment_job(id, kind, Tag::Wait, tx, region)?; + self.segmenter.worker.submit(job)?; + match rx.recv_timeout(WAIT) { + Ok(o) => o.result, + Err(_) => Err("the AI model took too long; try again".into()), + } + } + + /// Send a request whose result [`Session::segment_poll`] applies. + #[cfg(feature = "sam")] + fn segment_queue(&mut self, id: PhotoId, kind: worker::Kind, tag: Tag, region: Option<[f64; 4]>) -> Result<(), String> { + let reply = self.segmenter.results.0.clone(); + let job = self.segment_job(id, kind, tag, reply, region)?; + self.segmenter.worker.submit(job) + } + + /// Start preparing photo `id` for AI masks (load the model, encode the image), so the first + /// click is fast. In background mode it returns at once. + pub fn segment_prepare(&mut self, id: PhotoId) -> Result<(), String> { + self.segmenter.model_dir()?; + #[cfg(feature = "sam")] + { + if self.segmenter.background { + return self.segment_queue(id, worker::Kind::Prepare, Tag::Prepare, None); + } + self.segment_wait(id, worker::Kind::Prepare, None)?; + } + #[cfg(not(feature = "sam"))] + let _ = id; + Ok(()) + } + + /// Segment the object `clicks` pick on photo `id`, waiting for the model. + pub fn segment_clicks(&mut self, id: PhotoId, clicks: &[Click]) -> Result { + self.segmenter.model_dir()?; + if !clicks.iter().any(|c| c.include) { + return Err("click the object to include it first".into()); + } + #[cfg(feature = "sam")] + { + self.segment_wait(id, worker::Kind::Clicks(Self::model_clicks(clicks)), None)?.ok_or_else(|| "nothing was selected".to_string()) + } + #[cfg(not(feature = "sam"))] + { + let _ = id; + Err("AI masks are not available in this build".into()) + } + } + + /// Segment everything `text` describes on photo `id`, waiting for the model; `Ok(None)` + /// when nothing matches. + pub fn segment_text(&mut self, id: PhotoId, text: &str) -> Result, String> { + self.segmenter.model_dir()?; + let text = text.trim(); + if text.is_empty() { + return Err("describe what to select".into()); + } + #[cfg(feature = "sam")] + { + self.segment_wait(id, worker::Kind::Text(text.to_string()), None) + } + #[cfg(not(feature = "sam"))] + { + let _ = id; + Err("AI masks are not available in this build".into()) + } + } + + /// Background mode: queue Object clicks for component `comp` of mask `mask`; the result + /// (clicks and segmentation, one undo step) is applied by [`Session::segment_poll`]. + pub(crate) fn segment_clicks_later(&mut self, id: PhotoId, mask: u32, comp: usize, hint: Vec, exclude: Vec) -> Result<(), String> { + self.segmenter.model_dir()?; + if hint.is_empty() { + return Err("click the object to include it first".into()); + } + #[cfg(feature = "sam")] + { + let clicks = Self::model_clicks(&clicks_of(&hint, &exclude)); + let tag = Tag::Clicks { photo: id, mask, comp, hint: hint.clone(), exclude: exclude.clone() }; + self.segment_queue(id, worker::Kind::Clicks(clicks), tag, None)?; + self.segmenter.pending = Some(PendingClicks { photo: id, mask, comp, hint, exclude }); + Ok(()) + } + #[cfg(not(feature = "sam"))] + { + let _ = (id, mask, comp, exclude); + Err("AI masks are not available in this build".into()) + } + } + + /// Background mode: queue a Describe selection; the mask (or component, combined by `op` + /// into mask `mask`) is created by [`Session::segment_poll`] when something matches. + pub(crate) fn segment_text_later(&mut self, id: PhotoId, mask: Option, op: &str, name: Option, text: &str) -> Result<(), String> { + self.segmenter.model_dir()?; + let text = text.trim(); + if text.is_empty() { + return Err("describe what to select".into()); + } + #[cfg(feature = "sam")] + { + let tag = Tag::NewPrompt { photo: id, mask, op: op.to_string(), name, text: text.to_string() }; + self.segment_queue(id, worker::Kind::Text(text.to_string()), tag, None) + } + #[cfg(not(feature = "sam"))] + { + let _ = (id, mask, op, name); + Err("AI masks are not available in this build".into()) + } + } + + /// Start a zoomed-in pass for component `comp` of mask `mask` on photo `id` (an Object or + /// Describe selection): the photo around the selection is rendered at a higher resolution + /// and segmented again. In background mode [`Session::segment_poll`] applies the result; + /// otherwise it is applied before this returns. `Ok(false)` when there is nothing to + /// refine (no selection yet, or it is most of the photo) or a pass is running already. + pub fn segment_detail(&mut self, id: PhotoId, mask: u32, comp: usize) -> Result { + self.segmenter.model_dir()?; + let d = self.develop_of(id).ok_or("no such photo")?; + let shape = d.masks.iter().find(|m| m.id == mask).and_then(|m| m.components.get(comp)).map(|c| c.shape.clone()).ok_or("no such mask")?; + let (seg, prompt) = match &shape { + // (a document with absurdly many clicks isn't run through the model) + MaskShape::Object { hint, exclude, seg: Some(seg), .. } if !hint.is_empty() && hint.len() + exclude.len() <= MAX_CLICKS => { + (seg.clone(), Prompt::Clicks(clicks_of(hint, exclude))) + } + MaskShape::Prompt { text, seg: Some(seg), .. } => (seg.clone(), Prompt::Text(text.clone())), + _ => return Ok(false), + }; + let p = self.catalog.photo(id).ok_or("no such photo")?; + let (pw, ph) = (f64::from(p.width.max(1)), f64::from(p.height.max(1))); + let aspect = if d.orientation.swaps_axes() { ph / pw } else { pw / ph }; + let Some(region) = detail_region(&seg, aspect) else { return Ok(false) }; + #[cfg(feature = "sam")] + { + let kind = worker::Kind::Detail { prompt, region }; + let tag = Tag::Detail { photo: id, mask, comp, shape: strip_detail(shape) }; + if self.segmenter.background { + if self.segmenter.detail_busy() { + return Ok(false); + } + self.segment_queue(id, kind, tag, Some(region))?; + } else { + let patch = self.segment_wait(id, kind, Some(region))?; + let outcome = worker::Outcome { tag, result: Ok(patch), superseded: false }; + let mut polled = Polled::default(); + self.apply_outcome(outcome, &mut polled); + } + Ok(true) + } + #[cfg(not(feature = "sam"))] + { + let _ = (prompt, region, mask); + Ok(false) + } + } + + /// Apply finished background requests (to masks that haven't changed since) and collect + /// their messages. Cheap when nothing is pending: call it every frame. + pub fn segment_poll(&mut self) -> Polled { + #[cfg_attr(not(feature = "sam"), allow(unused_mut))] + let mut polled = Polled { messages: std::mem::take(&mut self.segmenter.messages), ..Default::default() }; + #[cfg(feature = "sam")] + { + // (before draining: a worker that is idle now has sent everything) + let idle = !self.segmenter.busy(); + let done: Vec = self.segmenter.results.1.try_iter().collect(); + for o in done { + self.apply_outcome(o, &mut polled); + } + // clicks whose request was lost (the worker restarted) aren't pending forever + if idle { + self.segmenter.pending = None; + } + } + polled + } + + #[cfg(feature = "sam")] + fn apply_outcome(&mut self, o: worker::Outcome, polled: &mut Polled) { + if o.superseded { + return; + } + match o.tag { + Tag::Wait => {} + Tag::Prepare => { + if let Err(e) = o.result { + polled.messages.push(e); + } + } + Tag::Clicks { photo, mask, comp, hint, exclude } => { + let latest = PendingClicks { photo, mask, comp, hint, exclude }; + if self.segmenter.pending.as_ref() != Some(&latest) { + return; // newer clicks are on their way + } + self.segmenter.pending = None; + let seg = match o.result { + Ok(Some(seg)) => seg, + Ok(None) => return, + Err(e) => return polled.messages.push(e), + }; + let Some(d) = self.develop_of(photo) else { return }; + let mut d = (*d).clone(); + let Some(c) = d.masks.iter_mut().find(|m| m.id == mask).and_then(|m| m.components.get_mut(comp)) else { return }; + let MaskShape::Object { edge, .. } = c.shape else { return }; + c.shape = MaskShape::Object { hint: latest.hint, exclude: latest.exclude, seg: Some(seg), detail: vec![], edge }; + match self.set_develop(photo, d, "Object Mask") { + Ok(()) => { + polled.changed = true; + polled.refine = Some(mask); + } + Err(e) => polled.messages.push(e.to_string()), + } + } + Tag::NewPrompt { photo, mask, op, name, text } => { + let seg = match o.result { + Ok(Some(seg)) => seg, + Ok(None) => return polled.messages.push(format!("Nothing matching “{text}” was found in this photo.")), + Err(e) => return polled.messages.push(e), + }; + if self.active() != Some(photo) { + return polled.messages.push(format!("The photo changed before “{text}” was found; try again.")); + } + let Ok(seg) = serde_json::to_value(seg) else { return }; + let mut p = serde_json::json!({"kind": "prompt", "text": text, "seg": seg}); + if let Some(n) = name { + p["name"] = serde_json::json!(n); + } + let r = match mask { + None => self.execute("mask.add", &p), + Some(m) => { + if !self.develop_of(photo).is_some_and(|d| d.masks.iter().any(|x| x.id == m)) { + return polled.messages.push("The mask was removed before the selection was found.".into()); + } + self.active_mask = Some(m); + p["op"] = serde_json::json!(op); + self.execute("mask.addComponent", &p) + } + }; + match r { + Ok(_) => { + polled.changed = true; + polled.refine = self.active_mask; + } + Err(e) => polled.messages.push(e.to_string()), + } + } + Tag::Detail { photo, mask, comp, shape } => { + let patch = match o.result { + Ok(Some(p)) => p, + Ok(None) => return, + Err(e) => return log::warn!("SAM 3 detail pass: {e}"), + }; + let Some(d) = self.develop_of(photo) else { return }; + let mut d = (*d).clone(); + let Some(c) = d.masks.iter_mut().find(|m| m.id == mask).and_then(|m| m.components.get_mut(comp)) else { return }; + if strip_detail(c.shape.clone()) != shape { + return; // clicked again meanwhile: a newer pass will follow + } + match &mut c.shape { + MaskShape::Object { detail, .. } | MaskShape::Prompt { detail, .. } => *detail = vec![patch], + _ => return, + } + match self.set_develop(photo, d, "AI Mask Detail") { + Ok(()) => polled.changed = true, + Err(e) => log::warn!("AI mask detail: {e}"), + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn seg_box(side: usize, x0: usize, y0: usize, x1: usize, y1: usize) -> SegMask { + let l: Vec = + (0..side * side).map(|i| if (x0..x1).contains(&(i % side)) && (y0..y1).contains(&(i / side)) { 8.0 } else { -8.0 }).collect(); + SegMask::from_logits(side, &l) + } + + #[test] + fn detail_region_frames_the_selection_with_a_margin() { + // a small object in the middle of a square photo + let r = detail_region(&seg_box(100, 40, 45, 50, 55), 1.0).unwrap(); + assert!(r[0] < 0.4 && r[2] > 0.5 && r[1] < 0.45 && r[3] > 0.55, "contains it: {r:?}"); + assert!(r[2] - r[0] < 0.25 && r[3] - r[1] < 0.25, "zoomed in: {r:?}"); + assert!(r.iter().all(|v| (0.0..=1.0).contains(v))); + // at the corner: shifted inside the photo, not cut + let r = detail_region(&seg_box(100, 0, 0, 5, 5), 1.5).unwrap(); + assert!(r[0] >= 0.0 && r[1] >= 0.0 && r[0] < 1e-9 && r[1] < 1e-9); + // a thin object stays at most 2:1 (in photo pixels, 1.5:1 photo) + let r = detail_region(&seg_box(100, 10, 50, 90, 52), 1.5).unwrap(); + let (w, h) = ((r[2] - r[0]) * 1.5, r[3] - r[1]); + assert!(w / h <= 2.0 + 1e-9, "{w} × {h}"); + // a hostile aspect ratio stays sane + assert!(detail_region(&seg_box(100, 40, 45, 50, 55), f64::NAN).is_some_and(|r| r.iter().all(|v| v.is_finite()))); + } + + #[test] + fn no_region_for_nothing_or_most_of_the_photo() { + assert!(detail_region(&seg_box(50, 0, 0, 0, 0), 1.0).is_none()); + assert!(detail_region(&seg_box(50, 2, 2, 48, 48), 1.0).is_none()); + let damaged = SegMask { side: 50, data: "?".into(), rect: None }; + assert!(detail_region(&damaged, 1.0).is_none()); + } +} diff --git a/crates/engine/src/segment/worker.rs b/crates/engine/src/segment/worker.rs new file mode 100644 index 00000000..9ded0fbe --- /dev/null +++ b/crates/engine/src/segment/worker.rs @@ -0,0 +1,426 @@ +//! The SAM 3 worker thread: it owns the model and the last encoded photo, and runs every +//! request (load, render the model input, encode, clicks, descriptions, detail passes) in +//! turn. The session only sends jobs and receives outcomes over channels, so no thread ever +//! waits on a lock the model holds, and the UI thread never runs the model. +//! +//! Every job runs under `catch_unwind`: a panic inside candle becomes an error outcome, drops +//! the (possibly inconsistent) model and leaves the worker ready for the next job. The busy +//! counters are decremented by drop guards, so they can't stay stuck. The model is unloaded +//! after [`IDLE_UNLOAD`] without requests (it holds several GB). + +use std::collections::VecDeque; +use std::panic::AssertUnwindSafe; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::mpsc::{self, RecvTimeoutError}; +use std::time::Duration; + +use lightcraft_develop::SegMask; +use lightcraft_segment::{Encoded, Sam3}; + +use super::{Input, Prompt, Tag}; +use crate::media::RenderJob; + +/// Unload the model after this long without a request. +pub const IDLE_UNLOAD: Duration = Duration::from_secs(10 * 60); + +/// What to compute. +pub enum Kind { + /// Load the model and encode the photo (so the first click is fast). + Prepare, + Clicks(Vec), + Text(String), + /// A zoomed-in pass over `region` (normalized) of the photo, rendered by the job's `render`. + Detail { + prompt: Prompt, + region: [f64; 4], + }, + /// Panics (tests of the worker's recovery). + #[cfg(test)] + Panic, +} + +/// One request. +pub struct Job { + /// The model folder. + pub dir: PathBuf, + /// Which photo and look the request is about (the encode cache's key). + pub key: u64, + /// Renders the model input when the cache doesn't hold `key` (always for a detail pass). + pub render: Option, + /// The photo's file, when it is missing on disk (for a clearer error if rendering fails). + pub missing: Option, + pub kind: Kind, + pub tag: Tag, + pub reply: mpsc::Sender, +} + +/// A finished request: the segmentation (`None`: nothing selected / nothing to report). +pub struct Outcome { + pub tag: Tag, + pub result: Result, String>, + /// A newer request for the same selection came in; this one was skipped. + pub superseded: bool, +} + +/// Counters the session reads without locking. +#[derive(Default)] +pub struct Shared { + /// Requests other than detail passes, queued or running. + pub pending: AtomicUsize, + /// Detail passes queued or running. + pub detail: AtomicUsize, + /// The model is in memory. + pub loaded: AtomicBool, + /// Loading the model or encoding a photo right now (the slow parts). + pub encoding: AtomicBool, +} + +/// Decrements a counter when dropped (also when unwinding). +struct Done<'a>(&'a AtomicUsize); + +impl Drop for Done<'_> { + fn drop(&mut self) { + let _ = self.0.try_update(Ordering::SeqCst, Ordering::SeqCst, |n| Some(n.saturating_sub(1))); + } +} + +/// Sets a flag while alive. +struct Flag<'a>(&'a AtomicBool); + +impl<'a> Flag<'a> { + fn raise(f: &'a AtomicBool) -> Self { + f.store(true, Ordering::SeqCst); + Flag(f) + } +} + +impl Drop for Flag<'_> { + fn drop(&mut self) { + self.0.store(false, Ordering::SeqCst); + } +} + +/// The session's handle on the worker thread (started on the first request, restarted if it +/// ever ends). +#[derive(Default)] +pub struct Worker { + tx: Option>, + thread: Option>, + pub shared: Arc, +} + +impl Worker { + fn counter(&self, job: &Job) -> &AtomicUsize { + if matches!(job.kind, Kind::Detail { .. }) { &self.shared.detail } else { &self.shared.pending } + } + + /// Whether the thread has ended (it never should); its counters are then meaningless. + fn dead(&self) -> bool { + self.thread.as_ref().is_some_and(|t| t.is_finished()) + } + + pub fn pending(&self) -> usize { + if self.dead() { 0 } else { self.shared.pending.load(Ordering::SeqCst) } + } + + pub fn detail(&self) -> usize { + if self.dead() { 0 } else { self.shared.detail.load(Ordering::SeqCst) } + } + + pub fn encoding(&self) -> bool { + !self.dead() && self.shared.encoding.load(Ordering::SeqCst) + } + + pub fn loaded(&self) -> bool { + !self.dead() && self.shared.loaded.load(Ordering::SeqCst) + } + + /// Queue `job`; returns at once. + pub fn submit(&mut self, job: Job) -> Result<(), String> { + if self.dead() { + // a thread that ended took its counters' meaning with it + self.tx = None; + self.thread = None; + self.shared = Arc::new(Shared::default()); + } + self.counter(&job).fetch_add(1, Ordering::SeqCst); + let job = match &self.tx { + Some(tx) => match tx.send(job) { + Ok(()) => return Ok(()), + Err(mpsc::SendError(job)) => job, + }, + None => job, + }; + // (re)start the thread + let (tx, rx) = mpsc::channel(); + let shared = self.shared.clone(); + match std::thread::Builder::new().name("sam3".into()).spawn(move || run(&rx, &shared, IDLE_UNLOAD)) { + Ok(t) => { + self.thread = Some(t); + self.tx = Some(tx.clone()); + tx.send(job).map_err(|_| "the AI mask worker stopped".to_string()) + } + Err(e) => { + let _ = self.counter(&job).try_update(Ordering::SeqCst, Ordering::SeqCst, |n| Some(n.saturating_sub(1))); + Err(format!("could not start the AI mask worker: {e}")) + } + } + } +} + +/// The model and the last encoded photo. +#[derive(Default)] +struct State { + model: Option<(PathBuf, Sam3)>, + cache: Option<(u64, Encoded)>, +} + +/// Whether `later` makes `job` pointless (it asks for the same selection again). +fn supersedes(later: &Job, job: &Job) -> bool { + match (&later.tag, &job.tag) { + (Tag::Prepare, Tag::Prepare) => later.key == job.key, + (Tag::Clicks { photo: a, mask: am, comp: ac, .. }, Tag::Clicks { photo: b, mask: bm, comp: bc, .. }) => (a, am, ac) == (b, bm, bc), + (Tag::Detail { photo: a, mask: am, comp: ac, .. }, Tag::Detail { photo: b, mask: bm, comp: bc, .. }) => (a, am, ac) == (b, bm, bc), + _ => false, + } +} + +fn run(rx: &mpsc::Receiver, shared: &Shared, idle: Duration) { + let mut state = State::default(); + let mut queue: VecDeque = VecDeque::new(); + loop { + if queue.is_empty() { + match rx.recv_timeout(idle) { + Ok(j) => queue.push_back(j), + Err(RecvTimeoutError::Timeout) => { + if state.model.is_some() { + state = State::default(); + shared.loaded.store(false, Ordering::SeqCst); + log::info!("SAM 3 unloaded after {} minutes without use", idle.as_secs() / 60); + } + continue; + } + Err(RecvTimeoutError::Disconnected) => return, + } + } + while let Ok(j) = rx.try_recv() { + queue.push_back(j); + } + let Some(mut job) = queue.pop_front() else { continue }; + let is_detail = matches!(job.kind, Kind::Detail { .. }); + let _done = Done(if is_detail { &shared.detail } else { &shared.pending }); + if queue.iter().any(|later| supersedes(later, &job)) { + let _ = job.reply.send(Outcome { tag: job.tag, result: Ok(None), superseded: true }); + continue; + } + let result = match std::panic::catch_unwind(AssertUnwindSafe(|| run_job(&mut state, shared, &mut job))) { + Ok(r) => r, + Err(p) => { + let why = p.downcast_ref::<&str>().map(|s| s.to_string()).or_else(|| p.downcast_ref::().cloned()).unwrap_or_default(); + log::error!("SAM 3 panicked: {why}"); + // don't trust what the model kept: load it again next time + state = State::default(); + Err("the AI model failed unexpectedly (it will be reloaded on the next try)".to_string()) + } + }; + if result.is_err() && !is_detail { + state.cache = None; + } + shared.loaded.store(state.model.is_some(), Ordering::SeqCst); + let _ = job.reply.send(Outcome { tag: job.tag, result, superseded: false }); + } +} + +/// The model for `dir`, loading it when needed. +fn model<'a>(state: &'a mut State, shared: &Shared, dir: &std::path::Path) -> Result<&'a mut Sam3, String> { + if state.model.as_ref().is_none_or(|(d, _)| d != dir) { + state.model = None; + state.cache = None; + let _busy = Flag::raise(&shared.encoding); + let t = web_time::Instant::now(); + let m = Sam3::load(dir).map_err(|e| e.to_string())?; + log::info!("SAM 3 loaded in {:?}", t.elapsed()); + state.model = Some((dir.to_path_buf(), m)); + shared.loaded.store(true, Ordering::SeqCst); + } + state.model.as_mut().map(|(_, m)| m).ok_or_else(|| "the model did not load".to_string()) +} + +/// Render the model input. +fn render(job: &mut Job) -> Result { + let r = job.render.take().ok_or("the photo is not available")?.run(); + let img = match r.rendered { + Ok(r) => r.image, + Err(e) => { + return Err(match &job.missing { + Some(path) => format!( + "This photo's file is missing (moved or deleted), so it can't be analyzed for AI masks: {path}. Library ▸ Find Missing Photos can relink it." + ), + None => format!("couldn't render the photo for AI masks: {e}"), + }); + } + }; + let rgb = img.data.iter().flat_map(|p| [p[0], p[1], p[2]]).collect(); + Ok(Input { rgb, w: img.width, h: img.height }) +} + +/// The model and the encoded photo `job.key` (rendering and encoding it when needed). +fn encoded<'a>(state: &'a mut State, shared: &Shared, job: &mut Job) -> Result<(&'a mut Sam3, &'a mut Encoded), String> { + model(state, shared, &job.dir)?; + if state.cache.as_ref().is_none_or(|(k, _)| *k != job.key) { + state.cache = None; + let input = render(job)?; + let _busy = Flag::raise(&shared.encoding); + let m = state.model.as_mut().map(|(_, m)| m).ok_or("the model did not load")?; + let t = web_time::Instant::now(); + let mut enc = m.encode(&input.rgb, input.w, input.h).map_err(|e| e.to_string())?; + m.prepare_clicks(&mut enc).map_err(|e| e.to_string())?; + log::info!("SAM 3 encoded {}×{} in {:?}", input.w, input.h, t.elapsed()); + state.cache = Some((job.key, enc)); + } + let State { model, cache } = state; + match (model.as_mut(), cache.as_mut()) { + (Some((_, m)), Some((_, e))) => Ok((m, e)), + _ => Err("the model is not ready".into()), + } +} + +/// Every phrase of `text` ("car, road"), merged (the per-pixel maximum); `None` when nothing +/// matches. +fn text_logits(model: &mut Sam3, enc: &mut Encoded, text: &str) -> Result>, String> { + let mut merged: Option> = None; + for phrase in text.split([',', ';']).map(str::trim).filter(|p| !p.is_empty()) { + if let Some(p) = model.segment_text(enc, phrase, 0.5).map_err(|e| e.to_string())? { + merged = Some(match merged { + None => p.logits, + Some(m) => m.iter().zip(&p.logits).map(|(a, b)| a.max(*b)).collect(), + }); + } + } + Ok(merged) +} + +fn run_job(state: &mut State, shared: &Shared, job: &mut Job) -> Result, String> { + let side = lightcraft_segment::MASK_SIDE; + match std::mem::replace(&mut job.kind, Kind::Prepare) { + Kind::Prepare => { + encoded(state, shared, job)?; + Ok(None) + } + Kind::Clicks(clicks) => { + let (m, enc) = encoded(state, shared, job)?; + let p = m.segment_clicks(enc, &clicks).map_err(|e| e.to_string())?; + Ok(Some(SegMask::from_logits(side, &p.logits))) + } + Kind::Text(text) => { + let (m, enc) = encoded(state, shared, job)?; + Ok(text_logits(m, enc, &text)?.map(|l| SegMask::from_logits(side, &l))) + } + #[cfg(test)] + #[allow(clippy::panic)] + Kind::Panic => panic!("a test panic"), + Kind::Detail { prompt, region } => { + let input = render(job)?; + let m = model(state, shared, &job.dir)?; + let t = web_time::Instant::now(); + let mut enc = m.encode(&input.rgb, input.w, input.h).map_err(|e| e.to_string())?; + let (rw, rh) = (region[2] - region[0], region[3] - region[1]); + if !(rw > 0.0 && rh > 0.0) { + return Ok(None); + } + let logits = match prompt { + Prompt::Clicks(clicks) => { + let inside: Vec = clicks + .iter() + .map(|c| ((c.at.x - region[0]) / rw, (c.at.y - region[1]) / rh, c.include)) + .filter(|(x, y, _)| (0.0..=1.0).contains(x) && (0.0..=1.0).contains(y)) + .map(|(x, y, positive)| lightcraft_segment::Click { x: x as f32, y: y as f32, positive }) + .collect(); + if !inside.iter().any(|c| c.positive) { + return Ok(None); + } + Some(m.segment_clicks(&mut enc, &inside).map_err(|e| e.to_string())?.logits) + } + Prompt::Text(text) => text_logits(m, &mut enc, &text)?, + }; + log::info!("SAM 3 detail pass ({}×{}) in {:?}", input.w, input.h, t.elapsed()); + Ok(logits.map(|l| SegMask::from_logits_in(side, &l, region))) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn job(tag: Tag, reply: &mpsc::Sender) -> Job { + // an empty folder: loading fails with an error (no model files) + let dir = std::env::temp_dir().join(format!("lc-sam3-worker-none-{}", std::process::id())); + Job { dir, key: 1, render: None, missing: None, kind: Kind::Prepare, tag, reply: reply.clone() } + } + + #[test] + fn errors_come_back_and_the_counters_return_to_zero() { + let mut w = Worker::default(); + let (tx, rx) = mpsc::channel(); + for _ in 0..3 { + w.submit(job(Tag::Wait, &tx)).unwrap(); + } + for _ in 0..3 { + let o = rx.recv_timeout(Duration::from_secs(30)).unwrap(); + assert!(o.result.unwrap_err().contains("not found")); + } + let t = web_time::Instant::now(); + while w.pending() > 0 && t.elapsed() < Duration::from_secs(5) { + std::thread::yield_now(); + } + assert_eq!(w.pending(), 0); + assert!(!w.encoding() && !w.loaded()); + } + + #[test] + fn a_panic_in_a_job_is_an_error_and_the_worker_carries_on() { + let shared = Shared::default(); + let (tx, rx) = mpsc::channel(); + let (jtx, jrx) = mpsc::channel(); + let mut j = job(Tag::Wait, &tx); + j.kind = Kind::Panic; + jtx.send(j).unwrap(); + jtx.send(job(Tag::Wait, &tx)).unwrap(); + shared.pending.store(2, Ordering::SeqCst); + drop(jtx); + run(&jrx, &shared, Duration::from_secs(60)); + let first = rx.recv().unwrap(); + assert!(first.result.unwrap_err().contains("failed unexpectedly")); + let second = rx.recv().unwrap(); + assert!(second.result.is_err(), "the next job still runs"); + assert_eq!(shared.pending.load(Ordering::SeqCst), 0, "counters reset"); + assert!(!shared.encoding.load(Ordering::SeqCst)); + } + + #[test] + fn a_newer_request_for_the_same_selection_supersedes_a_queued_one() { + let (tx, rx) = mpsc::channel(); + let (jtx, jrx) = mpsc::channel(); + let clicks = |n: usize| Tag::Clicks { + photo: lightcraft_catalog::PhotoId(1), + mask: 1, + comp: 0, + hint: vec![lightcraft_geom::Point::new(0.5, 0.5); n], + exclude: vec![], + }; + for n in 1..=3 { + jtx.send(job(clicks(n), &tx)).unwrap(); + } + drop(jtx); + let shared = Shared::default(); + shared.pending.store(3, Ordering::SeqCst); + run(&jrx, &shared, Duration::from_secs(60)); + let out: Vec = rx.try_iter().collect(); + assert_eq!(out.len(), 3); + assert!(out[0].superseded && out[1].superseded && !out[2].superseded); + assert_eq!(shared.pending.load(Ordering::SeqCst), 0); + } +} diff --git a/crates/engine/src/tests.rs b/crates/engine/src/tests.rs index 2837d672..d005d508 100644 --- a/crates/engine/src/tests.rs +++ b/crates/engine/src/tests.rs @@ -728,6 +728,40 @@ fn curve_reset_by_channel_and_whole() { 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()); +} + /// A People card's close-up is a square (in pixels) inside the photo, around the face, whatever the /// face rectangle says: corners, oversized and degenerate rectangles never leave the frame or panic. #[test] diff --git a/crates/engine/src/tests_segment.rs b/crates/engine/src/tests_segment.rs new file mode 100644 index 00000000..59aee18c --- /dev/null +++ b/crates/engine/src/tests_segment.rs @@ -0,0 +1,232 @@ +//! AI masks without the SAM 3 model (it is never required): every request is a clear error that +//! comes back at once, nothing waits on the network or the model, a damaged model is an error +//! (not a crash or a stuck "busy"), and the download command needs consent and fails cleanly. +//! No test touches the internet: downloads go to a local server. + +use std::time::{Duration, Instant}; + +use serde_json::json; + +use crate::Session; +use crate::segment::{NOT_INSTALLED, Segmenter}; + +fn tmp(name: &str) -> std::path::PathBuf { + let d = std::env::temp_dir().join(format!("lc-sam3-engine-{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&d); + d +} + +/// A folder that looks installed but holds no usable model. +fn damaged_model(name: &str) -> std::path::PathBuf { + let d = tmp(name); + std::fs::create_dir_all(&d).unwrap(); + std::fs::write(d.join("model.safetensors"), b"\x10\x00\x00\x00\x00\x00\x00\x00{not json at all}").unwrap(); + std::fs::write(d.join("vocab.json"), b"{}").unwrap(); + std::fs::write(d.join("merges.txt"), b"#version: 0.2\n").unwrap(); + d +} + +/// Poll until the worker is idle (or `limit` passes), collecting messages. +fn settle(s: &mut Session, limit: Duration) -> Vec { + let t = Instant::now(); + let mut messages = Vec::new(); + loop { + let p = s.segment_poll(); + messages.extend(p.messages); + if !s.segmenter.busy() && !s.segmenter.detail_busy() { + messages.extend(s.segment_poll().messages); + return messages; + } + assert!(t.elapsed() < limit, "the worker never finished"); + std::thread::sleep(Duration::from_millis(10)); + } +} + +#[test] +fn status_and_model_commands_without_the_model() { + let mut s = Session::with_demo(); + let st = s.execute("segment.model.status", &json!({})).unwrap(); + assert_eq!(st["available"], Segmenter::AVAILABLE); + assert_eq!(st["installed"], false); + assert_eq!(st["download"]["running"], false); + assert!(st["license"].as_str().unwrap().contains("SAM License")); + // downloading needs the user's consent, and names the size and the licence + let e = s.execute("segment.model.download", &json!({})).unwrap_err().to_string(); + assert!(e.contains("acknowledged") && e.contains("SAM License") && e.contains("GB"), "{e}"); + let e = s.execute("segment.model.download", &json!({"acknowledged": "yes"})).unwrap_err().to_string(); + assert!(e.contains("acknowledged"), "{e}"); + assert_eq!(s.execute("segment.model.cancel", &json!({})).unwrap()["cancelled"], false); + if !Segmenter::AVAILABLE { + assert!(s.execute("segment.model.download", &json!({"acknowledged": true})).is_err()); + return; + } + // no folder, then no download location: clear errors, nothing started + let e = s.execute("segment.model.download", &json!({"acknowledged": true})).unwrap_err().to_string(); + assert!(e.contains("no folder"), "{e}"); + let dir = tmp("nomirrors"); + s.segmenter.dir = Some(dir.clone()); + s.segmenter.mirrors_file = Some(dir.join("none.txt")); + if s.segmenter.mirrors().is_empty() { + let e = s.execute("segment.model.download", &json!({"acknowledged": true})).unwrap_err().to_string(); + assert!(e.contains("LIGHTCRAFT_SAM3_MIRRORS"), "{e}"); + } + assert!(!s.segmenter.download_status().running); +} + +#[cfg(feature = "sam")] +#[test] +fn a_failing_download_ends_with_an_error_and_never_blocks() { + use std::io::{BufRead, BufReader, Write}; + // a local mirror that has nothing + let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}/sam3", l.local_addr().unwrap()); + std::thread::spawn(move || { + for mut c in l.incoming().flatten() { + let mut r = BufReader::new(c.try_clone().unwrap()); + let mut line = String::new(); + while r.read_line(&mut line).unwrap_or(0) > 2 { + line.clear(); + } + let _ = c.write_all(b"HTTP/1.1 404 Not Found\r\nContent-Length: 13\r\nConnection: close\r\n\r\n

404

\n"); + } + }); + let dir = tmp("download404"); + let mirrors = dir.with_extension("mirrors.txt"); + std::fs::write(&mirrors, format!("# test mirror\n{base}\n")).unwrap(); + let mut s = Session::with_demo(); + s.segmenter.dir = Some(dir.clone()); + s.segmenter.mirrors_file = Some(mirrors.clone()); + let t = Instant::now(); + let r = s.execute("segment.model.download", &json!({"acknowledged": true})).unwrap(); + assert_eq!(r["started"], true); + assert!(t.elapsed() < Duration::from_secs(1), "starting a download returns at once"); + // a second request while it runs starts nothing + let again = s.execute("segment.model.download", &json!({"acknowledged": true})).unwrap(); + assert!(again["started"] == false || again["downloading"] == false); + let t = Instant::now(); + while s.segmenter.download_status().running { + assert!(t.elapsed() < Duration::from_secs(30), "the download never ended"); + // AI mask requests meanwhile: an error at once, never a wait + let q = Instant::now(); + let e = s.execute("mask.add", &json!({"kind": "prompt", "text": "sky"})).unwrap_err().to_string(); + assert!(e.contains(NOT_INSTALLED), "{e}"); + assert!(q.elapsed() < Duration::from_millis(500)); + std::thread::sleep(Duration::from_millis(20)); + } + let st = s.segmenter.download_status(); + assert!(!st.finished); + let e = st.error.unwrap(); + assert!(e.contains("not found"), "{e}"); + assert!(!s.segmenter.installed()); + // nothing was left behind under a real name, and no error page was saved + for f in ["model.safetensors", "vocab.json", "merges.txt", "vocab.json.part"] { + assert!(!dir.join(f).exists(), "{f}"); + } + let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::remove_file(&mirrors); +} + +/// The app (background mode) never waits: with the model missing every request fails at once; +/// with a damaged model the requests are queued, fail on the worker, and come back as messages. +#[test] +fn background_requests_never_wait_and_failures_come_back_as_messages() { + let mut s = Session::with_demo(); + s.segmenter.background = true; + let id = s.active().unwrap(); + s.segmenter.dir = Some(tmp("missing")); + for (cmd, p) in [("mask.add", json!({"kind": "object"})), ("mask.add", json!({"kind": "prompt", "text": "sky"})), ("segment.prepare", json!({}))] + { + let t = Instant::now(); + let e = s.execute(cmd, &p).unwrap_err().to_string(); + assert!(t.elapsed() < Duration::from_millis(500), "{cmd}"); + if Segmenter::AVAILABLE { + assert!(e.starts_with(NOT_INSTALLED), "{cmd}: {e}"); + } else { + assert!(e.contains("not available"), "{cmd}: {e}"); + } + } + assert!(s.develop_of(id).unwrap().masks.is_empty()); + if !Segmenter::AVAILABLE { + return; + } + // a model folder whose weights are damaged + let dir = damaged_model("damaged"); + s.segmenter.dir = Some(dir.clone()); + assert!(s.segmenter.installed()); + let t = Instant::now(); + s.execute("mask.add", &json!({"kind": "object"})).unwrap(); + let r = s.execute("mask.objectPoint", &json!({"x": 0.5, "y": 0.5})).unwrap(); + assert_eq!(r["pending"], true); + assert!(s.segmenter.pending_clicks().is_some_and(|p| p.hint.len() == 1)); + // a second click while the first is on its way builds on it + let r = s.execute("mask.objectPoint", &json!({"x": 0.6, "y": 0.5, "exclude": true})).unwrap(); + assert_eq!((r["include"].as_u64(), r["exclude"].as_u64()), (Some(1), Some(1))); + let r = s.execute("mask.add", &json!({"kind": "prompt", "text": "the sky"})).unwrap(); + assert_eq!(r["pending"], true); + assert!(t.elapsed() < Duration::from_secs(1), "nothing waited for the model: {:?}", t.elapsed()); + let messages = settle(&mut s, Duration::from_secs(60)); + assert!(!messages.is_empty() && messages.iter().all(|m| !m.is_empty()), "{messages:?}"); + assert!(messages.iter().any(|m| m.contains("model.safetensors") || m.contains("SAM 3")), "{messages:?}"); + assert!(s.segmenter.pending_clicks().is_none()); + // the Object mask is there, without a selection; no Describe mask was made from nothing + let d = s.develop_of(id).unwrap(); + assert_eq!(d.masks.len(), 1); + assert!(matches!(&d.masks[0].components[0].shape, lightcraft_develop::MaskShape::Object { seg: None, .. })); + assert!(!s.segmenter.busy() && !s.segmenter.loaded()); + let _ = std::fs::remove_dir_all(&dir); +} + +/// CLI / MCP (waiting mode): a damaged model is an error from the command, quickly. +#[test] +fn waiting_requests_with_a_damaged_model_are_errors() { + if !Segmenter::AVAILABLE { + return; + } + let mut s = Session::with_demo(); + let id = s.active().unwrap(); + let dir = damaged_model("damaged-wait"); + s.segmenter.dir = Some(dir.clone()); + let e = s.execute("mask.add", &json!({"kind": "prompt", "text": "sky"})).unwrap_err().to_string(); + assert!(e.contains("model.safetensors") || e.contains("SAM 3"), "{e}"); + let e = s.execute("mask.add", &json!({"kind": "object", "points": [[0.5, 0.5]]})).unwrap_err().to_string(); + assert!(!e.is_empty()); + assert!(s.develop_of(id).unwrap().masks.is_empty()); + assert!(!s.segmenter.busy()); + // a stored segmentation needs no model at all (replayed or pasted masks) + let seg = serde_json::to_value(lightcraft_develop::SegMask::from_logits(4, &[5.0; 16])).unwrap(); + s.segmenter.dir = None; + s.execute("mask.add", &json!({"kind": "prompt", "text": "sky", "seg": seg})).unwrap(); + assert_eq!(s.develop_of(id).unwrap().masks.len(), 1); + let _ = std::fs::remove_dir_all(&dir); +} + +/// Too many clicks on one Object selection are refused, not sent to the model. +#[test] +fn object_clicks_are_capped() { + if !Segmenter::AVAILABLE { + return; + } + let mut s = Session::with_demo(); + s.segmenter.background = true; + let dir = damaged_model("clicks"); + s.segmenter.dir = Some(dir.clone()); + let id = s.active().unwrap(); + let mut d = (*s.develop_of(id).unwrap()).clone(); + let many = vec![lightcraft_geom::Point::new(0.5, 0.5); crate::segment::MAX_CLICKS]; + d.masks.push(lightcraft_develop::Mask { + id: 1, + components: vec![lightcraft_develop::MaskComponent { + name: None, + op: lightcraft_develop::MaskOp::Add, + invert: false, + shape: lightcraft_develop::MaskShape::Object { hint: many, exclude: vec![], seg: None, detail: vec![], edge: 0.0 }, + }], + ..Default::default() + }); + s.set_develop(id, d, "test").unwrap(); + s.active_mask = Some(1); + let e = s.execute("mask.objectPoint", &json!({"x": 0.5, "y": 0.5})).unwrap_err().to_string(); + assert!(e.contains("at most"), "{e}"); + assert!(!s.segmenter.busy()); + let _ = std::fs::remove_dir_all(&dir); +} diff --git a/crates/pipeline/src/masks.rs b/crates/pipeline/src/masks.rs index 5c6e003c..a71667bc 100644 --- a/crates/pipeline/src/masks.rs +++ b/crates/pipeline/src/masks.rs @@ -3,7 +3,7 @@ //! Shapes are defined in normalized oriented-image coordinates; sizes are in "long-edge units". //! Components combine with Add (max), Subtract (a·(1−c)) and Intersect (a·c). -use lightcraft_develop::{BrushStroke, LocalAdjustments, Mask, MaskOp, MaskShape}; +use lightcraft_develop::{BrushStroke, LocalAdjustments, Mask, MaskOp, MaskShape, SegMask}; use lightcraft_geom::Point; use lightcraft_raster::{Plane, Rgb32f}; @@ -153,6 +153,19 @@ pub fn shape_alpha(shape: &MaskShape, frame: &Frame, w: usize, h: usize, img: &R } smooth_plane(&mut out, 0.01 * frame_px(frame, w)); } + MaskShape::Object { seg, detail, edge, .. } | MaskShape::Prompt { seg, detail, edge, .. } if seg.is_some() || !detail.is_empty() => { + // Edge: below 0 a steeper transition (up to 8× the logits), above 0 feathered (up to + // 2 % of the long edge, so previews and exports match) + let e = (*edge / 100.0).clamp(-1.0, 1.0) as f32; + let gain = if e < 0.0 { 1.0 - 7.0 * e } else { 1.0 }; + sample_seg(seg.as_ref(), detail, gain, frame, w, h, &mut out); + if e > 0.0 { + smooth_plane(&mut out, e * 0.02 * frame_px(frame, w)); + } + } + // not computed yet (no clicks, or a build without the model): nothing selected + MaskShape::Object { hint, .. } if hint.is_empty() => {} + MaskShape::Prompt { .. } => {} MaskShape::Subject | MaskShape::Object { .. } | MaskShape::People { .. } => { // Saliency heuristic: centre-weighted local contrast (replaced by the segmenter in M12). let m = frame.out_to_norm(w, h); @@ -176,6 +189,61 @@ pub fn shape_alpha(shape: &MaskShape, frame: &Frame, w: usize, h: usize, img: &R out } +/// A decoded [`SegMask`]: its logits and the normalized rectangle they cover. +struct SegGrid { + logits: Vec, + side: usize, + r: [f64; 4], +} + +impl SegGrid { + fn new(seg: &SegMask) -> Option { + Some(SegGrid { logits: seg.logits()?, side: seg.side as usize, r: seg.bounds()? }) + } + + fn contains(&self, n: Point) -> bool { + n.x >= self.r[0] && n.x <= self.r[2] && n.y >= self.r[1] && n.y <= self.r[3] + } + + /// The bilinearly sampled logit at normalized image point `n` (cell centres at + /// `(i + 0.5) / side` of the rectangle). + fn logit(&self, n: Point) -> f32 { + let side = self.side; + let at = |x: usize, y: usize| self.logits.get(y.min(side - 1) * side + x.min(side - 1)).copied().unwrap_or(-16.0); + let u = (n.x - self.r[0]) / (self.r[2] - self.r[0]); + let v = (n.y - self.r[1]) / (self.r[3] - self.r[1]); + let (fx, fy) = ((u * side as f64 - 0.5) as f32, (v * side as f64 - 0.5) as f32); + let (fx, fy) = (fx.clamp(0.0, (side - 1) as f32), fy.clamp(0.0, (side - 1) as f32)); + let (x0, y0) = (fx.floor() as usize, fy.floor() as usize); + let (tx, ty) = (fx - x0 as f32, fy - y0 as f32); + let top = at(x0, y0) * (1.0 - tx) + at(x0 + 1, y0) * tx; + let bot = at(x0, y0 + 1) * (1.0 - tx) + at(x0 + 1, y0 + 1) * tx; + top * (1.0 - ty) + bot * ty + } +} + +/// Stored model segmentations at each output pixel: the zoomed-in `detail` passes inside +/// their rectangles (the highest of them where they overlap), the whole-photo `seg` +/// elsewhere; the logit (times `gain`: > 1 is a harder edge) through a sigmoid. Damaged +/// data selects nothing. +fn sample_seg(seg: Option<&SegMask>, detail: &[SegMask], gain: f32, frame: &Frame, w: usize, h: usize, out: &mut Plane) { + let coarse = seg.and_then(SegGrid::new); + // (capped: a hostile document can hold many) + let patches: Vec = detail.iter().take(lightcraft_develop::segmask::MAX_DETAIL).filter_map(SegGrid::new).collect(); + if coarse.is_none() && patches.is_empty() { + return; + } + 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 +424,93 @@ mod tests { assert!((a.get(50, 25) - 0.5).abs() < 0.05); } + #[test] + fn stored_segmentations_render_at_any_size_and_follow_the_crop() { + // the left half selected, on a 16² logit grid + let side = 16; + let logits: Vec = (0..side * side).map(|i| if i % side < side / 2 { 12.0 } else { -12.0 }).collect(); + let seg = SegMask::from_logits(side, &logits); + for shape in [ + MaskShape::Object { hint: vec![Point::new(0.25, 0.5)], exclude: vec![], seg: Some(seg.clone()), detail: vec![], edge: 0.0 }, + MaskShape::Prompt { text: "left".into(), seg: Some(seg.clone()), detail: vec![], edge: 0.0 }, + ] { + for (w, h) in [(40, 20), (400, 200)] { + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + assert!(a.get(1, h / 2) > 0.99 && a.get(w - 2, h / 2) < 0.01, "{w}×{h}"); + if w >= 400 { + let soft = (0..w).filter(|x| (0.05..0.95).contains(&a.get(*x, h / 2))).count(); + assert!((2..30).contains(&soft), "an anti-aliased edge, not a ramp: {soft} px"); + } + } + } + // cropped to the right half: nothing left selected + let mut d = DevelopSettings::default(); + d.crop.geometry.rect = lightcraft_geom::Rect::from_center(Point::new(0.75, 0.5), 0.5, 1.0); + let f = Frame::new(40, 20, &d, true); + let shape = MaskShape::Prompt { text: "left".into(), seg: Some(seg), detail: vec![], edge: 0.0 }; + let a = shape_alpha(&shape, &f, 20, 20, &Rgb32f::new(20, 20), &Plane::new(20, 20), 0.0); + assert!(a.data.iter().all(|v| *v < 0.05)); + } + + #[test] + fn detail_patches_replace_the_coarse_mask_inside_their_rectangle() { + // coarse: nothing; a patch over the right half selects its own left half (x 0.5..0.75) + let coarse = SegMask::from_logits(4, &[-12.0; 16]); + let side = 16; + let l: Vec = (0..side * side).map(|i| if i % side < side / 2 { 12.0 } else { -12.0 }).collect(); + let patch = SegMask::from_logits_in(side, &l, [0.5, 0.0, 1.0, 1.0]); + let shape = MaskShape::Object { hint: vec![Point::new(0.6, 0.5)], exclude: vec![], seg: Some(coarse), detail: vec![patch], edge: 0.0 }; + let (w, h) = (200, 100); + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + assert!(a.get(20, 50) < 0.01, "outside the patch: the coarse mask"); + assert!(a.get(120, 50) > 0.99, "inside the patch, its selected half"); + assert!(a.get(190, 50) < 0.01, "inside the patch, its unselected half"); + } + + #[test] + fn a_hostile_number_of_detail_patches_is_capped() { + // built in memory (a deserialized document is capped already): only the first + // MAX_DETAIL are decoded and sampled, so this needs 4 × 4 MB, not 20 000 × 4 MB + let side = 1024; + let l: Vec = vec![12.0; side * side]; + let patch = SegMask::from_logits_in(side, &l, [0.0, 0.0, 0.5, 1.0]); + let shape = MaskShape::Object { hint: vec![Point::new(0.2, 0.5)], exclude: vec![], seg: None, detail: vec![patch; 20_000], edge: 0.0 }; + let (w, h) = (64, 32); + let t = std::time::Instant::now(); + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + assert!(a.get(5, 16) > 0.99 && a.get(60, 16) < 0.01); + assert!(t.elapsed() < std::time::Duration::from_secs(20), "{:?}", t.elapsed()); + } + + #[test] + fn the_edge_setting_hardens_or_feathers_the_transition() { + let side = 16; + let l: Vec = (0..side * side).map(|i| ((i % side) as f32 - 7.5) * -0.6).collect(); // a gentle ramp + let seg = SegMask::from_logits(side, &l); + let (w, h) = (400, 40); + let soft_px = |edge: f64| { + let shape = MaskShape::Prompt { text: "x".into(), seg: Some(seg.clone()), detail: vec![], edge }; + let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0); + assert!(a.get(2, h / 2) > 0.9 && a.get(w - 3, h / 2) < 0.1, "still the same selection at {edge}"); + (0..w).filter(|x| (0.1..0.9).contains(&a.get(*x, h / 2))).count() + }; + let (hard, plain, soft) = (soft_px(-100.0), soft_px(0.0), soft_px(100.0)); + assert!(hard < plain && plain < soft, "transition widths {hard} < {plain} < {soft}"); + } + + #[test] + fn ai_shapes_without_a_segmentation_select_nothing() { + let f = frame(30, 30); + let (img, l) = (Rgb32f::new(30, 30), Plane::new(30, 30)); + for shape in [ + MaskShape::Object { hint: vec![], exclude: vec![], seg: None, detail: vec![], edge: 0.0 }, + MaskShape::Prompt { text: "sky".into(), seg: None, detail: vec![], edge: 0.0 }, + MaskShape::Prompt { text: "sky".into(), seg: Some(SegMask { side: 4, data: "damaged!".into(), rect: None }), detail: vec![], edge: 0.0 }, + ] { + assert!(shape_alpha(&shape, &f, 30, 30, &img, &l, 0.0).data.iter().all(|v| *v == 0.0)); + } + } + #[test] fn radial_inside_and_invert() { let f = frame(100, 100); diff --git a/crates/segment/Cargo.toml b/crates/segment/Cargo.toml new file mode 100644 index 00000000..bd7b0d66 --- /dev/null +++ b/crates/segment/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "lightcraft-segment" +version.workspace = true +edition.workspace = true +license = "Apache-2.0" +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 } +rustls = { version = "0.23", default-features = false, features = ["std", "tls12"] } +rustls-rustcrypto = "0.0.2-alpha" +webpki-roots = "1" +sha2 = { version = "0.10", default-features = false, features = ["std"] } + +[target.'cfg(target_os = "macos")'.dependencies] +candle-core = { workspace = true, features = ["metal"] } +candle-nn = { workspace = true, features = ["metal"] } diff --git a/crates/segment/examples/sam3_bench.rs b/crates/segment/examples/sam3_bench.rs new file mode 100644 index 00000000..982f482e --- /dev/null +++ b/crates/segment/examples/sam3_bench.rs @@ -0,0 +1,25 @@ +//! Times SAM 3 on this machine (a synthetic gradient image; the first run includes compiling +//! the GPU kernels): `cargo run -p lightcraft-segment --release --example sam3_bench -- `. + +fn main() -> Result<(), Box> { + let mut args = std::env::args().skip(1); + let dir = std::path::PathBuf::from(args.next().ok_or("usage: sam3_bench ")?); + let t = std::time::Instant::now(); + let mut model = lightcraft_segment::Sam3::load(&dir)?; + eprintln!("load: {:?} on {:?}", t.elapsed(), model.device()); + let (w, h) = (1200, 800); + let rgb: Vec = (0..w * h * 3).map(|i| ((i / 3 % w) * 255 / w) as u8).collect(); + for run in 0..3 { + let t = std::time::Instant::now(); + let mut enc = model.encode(&rgb, w, h)?; + lightcraft_segment::sync(&model)?; + let e = t.elapsed(); + let t = std::time::Instant::now(); + let p = model.segment_clicks(&mut enc, &[lightcraft_segment::Click { x: 0.5, y: 0.5, positive: true }])?; + let c = t.elapsed(); + let t = std::time::Instant::now(); + let d = model.segment_text(&mut enc, "sky", 0.5)?; + eprintln!("run {run}: encode {e:?}, click {c:?} (score {:.2}), text {:?} ({})", p.score, t.elapsed(), d.is_some()); + } + Ok(()) +} diff --git a/crates/segment/src/clip.rs b/crates/segment/src/clip.rs new file mode 100644 index 00000000..b69c148a --- /dev/null +++ b/crates/segment/src/clip.rs @@ -0,0 +1,81 @@ +//! The CLIP text encoder (24 layers, 1024 wide, causal) that turns a prompt into the token +//! features SAM 3's detector attends to. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`, CLIP text model), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. + +use candle_core::{Device, Module, Tensor}; +use candle_nn::{Embedding, LayerNorm, VarBuilder}; + +use crate::Result; +use crate::nn::{Attention, Mlp, ln}; +use crate::tokenizer::CONTEXT; + +const HIDDEN: usize = 1024; +const LAYERS: usize = 24; +const HEADS: usize = 16; +const EPS: f64 = 1e-5; + +struct Layer { + ln1: LayerNorm, + attn: Attention, + ln2: LayerNorm, + mlp: Mlp, +} + +pub struct TextEncoder { + tokens: Embedding, + positions: Tensor, + layers: Vec, + final_ln: LayerNorm, + device: Device, +} + +impl TextEncoder { + /// `vb` at `detector_model.text_encoder.text_model`. + pub fn new(vb: VarBuilder) -> Result { + let e = vb.pp("embeddings"); + let mut layers = Vec::with_capacity(LAYERS); + for i in 0..LAYERS { + let l = vb.pp(format!("encoder.layers.{i}")); + layers.push(Layer { + ln1: ln(HIDDEN, EPS, l.pp("layer_norm1"))?, + attn: Attention::named(HIDDEN, HIDDEN, HEADS, "out_proj", l.pp("self_attn"))?, + ln2: ln(HIDDEN, EPS, l.pp("layer_norm2"))?, + mlp: Mlp::new(HIDDEN, 4 * HIDDEN, true, l.pp("mlp"))?, + }); + } + Ok(Self { + tokens: candle_nn::embedding(49408, HIDDEN, e.pp("token_embedding"))?, + positions: e.get((CONTEXT, HIDDEN), "position_embedding.weight")?.unsqueeze(0)?, + layers, + final_ln: ln(HIDDEN, EPS, vb.pp("final_layer_norm"))?, + device: vb.device().clone(), + }) + } + + /// `[1, 32, 1024]` features for token `ids` (`valid` marks real tokens). + pub fn forward(&self, ids: &[u32], valid: &[bool]) -> Result { + let ids = Tensor::from_slice(ids, (1, ids.len()), &self.device)?; + let mut x = self.tokens.forward(&ids)?.broadcast_add(&self.positions)?; + // causal, and padding keys masked + let n = valid.len(); + let mut bias = vec![0f32; n * n]; + for (q, row) in bias.chunks_mut(n).enumerate() { + for (k, v) in row.iter_mut().enumerate() { + if k > q || !valid.get(k).copied().unwrap_or(false) { + *v = f32::MIN; + } + } + } + let bias = Tensor::from_vec(bias, (1, 1, n, n), &self.device)?; + for l in &self.layers { + let h = l.ln1.forward(&x)?; + x = (&x + l.attn.forward(&h, &h, &h, Some(&bias))?)?; + let h = l.ln2.forward(&x)?; + x = (&x + l.mlp.forward(&h)?)?; + } + Ok(self.final_ln.forward(&x)?) + } +} diff --git a/crates/segment/src/detector.rs b/crates/segment/src/detector.rs new file mode 100644 index 00000000..7c4eccbd --- /dev/null +++ b/crates/segment/src/detector.rs @@ -0,0 +1,308 @@ +//! Text prompts (SAM 3's concept detector): the prompt's CLIP features fused into the image +//! features by a DETR encoder, 200 object queries refined by a DETR decoder (with box +//! relative-position bias and a presence token), per-query masks from a pixel decoder. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. + +use candle_core::{D, Device, IndexOp, Module, Tensor}; +use candle_nn::{Conv2d, Conv2dConfig, GroupNorm, LayerNorm, Linear, VarBuilder}; + +use crate::clip::TextEncoder; +use crate::neck::{FPN, Neck}; +use crate::nn::{Attention, DecoderMlp, Mlp, TORCH_LN_EPS, cxcywh_to_xyxy, inverse_sigmoid, linear, ln, padding_bias, sine_grid, to_vec}; +use crate::{Error, Prediction, Result}; + +const HIDDEN: usize = FPN; +const HEADS: usize = 8; +const QUERIES: usize = 200; +const LAYERS: usize = 6; + +/// Detector image features, computed once per image. +pub struct Features { + /// Levels 0..=2 (288, 144, 72 per side). + levels: Vec, + /// Level 2 flattened `[1, 5184, 256]` and its sine position encoding. + flat: Tensor, + pos: Tensor, + side: usize, +} + +/// Everything the detector predicts for one prompt. +pub struct Output { + /// Per-query mask logits `[200, 288, 288]` flattened. + pub masks: Tensor, + /// Per-query logits (before the presence score). + pub logits: Vec, + /// Whether the concept is in the image at all (logit). + pub presence: f32, + /// Boxes (x0, y0, x1, y1), normalized. + pub boxes: Vec<[f32; 4]>, + /// The text features `[1, 32, 256]` (tests). + pub text: Tensor, +} + +struct EncoderLayer { + ln1: LayerNorm, + self_attn: Attention, + ln2: LayerNorm, + cross_attn: Attention, + ln3: LayerNorm, + mlp: Mlp, +} + +struct DecoderLayer { + self_attn: Attention, + self_ln: LayerNorm, + text_attn: Attention, + text_ln: LayerNorm, + vision_attn: Attention, + vision_ln: LayerNorm, + mlp: Mlp, + mlp_ln: LayerNorm, +} + +pub struct Detector { + device: Device, + neck: Neck, + text: TextEncoder, + text_proj: Linear, + encoder: Vec, + decoder: Vec, + out_ln: LayerNorm, + box_head: DecoderMlp, + query_embed: Tensor, + reference_points: Tensor, + presence_token: Tensor, + presence_head: DecoderMlp, + presence_ln: LayerNorm, + ref_point_head: DecoderMlp, + rpb_x: DecoderMlp, + rpb_y: DecoderMlp, + // dot-product scoring + score_text_mlp: DecoderMlp, + score_text_ln: LayerNorm, + score_text_proj: Linear, + score_query_proj: Linear, + // mask decoder + prompt_attn: Attention, + prompt_ln: LayerNorm, + pixel_convs: Vec, + pixel_norms: Vec, + instance_proj: Conv2d, + mask_embedder: Vec, +} + +impl Detector { + /// `vb` at `detector_model`. + pub fn new(vb: VarBuilder) -> Result { + let e = vb.pp("detr_encoder"); + let mut encoder = Vec::with_capacity(LAYERS); + for i in 0..LAYERS { + let l = e.pp(format!("layers.{i}")); + encoder.push(EncoderLayer { + ln1: ln(HIDDEN, TORCH_LN_EPS, l.pp("layer_norm1"))?, + self_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("self_attn"))?, + ln2: ln(HIDDEN, TORCH_LN_EPS, l.pp("layer_norm2"))?, + cross_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("cross_attn"))?, + ln3: ln(HIDDEN, TORCH_LN_EPS, l.pp("layer_norm3"))?, + mlp: Mlp::new(HIDDEN, 2048, false, l.pp("mlp"))?, + }); + } + let d = vb.pp("detr_decoder"); + let mut decoder = Vec::with_capacity(LAYERS); + for i in 0..LAYERS { + let l = d.pp(format!("layers.{i}")); + decoder.push(DecoderLayer { + self_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("self_attn"))?, + self_ln: ln(HIDDEN, TORCH_LN_EPS, l.pp("self_attn_layer_norm"))?, + text_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("text_cross_attn"))?, + text_ln: ln(HIDDEN, TORCH_LN_EPS, l.pp("text_cross_attn_layer_norm"))?, + vision_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("vision_cross_attn"))?, + vision_ln: ln(HIDDEN, TORCH_LN_EPS, l.pp("vision_cross_attn_layer_norm"))?, + mlp: Mlp::new(HIDDEN, 2048, false, l.pp("mlp"))?, + mlp_ln: ln(HIDDEN, TORCH_LN_EPS, l.pp("mlp_layer_norm"))?, + }); + } + let s = vb.pp("dot_product_scoring"); + let m = vb.pp("mask_decoder"); + let p = m.pp("pixel_decoder"); + let conv3 = Conv2dConfig { padding: 1, ..Default::default() }; + Ok(Self { + device: vb.device().clone(), + neck: Neck::new(vb.pp("vision_encoder.neck"))?, + text: TextEncoder::new(vb.pp("text_encoder.text_model"))?, + text_proj: linear(1024, HIDDEN, vb.pp("text_projection"))?, + encoder, + decoder, + out_ln: ln(HIDDEN, TORCH_LN_EPS, d.pp("output_layer_norm"))?, + box_head: DecoderMlp::new(HIDDEN, HIDDEN, 4, 3, d.pp("box_head"))?, + query_embed: d.get((QUERIES, HIDDEN), "query_embed.weight")?.unsqueeze(0)?, + reference_points: d.get((QUERIES, 4), "reference_points.weight")?.unsqueeze(0)?, + presence_token: d.get((1, HIDDEN), "presence_token.weight")?.unsqueeze(0)?, + presence_head: DecoderMlp::new(HIDDEN, HIDDEN, 1, 3, d.pp("presence_head"))?, + presence_ln: ln(HIDDEN, TORCH_LN_EPS, d.pp("presence_layer_norm"))?, + ref_point_head: DecoderMlp::new(2 * HIDDEN, HIDDEN, HIDDEN, 2, d.pp("ref_point_head"))?, + rpb_x: DecoderMlp::new(2, HIDDEN, HEADS, 2, d.pp("box_rpb_embed_x"))?, + rpb_y: DecoderMlp::new(2, HIDDEN, HEADS, 2, d.pp("box_rpb_embed_y"))?, + score_text_mlp: DecoderMlp::new(HIDDEN, 2048, HIDDEN, 2, s.pp("text_mlp"))?, + score_text_ln: ln(HIDDEN, TORCH_LN_EPS, s.pp("text_mlp_out_norm"))?, + score_text_proj: linear(HIDDEN, HIDDEN, s.pp("text_proj"))?, + score_query_proj: linear(HIDDEN, HIDDEN, s.pp("query_proj"))?, + prompt_attn: Attention::new(HIDDEN, HIDDEN, HEADS, m.pp("prompt_cross_attn"))?, + prompt_ln: ln(HIDDEN, TORCH_LN_EPS, m.pp("prompt_cross_attn_norm"))?, + pixel_convs: (0..3).map(|i| Ok(candle_nn::conv2d(HIDDEN, HIDDEN, 3, conv3, p.pp(format!("conv_layers.{i}")))?)).collect::>()?, + pixel_norms: (0..3).map(|i| Ok(candle_nn::group_norm(8, HIDDEN, 1e-5, p.pp(format!("norms.{i}")))?)).collect::>()?, + instance_proj: candle_nn::conv2d(HIDDEN, HIDDEN, 1, Conv2dConfig::default(), m.pp("instance_projection"))?, + mask_embedder: (0..3).map(|i| linear(HIDDEN, HIDDEN, m.pp(format!("mask_embedder.layers.{i}")))).collect::>()?, + }) + } + + /// The detector's feature pyramid for the backbone output `[1, 1024, 72, 72]`. + pub fn features(&self, backbone: &Tensor) -> Result { + let levels = self.neck.forward(backbone, 3)?; + let last = levels.last().ok_or_else(|| Error::Model("detector neck".into()))?; + let (_, _, side, _) = last.dims4()?; + let flat = last.flatten_from(2)?.transpose(1, 2)?.contiguous()?; + let pos = sine_grid(side, side, HIDDEN / 2, &self.device)?.flatten_from(2)?.transpose(1, 2)?.contiguous()?; + Ok(Features { levels, flat, pos, side }) + } + + /// Box relative-position bias `[1, heads, 1 + 200, side²]` for reference boxes (cx, cy, w, h). + fn rpb(&self, boxes: &Tensor, side: usize) -> Result { + let xyxy = cxcywh_to_xyxy(boxes)?.squeeze(0)?; // [200, 4] + let coords = Tensor::arange(0u32, side as u32, &self.device)?.to_dtype(candle_core::DType::F32)?.affine(1.0 / side as f64, 0.0)?; + let coords = coords.reshape((1, side, 1))?; + let log_scale = |d: Tensor| -> Result { + let d = (d * 8.0)?; + Ok(d.sign()?.mul(&((d.abs()? + 1.0)?.log()? / (8f64.ln()))?)?) + }; + let pick = |a: usize, b: usize| -> Result { Ok(Tensor::cat(&[xyxy.narrow(1, a, 1)?, xyxy.narrow(1, b, 1)?], 1)?.unsqueeze(1)?) }; + let dy = log_scale(coords.broadcast_sub(&pick(1, 3)?)?)?; // [200, side, 2] + let dx = log_scale(coords.broadcast_sub(&pick(0, 2)?)?)?; + let ey = self.rpb_y.forward(&dy)?; // [200, side, heads] + let ex = self.rpb_x.forward(&dx)?; + let m = ey.unsqueeze(2)?.broadcast_add(&ex.unsqueeze(1)?)?; // [200, side, side, heads] + let m = m.reshape((QUERIES, side * side, HEADS))?.permute((2, 0, 1))?; // [heads, 200, side²] + let m = m.pad_with_zeros(1, 1, 0)?; // the presence token attends everywhere + Ok(m.unsqueeze(0)?.contiguous()?) + } + + /// Sine embedding of boxes (cx, cy, w, h) `[1, 200, 512]` (y, x, w, h features). + fn box_sine(&self, boxes: &Tensor) -> Result { + let feats = HIDDEN / 2; + let dim_t: Vec = (0..feats).map(|i| 10000f64.powf(2.0 * (i / 2) as f64 / feats as f64) as f32).collect(); + let dim_t = Tensor::from_vec(dim_t, (1, 1, feats), &self.device)?; + let even: Vec = (0..feats).map(|i| if i % 2 == 0 { 1.0 } else { 0.0 }).collect(); + let even = Tensor::from_vec(even, (1, 1, feats), &self.device)?; + let one = |i: usize| -> Result { + let p = (boxes.narrow(D::Minus1, i, 1)? * (2.0 * std::f64::consts::PI))?.broadcast_div(&dim_t)?; + Ok((p.sin()?.broadcast_mul(&even)? + p.cos()?.broadcast_mul(&(1.0 - &even)?)?)?) + }; + Ok(Tensor::cat(&[one(1)?, one(0)?, one(2)?, one(3)?], D::Minus1)?) + } + + /// Run the detector for one tokenized prompt. + pub fn predict(&self, f: &Features, ids: &[u32], valid: &[bool]) -> Result { + let text = self.text_proj.forward(&self.text.forward(ids, valid)?)?; // [1, 32, 256] + let text_bias = padding_bias(valid, &self.device)?; + + // DETR encoder: image tokens attend to themselves and to the prompt + let mut x = f.flat.clone(); + for l in &self.encoder { + let h = l.ln1.forward(&x)?; + let hp = (&h + &f.pos)?; + x = (&x + l.self_attn.forward(&hp, &hp, &h, None)?)?; + let h = l.ln2.forward(&x)?; + x = (&x + l.cross_attn.forward(&h, &text, &text, Some(&text_bias))?)?; + let h = l.ln3.forward(&x)?; + x = (&x + l.mlp.forward(&h)?)?; + } + let memory = x; + + // DETR decoder: presence token + 200 queries with iterative box refinement + let mut boxes = candle_nn::ops::sigmoid(&self.reference_points)?; + let mut h = Tensor::cat(&[self.presence_token.clone(), self.query_embed.clone()], 1)?; + let key_pos = (&memory + &f.pos)?; + let mut queries = None; + let mut presence = None; + for l in &self.decoder { + let qpos = self.ref_point_head.forward(&self.box_sine(&boxes)?)?.pad_with_zeros(1, 1, 0)?; + let bias = self.rpb(&boxes, f.side)?; + let q = (&h + &qpos)?; + h = l.self_ln.forward(&(&h + l.self_attn.forward(&q, &q, &h, None)?)?)?; + let q = (&h + &qpos)?; + h = l.text_ln.forward(&(&h + l.text_attn.forward(&q, &text, &text, Some(&text_bias))?)?)?; + let q = (&h + &qpos)?; + h = l.vision_ln.forward(&(&h + l.vision_attn.forward(&q, &key_pos, &memory, Some(&bias))?)?)?; + h = l.mlp_ln.forward(&(&h + l.mlp.forward(&h)?)?)?; + let qh = self.out_ln.forward(&h.narrow(1, 1, QUERIES)?)?; + boxes = candle_nn::ops::sigmoid(&(self.box_head.forward(&qh)? + inverse_sigmoid(&boxes)?)?)?; + let p = self.presence_head.forward(&self.presence_ln.forward(&h.narrow(1, 0, 1)?)?)?; + presence = Some(p.clamp(-10f32, 10f32)?); + queries = Some(qh); + } + let (Some(queries), Some(presence)) = (queries, presence) else { return Err(Error::Model("decoder has no layers".into())) }; + + // dot-product scoring against the mean of the (refined) prompt tokens + let t = self.score_text_ln.forward(&(self.score_text_mlp.forward(&text)? + &text)?)?; + let w: Vec = valid.iter().map(|v| if *v { 1.0 } else { 0.0 }).collect(); + let count = w.iter().sum::().max(1.0); + let w = Tensor::from_vec(w, (1, valid.len(), 1), &self.device)?; + let pooled = (t.broadcast_mul(&w)?.sum(1)? / f64::from(count))?; // [1, 256] + let pt = self.score_text_proj.forward(&pooled)?; // [1, 256] + let pq = self.score_query_proj.forward(&queries)?; // [1, 200, 256] + let logits = (pq.matmul(&pt.unsqueeze(2)?)? / (HIDDEN as f64).sqrt())?.clamp(-12f32, 12f32)?; + + // masks: the encoder memory attends to the prompt, then a pixel decoder up to 288² + let mem = (&memory + self.prompt_attn.forward(&self.prompt_ln.forward(&memory)?, &text, &text, Some(&text_bias))?)?; + let [l0, l1, _] = &f.levels[..] else { return Err(Error::Model("detector levels".into())) }; + let mut prev = mem.transpose(1, 2)?.reshape((1, HIDDEN, f.side, f.side))?; + for (i, skip) in [l1, l0].into_iter().enumerate() { + let (_, _, sh, sw) = skip.dims4()?; + prev = (prev.upsample_nearest2d(sh, sw)? + skip)?; + let (Some(conv), Some(norm)) = (self.pixel_convs.get(i), self.pixel_norms.get(i)) else { + return Err(Error::Model("pixel decoder".into())); + }; + prev = norm.forward(&conv.forward(&prev)?)?.relu()?; + } + let inst = self.instance_proj.forward(&prev)?; // [1, 256, 288, 288] + let (_, c, mh, mw) = inst.dims4()?; + let mut e = queries.clone(); + for (i, l) in self.mask_embedder.iter().enumerate() { + e = l.forward(&e)?; + if i + 1 < self.mask_embedder.len() { + e = e.relu()?; + } + } + let masks = e.matmul(&inst.reshape((1, c, mh * mw))?)?.reshape((QUERIES, mh, mw))?; + let boxes = cxcywh_to_xyxy(&boxes)?; + let bv = to_vec(&boxes)?; + Ok(Output { + masks, + logits: to_vec(&logits)?, + presence: to_vec(&presence)?.first().copied().unwrap_or(-10.0), + boxes: bv.as_chunks::<4>().0.to_vec(), + text, + }) + } +} + +/// The instances of `out` scoring above `threshold` (sigmoid(logit) × sigmoid(presence)), +/// merged into one mask (the per-pixel maximum of their probabilities, as logits). +pub fn merge(out: &Output, threshold: f32) -> Option { + let s = |x: f32| 1.0 / (1.0 + (-x).exp()); + let presence = s(out.presence); + let keep: Vec<(usize, f32)> = out.logits.iter().enumerate().map(|(i, l)| (i, s(*l) * presence)).filter(|(_, p)| *p > threshold).collect(); + let best = keep.iter().map(|k| k.1).fold(0f32, f32::max); + let mut merged: Option> = None; + for (i, _) in &keep { + let m = out.masks.i(*i).ok().and_then(|t| to_vec(&t).ok())?; + merged = Some(match merged { + None => m, + Some(acc) => acc.iter().zip(&m).map(|(a, b)| a.max(*b)).collect(), + }); + } + merged.map(|logits| Prediction { logits, score: best }) +} diff --git a/crates/segment/src/fetch/http.rs b/crates/segment/src/fetch/http.rs new file mode 100644 index 00000000..332f0147 --- /dev/null +++ b/crates/segment/src/fetch/http.rs @@ -0,0 +1,541 @@ +//! A small HTTP/1.1 GET client for model downloads, in pure Rust: `std::net` sockets and, for +//! `https://`, rustls with the RustCrypto provider (`rustls-rustcrypto`; no `ring` or +//! `aws-lc-rs`, so nothing is compiled from C or assembly) and the Mozilla root certificates +//! (`webpki-roots`). +//! +//! Only what a large-file download needs: one request per connection (`Connection: close`), +//! `Range` requests for resuming, redirects (followed by the caller), `Content-Length`, +//! chunked and read-to-close bodies. Every read waits at most `stall` for data and checks the +//! cancel flag at least twice a second, so a dead server or a cancelled download never hangs +//! the downloading thread. No proxies (see docs/ai-masks.md). + +use std::io::{ErrorKind, Read, Write}; +use std::net::{TcpStream, ToSocketAddrs}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::{Duration, Instant}; + +/// Why a request failed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum HttpError { + /// The URL can't be used (scheme, host, characters). + BadUrl(String), + /// Name resolution or connecting failed (or took longer than the connect timeout). + Connect(String), + /// No data arrived for the stall timeout. + Stalled, + /// The cancel flag was raised. + Cancelled, + /// TLS failed (certificate, handshake). + Tls(String), + /// The server's response could not be understood. + Protocol(String), + /// Any other I/O error. + Io(String), +} + +impl std::fmt::Display for HttpError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + HttpError::BadUrl(e) => write!(f, "unusable URL: {e}"), + HttpError::Connect(e) => write!(f, "could not connect: {e}"), + HttpError::Stalled => write!(f, "the server stopped sending data"), + HttpError::Cancelled => write!(f, "cancelled"), + HttpError::Tls(e) => write!(f, "secure connection failed: {e}"), + HttpError::Protocol(e) => write!(f, "unexpected response: {e}"), + HttpError::Io(e) => write!(f, "{e}"), + } + } +} + +/// A parsed `http://` or `https://` URL. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Url { + pub tls: bool, + pub host: String, + pub port: u16, + /// Path and query, starting with `/`. + pub path: String, +} + +impl Url { + pub fn parse(s: &str) -> Result { + let bad = |why: &str| HttpError::BadUrl(format!("{s}: {why}")); + let s = s.trim(); + if s.chars().any(|c| c.is_control() || c == ' ') { + return Err(bad("contains spaces or control characters")); + } + let (tls, rest) = if let Some(r) = s.strip_prefix("https://") { + (true, r) + } else if let Some(r) = s.strip_prefix("http://") { + (false, r) + } else { + return Err(bad("only http:// and https:// are supported")); + }; + let rest = rest.split('#').next().unwrap_or_default(); + let (authority, path) = match rest.find(['/', '?']) { + Some(i) => (rest.get(..i).unwrap_or_default(), rest.get(i..).unwrap_or_default()), + None => (rest, ""), + }; + let path = if path.starts_with('/') { path.to_string() } else { format!("/{path}") }; + if authority.contains('@') { + return Err(bad("user names and passwords in URLs are not supported")); + } + let (host, port) = if let Some(v6) = authority.strip_prefix('[') { + let end = v6.find(']').ok_or_else(|| bad("unterminated IPv6 address"))?; + let host = v6.get(..end).unwrap_or_default(); + let after = v6.get(end + 1..).unwrap_or_default(); + (host.to_string(), after.strip_prefix(':')) + } else { + match authority.rsplit_once(':') { + Some((h, p)) => (h.to_string(), Some(p)), + None => (authority.to_string(), None), + } + }; + let port = match port { + Some(p) => p.parse::().ok().filter(|p| *p != 0).ok_or_else(|| bad("bad port"))?, + None if tls => 443, + None => 80, + }; + if host.is_empty() || !host.chars().all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '.' | ':' | '_')) { + return Err(bad("bad host name")); + } + Ok(Url { tls, host: host.to_ascii_lowercase(), port, path }) + } + + /// `Location` of a redirect, resolved against this URL. + pub fn join(&self, location: &str) -> Result { + let l = location.trim(); + if l.starts_with("http://") || l.starts_with("https://") { + return Url::parse(l); + } + // another scheme (`ftp:`, `file:`, `data:`…) + if let Some((scheme, _)) = l.split_once(':') + && !scheme.is_empty() + && scheme.chars().all(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '-' | '.')) + { + return Err(HttpError::BadUrl(format!("redirect to an unsupported location `{scheme}:`"))); + } + let scheme = if self.tls { "https" } else { "http" }; + let host = if self.host.contains(':') { format!("[{}]", self.host) } else { self.host.clone() }; + let base = format!("{scheme}://{host}:{}", self.port); + if let Some(rest) = l.strip_prefix("//") { + return Url::parse(&format!("{scheme}://{rest}")); + } + if l.starts_with('/') { + return Url::parse(&format!("{base}{l}")); + } + // relative to the current directory + let dir = self.path.split('?').next().unwrap_or("/"); + let dir = dir.rsplit_once('/').map_or("", |(d, _)| d); + Url::parse(&format!("{base}{dir}/{l}")) + } + + fn host_header(&self) -> String { + let host = if self.host.contains(':') { format!("[{}]", self.host) } else { self.host.clone() }; + let default = if self.tls { 443 } else { 80 }; + if self.port == default { host } else { format!("{host}:{}", self.port) } + } + + /// The URL without its query (for messages: a query may hold a signed token). + pub fn display(&self) -> String { + let scheme = if self.tls { "https" } else { "http" }; + format!("{scheme}://{}{}", self.host_header(), self.path.split('?').next().unwrap_or("/")) + } +} + +/// Timeouts and the cancel flag for one request. +pub struct Limits<'a> { + pub connect: Duration, + pub stall: Duration, + pub cancel: &'a AtomicBool, +} + +/// How often a blocked read wakes up to check the cancel flag. +const TICK: Duration = Duration::from_millis(250); +/// Longest status line + headers accepted. +const MAX_HEAD: usize = 64 * 1024; + +enum Stream { + Plain(TcpStream), + Tls(Box, TcpStream), +} + +fn is_timeout(e: &std::io::Error) -> bool { + matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut | ErrorKind::Interrupted) +} + +fn tls_config() -> Result, HttpError> { + let roots = rustls::RootCertStore { roots: webpki_roots::TLS_SERVER_ROOTS.to_vec() }; + let config = rustls::ClientConfig::builder_with_provider(Arc::new(rustls_rustcrypto::provider())) + .with_safe_default_protocol_versions() + .map_err(|e| HttpError::Tls(e.to_string()))? + .with_root_certificates(roots) + .with_no_client_auth(); + Ok(Arc::new(config)) +} + +impl Stream { + fn connect(url: &Url, limits: &Limits) -> Result { + let addrs: Vec<_> = (url.host.as_str(), url.port).to_socket_addrs().map_err(|e| HttpError::Connect(format!("{}: {e}", url.host)))?.collect(); + if addrs.is_empty() { + return Err(HttpError::Connect(format!("{}: no address", url.host))); + } + let mut last = String::new(); + let mut sock = None; + for a in addrs { + if limits.cancel.load(Ordering::Relaxed) { + return Err(HttpError::Cancelled); + } + match TcpStream::connect_timeout(&a, limits.connect) { + Ok(s) => { + sock = Some(s); + break; + } + Err(e) => last = format!("{}: {e}", url.host), + } + } + let sock = sock.ok_or(HttpError::Connect(last))?; + sock.set_read_timeout(Some(TICK)).map_err(|e| HttpError::Io(e.to_string()))?; + sock.set_write_timeout(Some(TICK)).map_err(|e| HttpError::Io(e.to_string()))?; + let _ = sock.set_nodelay(true); + if !url.tls { + return Ok(Stream::Plain(sock)); + } + let name = rustls::pki_types::ServerName::try_from(url.host.clone()).map_err(|e| HttpError::Tls(format!("{}: {e}", url.host)))?; + let conn = rustls::ClientConnection::new(tls_config()?, name).map_err(|e| HttpError::Tls(e.to_string()))?; + let mut s = Stream::Tls(Box::new(conn), sock); + s.handshake(limits)?; + Ok(s) + } + + /// Wait for `f` to make progress, retrying timeouts until `stall` passes without any. + fn retry(limits: &Limits, mut f: impl FnMut() -> std::io::Result) -> Result { + let start = Instant::now(); + loop { + if limits.cancel.load(Ordering::Relaxed) { + return Err(HttpError::Cancelled); + } + match f() { + Ok(v) => return Ok(v), + Err(e) if is_timeout(&e) => { + if start.elapsed() >= limits.stall { + return Err(HttpError::Stalled); + } + } + Err(e) if e.kind() == ErrorKind::InvalidData => return Err(HttpError::Tls(e.to_string())), + Err(e) => return Err(HttpError::Io(e.to_string())), + } + } + } + + fn handshake(&mut self, limits: &Limits) -> Result<(), HttpError> { + let Stream::Tls(conn, sock) = self else { return Ok(()) }; + while conn.is_handshaking() { + Self::retry(limits, || conn.complete_io(sock))?; + } + Ok(()) + } + + fn write_all(&mut self, data: &[u8], limits: &Limits) -> Result<(), HttpError> { + match self { + Stream::Plain(sock) => { + let mut rest = data; + while !rest.is_empty() { + let n = Self::retry(limits, || sock.write(rest))?; + if n == 0 { + return Err(HttpError::Io("connection closed while sending the request".into())); + } + rest = rest.get(n..).unwrap_or_default(); + } + Ok(()) + } + Stream::Tls(conn, sock) => { + // buffered by rustls, then sent as TLS records + conn.writer().write_all(data).map_err(|e| HttpError::Tls(e.to_string()))?; + while conn.wants_write() { + Self::retry(limits, || conn.write_tls(sock))?; + } + Ok(()) + } + } + } + + /// Read some bytes (0 at the end of the stream), waiting at most `stall` for them. + fn read(&mut self, buf: &mut [u8], limits: &Limits) -> Result { + match self { + Stream::Plain(sock) => Self::retry(limits, || sock.read(buf)), + Stream::Tls(conn, sock) => { + let start = Instant::now(); + loop { + if limits.cancel.load(Ordering::Relaxed) { + return Err(HttpError::Cancelled); + } + match conn.reader().read(buf) { + Ok(n) => return Ok(n), + // closed without close_notify: the caller checks the length it got + Err(e) if e.kind() == ErrorKind::UnexpectedEof => return Ok(0), + Err(e) if e.kind() == ErrorKind::WouldBlock => {} + Err(e) => return Err(HttpError::Tls(e.to_string())), + } + // no plaintext buffered: more TLS records from the socket + match conn.read_tls(sock) { + Ok(0) => return Ok(0), + Ok(_) => { + conn.process_new_packets().map_err(|e| HttpError::Tls(e.to_string()))?; + while conn.wants_write() { + Self::retry(limits, || conn.write_tls(sock))?; + } + } + Err(e) if is_timeout(&e) => { + if start.elapsed() >= limits.stall { + return Err(HttpError::Stalled); + } + } + Err(e) => return Err(HttpError::Io(e.to_string())), + } + } + } + } + } +} + +/// The connection with a read buffer (headers and chunk sizes are read line by line). +struct Conn { + stream: Stream, + buf: Vec, + pos: usize, +} + +impl Conn { + fn fill(&mut self, limits: &Limits) -> Result { + if self.pos >= self.buf.len() { + self.buf.clear(); + self.pos = 0; + } + let mut tmp = [0u8; 64 * 1024]; + let n = self.stream.read(&mut tmp, limits)?; + self.buf.extend_from_slice(tmp.get(..n).unwrap_or_default()); + Ok(n) + } + + /// One line without its CRLF; at most `max` bytes. + fn line(&mut self, max: usize, limits: &Limits) -> Result { + loop { + let avail = self.buf.get(self.pos..).unwrap_or_default(); + if let Some(i) = avail.iter().position(|b| *b == b'\n') { + let line = avail.get(..i).unwrap_or_default(); + let line = line.strip_suffix(b"\r").unwrap_or(line); + let s = String::from_utf8_lossy(line).into_owned(); + self.pos = self.pos.saturating_add(i + 1); + return Ok(s); + } + if avail.len() > max { + return Err(HttpError::Protocol("header line too long".into())); + } + // keep the partial line, read more after it + if self.pos > 0 { + self.buf.drain(..self.pos.min(self.buf.len())); + self.pos = 0; + } + let mut tmp = [0u8; 16 * 1024]; + let n = self.stream.read(&mut tmp, limits)?; + if n == 0 { + return Err(HttpError::Protocol("connection closed in the middle of the response".into())); + } + self.buf.extend_from_slice(tmp.get(..n).unwrap_or_default()); + } + } + + /// Up to `out.len()` body bytes (0 at the end of the stream). + fn read(&mut self, out: &mut [u8], limits: &Limits) -> Result { + if self.pos >= self.buf.len() && self.fill(limits)? == 0 { + return Ok(0); + } + let avail = self.buf.get(self.pos..).unwrap_or_default(); + let n = avail.len().min(out.len()); + if let (Some(dst), Some(src)) = (out.get_mut(..n), avail.get(..n)) { + dst.copy_from_slice(src); + } + self.pos = self.pos.saturating_add(n); + Ok(n) + } +} + +/// A response's status, headers and body. +pub struct Response { + pub status: u16, + headers: Vec<(String, String)>, + conn: Conn, + body: BodyMode, +} + +enum BodyMode { + Length(u64), + Chunked { left: u64, done: bool }, + Close, + Done, +} + +impl Response { + /// The first header named `name` (case-insensitive). + pub fn header(&self, name: &str) -> Option<&str> { + self.headers.iter().find(|(k, _)| k.eq_ignore_ascii_case(name)).map(|(_, v)| v.as_str()) + } + + /// The body's length when the server announced it. + pub fn content_length(&self) -> Option { + match self.body { + BodyMode::Length(n) => Some(n), + _ => None, + } + } + + /// Read body bytes into `out` (0 at the end of the body). + pub fn read(&mut self, out: &mut [u8], limits: &Limits) -> Result { + match &mut self.body { + BodyMode::Done => Ok(0), + BodyMode::Length(left) => { + if *left == 0 { + self.body = BodyMode::Done; + return Ok(0); + } + let want = usize::try_from(*left).unwrap_or(usize::MAX).min(out.len()); + let n = self.conn.read(out.get_mut(..want).unwrap_or_default(), limits)?; + if n == 0 { + return Err(HttpError::Protocol("the connection closed before the end of the file".into())); + } + *left = left.saturating_sub(n as u64); + Ok(n) + } + BodyMode::Close => { + let n = self.conn.read(out, limits)?; + if n == 0 { + self.body = BodyMode::Done; + } + Ok(n) + } + BodyMode::Chunked { left, done } => { + if *done { + return Ok(0); + } + if *left == 0 { + let line = self.conn.line(1024, limits)?; + let hex = line.split(';').next().unwrap_or_default().trim(); + let size = u64::from_str_radix(hex, 16).map_err(|_| HttpError::Protocol(format!("bad chunk size `{hex}`")))?; + if size == 0 { + // trailers, up to the empty line + let mut n = 0; + while !self.conn.line(MAX_HEAD, limits)?.is_empty() { + n += 1; + if n > 100 { + return Err(HttpError::Protocol("too many trailers".into())); + } + } + *done = true; + return Ok(0); + } + *left = size; + } + let want = usize::try_from(*left).unwrap_or(usize::MAX).min(out.len()); + let n = self.conn.read(out.get_mut(..want).unwrap_or_default(), limits)?; + if n == 0 { + return Err(HttpError::Protocol("the connection closed in the middle of a chunk".into())); + } + *left = left.saturating_sub(n as u64); + if *left == 0 && !self.conn.line(16, limits)?.is_empty() { + return Err(HttpError::Protocol("missing CRLF after a chunk".into())); + } + Ok(n) + } + } + } +} + +/// Send `GET url` with extra `headers` and read the status line and headers. +pub fn get(url: &Url, headers: &[(&str, String)], limits: &Limits) -> Result { + let mut req = format!( + "GET {} HTTP/1.1\r\nHost: {}\r\nUser-Agent: LightCraft/{}\r\nAccept: */*\r\nAccept-Encoding: identity\r\nConnection: close\r\n", + url.path, + url.host_header(), + env!("CARGO_PKG_VERSION") + ); + for (k, v) in headers { + if v.chars().any(|c| c == '\r' || c == '\n') || k.chars().any(|c| c == '\r' || c == '\n' || c == ':') { + return Err(HttpError::BadUrl("header with a line break".into())); + } + req.push_str(&format!("{k}: {v}\r\n")); + } + req.push_str("\r\n"); + let mut stream = Stream::connect(url, limits)?; + stream.write_all(req.as_bytes(), limits)?; + let mut conn = Conn { stream, buf: Vec::new(), pos: 0 }; + let status_line = conn.line(MAX_HEAD, limits)?; + let mut parts = status_line.split_whitespace(); + let version = parts.next().unwrap_or_default(); + if !version.starts_with("HTTP/1.") { + return Err(HttpError::Protocol(format!("not an HTTP/1.x response: {}", status_line.chars().take(60).collect::()))); + } + let status = parts.next().and_then(|s| s.parse::().ok()).ok_or_else(|| HttpError::Protocol("bad status line".into()))?; + let mut headers = Vec::new(); + let mut total = status_line.len(); + loop { + let line = conn.line(MAX_HEAD, limits)?; + if line.is_empty() { + break; + } + total = total.saturating_add(line.len()); + if total > MAX_HEAD || headers.len() > 200 { + return Err(HttpError::Protocol("response headers too large".into())); + } + if let Some((k, v)) = line.split_once(':') { + headers.push((k.trim().to_string(), v.trim().to_string())); + } + } + let find = |name: &str| headers.iter().find(|(k, _)| k.eq_ignore_ascii_case(name)).map(|(_, v)| v.as_str()); + let body = if status == 204 || status == 304 || (100..200).contains(&status) { + BodyMode::Done + } else if find("transfer-encoding").is_some_and(|v| v.to_ascii_lowercase().contains("chunked")) { + BodyMode::Chunked { left: 0, done: false } + } else if let Some(n) = find("content-length") { + BodyMode::Length(n.trim().parse::().map_err(|_| HttpError::Protocol(format!("bad Content-Length `{n}`")))?) + } else { + BodyMode::Close + }; + Ok(Response { status, headers, conn, body }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_urls() { + let u = Url::parse("https://cdn.example.com/models/sam3/model.safetensors?sig=1#x").unwrap(); + assert_eq!((u.tls, u.host.as_str(), u.port, u.path.as_str()), (true, "cdn.example.com", 443, "/models/sam3/model.safetensors?sig=1")); + assert_eq!(u.display(), "https://cdn.example.com/models/sam3/model.safetensors"); + let u = Url::parse("http://127.0.0.1:8080").unwrap(); + assert_eq!((u.tls, u.port, u.path.as_str()), (false, 8080, "/")); + let u = Url::parse("http://[::1]:9/a").unwrap(); + assert_eq!((u.host.as_str(), u.port), ("::1", 9)); + for bad in ["ftp://x/y", "https://", "https://a b/", "https://u:p@h/", "http://h:0/", "http://h:99999/", "https://h/\r\nX: y", "file:///etc"] + { + assert!(Url::parse(bad).is_err(), "{bad}"); + } + } + + #[test] + fn resolves_redirects() { + let u = Url::parse("https://a.example/x/y/file?q").unwrap(); + assert_eq!(u.join("https://b.example/z").unwrap().host, "b.example"); + assert_eq!(u.join("/root").unwrap().path, "/root"); + assert_eq!(u.join("other").unwrap().path, "/x/y/other"); + assert_eq!(u.join("//c.example/p").unwrap().host, "c.example"); + assert!(u.join("ftp://nope").is_err()); + } + + #[test] + fn the_tls_config_builds_with_the_pure_rust_provider() { + assert!(tls_config().is_ok()); + } +} diff --git a/crates/segment/src/fetch/mod.rs b/crates/segment/src/fetch/mod.rs new file mode 100644 index 00000000..92c5b6ea --- /dev/null +++ b/crates/segment/src/fetch/mod.rs @@ -0,0 +1,424 @@ +//! Downloading the SAM 3 model, only when the user asked for it: from an ordered list of mirrors +//! (the next one when a mirror fails), resuming partial files, with connect and stall timeouts, +//! a size cap, progress and cancellation, and checked (exact size, SHA-256 where pinned) before +//! each file is moved into the model folder. Runs on whatever thread calls [`download`] (the +//! engine uses a background thread; the UI thread never waits for the network). +//! +//! Files are written to `.part` next to their final place and renamed when verified, so a +//! model folder never holds a half-written file under a real name, and an interrupted download +//! resumes from where it stopped. A file that fails its check is deleted, never kept. +//! +//! The weights are never part of LightCraft (SAM License, see docs/ai-masks.md); where they are +//! downloaded from is configured by [`DEFAULT_MIRRORS`] and the user's own list. + +pub mod http; + +use std::fs::{File, OpenOptions}; +use std::io::{Read, Write}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; + +use sha2::{Digest, Sha256}; + +use http::{HttpError, Limits, Url}; + +/// One file of the model. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct FileSpec { + pub name: &'static str, + /// Exact size in bytes, when pinned (then the download resumes and is checked against it). + pub size: Option, + /// SHA-256 (lower-case hex), when pinned. + pub sha256: Option<&'static str>, + /// Largest size accepted when `size` is not pinned. + pub max: u64, +} + +/// The `facebook/sam3` checkpoint files [`crate::Sam3::load`] needs. +/// +/// `model.safetensors` is pinned to the official checkpoint (Hugging Face LFS SHA-256). The two +/// tokenizer files are small text files without a pinned hash yet: they are only parsed by the +/// tokenizer (never executed) and capped in size. +// TODO(maintainer): pin `vocab.json` and `merges.txt` (size + SHA-256) when the files are +// uploaded to LightCraft's CDN. +pub const SAM3_FILES: &[FileSpec] = &[ + FileSpec { name: "vocab.json", size: None, sha256: None, max: 16 << 20 }, + FileSpec { name: "merges.txt", size: None, sha256: None, max: 16 << 20 }, + FileSpec { + name: crate::WEIGHTS_FILE, + size: Some(SAM3_WEIGHTS_SIZE), + sha256: Some("6d06f0a5f84e435071fe6603e61d0b4cc7b40e0d39d487cfd4d67d8cc11cc14a"), + max: SAM3_WEIGHTS_SIZE, + }, +]; + +/// Size of the official `model.safetensors` (3.44 GB). +pub const SAM3_WEIGHTS_SIZE: u64 = 3_439_938_512; + +/// Where LightCraft downloads the model from, in order: each is a base URL, and file `f` is at +/// `/`. **Empty for now**: LightCraft has no download location of its own yet, so the +/// in-app download needs the user's list (`LIGHTCRAFT_SAM3_MIRRORS` or the mirrors file, see +/// [`mirrors`] and docs/ai-masks.md) until this is filled in. +/// +/// Hugging Face's `facebook/sam3` can't be a default: it is gated (each user must accept the +/// SAM License there and download with their own token). +// TODO(maintainer): add LightCraft's CDN locations (https, primary first), e.g. +// "https:///models/sam3/" — the files there must match `SAM3_FILES`. +pub const DEFAULT_MIRRORS: &[&str] = &[]; + +/// Environment variable with extra mirrors (base URLs separated by commas, spaces or newlines), +/// tried before the defaults. +pub const MIRRORS_ENV: &str = "LIGHTCRAFT_SAM3_MIRRORS"; + +/// Most mirrors used (a hostile list can't make a download loop forever). +const MAX_MIRRORS: usize = 16; +/// Redirects followed per request. +const MAX_REDIRECTS: usize = 8; + +/// The mirrors to try, in order: the environment variable's, then the mirrors file's (one base +/// URL per line, `#` comments), then [`DEFAULT_MIRRORS`]; duplicates and unusable URLs dropped. +pub fn mirrors(env: Option<&str>, file: Option<&Path>) -> Vec { + let mut out: Vec = Vec::new(); + let mut add = |s: &str| { + let s = s.trim().trim_end_matches('/'); + if !s.is_empty() && Url::parse(s).is_ok() && !out.iter().any(|m| m == s) && out.len() < MAX_MIRRORS { + out.push(s.to_string()); + } + }; + if let Some(env) = env { + env.split([',', ' ', '\n', '\t']).for_each(&mut add); + } + // a small text file; anything larger isn't a mirror list + if let Some(text) = file.and_then(|f| std::fs::metadata(f).ok().filter(|m| m.len() < 64 * 1024).and_then(|_| std::fs::read_to_string(f).ok())) { + text.lines().map(|l| l.split('#').next().unwrap_or_default()).for_each(&mut add); + } + DEFAULT_MIRRORS.iter().for_each(|m| add(m)); + out +} + +/// Download settings. +#[derive(Clone, Debug)] +pub struct Options { + /// Longest wait to connect to a mirror. + pub connect_timeout: Duration, + /// Longest wait for the next bytes once connected. + pub stall_timeout: Duration, + /// Tries per mirror and file (each resumes where the last one stopped). + pub attempts: u32, +} + +impl Default for Options { + fn default() -> Self { + Options { connect_timeout: Duration::from_secs(15), stall_timeout: Duration::from_secs(30), attempts: 3 } + } +} + +/// Where a download is. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct Progress { + pub file: String, + /// Bytes of all files on disk so far, and the total expected (pinned sizes + announced ones). + pub done: u64, + pub total: u64, + /// The mirror in use (without query). + pub mirror: String, +} + +/// Why a download failed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum DownloadError { + /// No mirror is configured. + NoMirrors, + Cancelled, + /// A local file problem (folder not writable, disk full…). + Disk(String), + /// Every mirror failed for `file`; the reason for each. + AllMirrorsFailed { + file: String, + errors: Vec, + }, +} + +impl std::fmt::Display for DownloadError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + DownloadError::NoMirrors => write!( + f, + "no download location is configured for the SAM 3 model in this build (set {MIRRORS_ENV}, or put the files in the model folder yourself; see docs/ai-masks.md)" + ), + DownloadError::Cancelled => write!(f, "the download was cancelled"), + DownloadError::Disk(e) => write!(f, "{e}"), + DownloadError::AllMirrorsFailed { file, errors } => write!(f, "could not download {file}: {}", errors.join("; ")), + } + } +} + +/// One mirror's failure for one file; `Retry` resumes on the same mirror. +enum Fail { + Retry(String), + NextMirror(String), + Cancelled, + Disk(String), +} + +/// Download `files` into `dir` from `mirrors` (in order; see [`mirrors`]). Files already there +/// and verified are kept. `progress` is called as bytes arrive (often: throttle in it). +pub fn download( + files: &[FileSpec], + mirrors: &[String], + dir: &Path, + opts: &Options, + cancel: &AtomicBool, + progress: &mut dyn FnMut(&Progress), +) -> Result<(), DownloadError> { + std::fs::create_dir_all(dir).map_err(|e| DownloadError::Disk(format!("can't create {}: {e}", dir.display())))?; + // what's left to fetch + let mut todo = Vec::new(); + for f in files { + if !verified(f, &dir.join(f.name))? { + todo.push(f); + } + } + if todo.is_empty() { + return Ok(()); + } + if mirrors.is_empty() { + return Err(DownloadError::NoMirrors); + } + let total: u64 = files.iter().map(|f| f.size.unwrap_or(0)).sum(); + let mut before: u64 = files.iter().filter(|f| !todo.contains(f)).filter_map(|f| f.size).sum(); + for f in todo { + let mut errors = Vec::new(); + let mut ok = false; + 'mirrors: for m in mirrors { + let url = format!("{m}/{}", f.name); + for _ in 0..opts.attempts.max(1) { + if cancel.load(Ordering::Relaxed) { + return Err(DownloadError::Cancelled); + } + let mut report = |n: u64, announced: u64| { + let mirror = Url::parse(m).map(|u| u.display()).unwrap_or_default(); + progress(&Progress { + file: f.name.to_string(), + done: before.saturating_add(n), + total: total.max(before.saturating_add(announced)), + mirror, + }) + }; + match fetch_file(f, &url, dir, opts, cancel, &mut report) { + Ok(()) => { + ok = true; + break 'mirrors; + } + Err(Fail::Cancelled) => return Err(DownloadError::Cancelled), + Err(Fail::Disk(e)) => return Err(DownloadError::Disk(e)), + Err(Fail::Retry(e)) => { + log::warn!("SAM 3 download of {url}: {e} (retrying)"); + errors.push(format!("{}: {e}", short(m))); + } + Err(Fail::NextMirror(e)) => { + log::warn!("SAM 3 download of {url}: {e}"); + errors.push(format!("{}: {e}", short(m))); + continue 'mirrors; + } + } + } + } + if !ok { + return Err(DownloadError::AllMirrorsFailed { file: f.name.to_string(), errors }); + } + before = before.saturating_add(std::fs::metadata(dir.join(f.name)).map(|m| m.len()).unwrap_or(0)); + } + Ok(()) +} + +/// A mirror for messages (no query string: it may hold a token). +fn short(m: &str) -> String { + Url::parse(m).map(|u| u.display()).unwrap_or_else(|_| "mirror".into()) +} + +/// Whether `path` is already the right file (exact size and hash when pinned). +fn verified(f: &FileSpec, path: &Path) -> Result { + let Ok(meta) = std::fs::metadata(path) else { return Ok(false) }; + if !meta.is_file() || meta.len() == 0 { + return Ok(false); + } + if f.size.is_some_and(|s| s != meta.len()) || meta.len() > f.max { + return Ok(false); + } + match f.sha256 { + Some(want) => Ok(sha256_file(path).map_err(DownloadError::Disk)? == want), + None => Ok(true), + } +} + +fn sha256_file(path: &Path) -> Result { + let mut file = File::open(path).map_err(|e| format!("{}: {e}", path.display()))?; + let mut h = Sha256::new(); + let mut buf = vec![0u8; 1 << 20]; + loop { + let n = file.read(&mut buf).map_err(|e| format!("{}: {e}", path.display()))?; + if n == 0 { + break; + } + h.update(buf.get(..n).unwrap_or_default()); + } + Ok(h.finalize().iter().map(|b| format!("{b:02x}")).collect()) +} + +fn part_path(dir: &Path, name: &str) -> PathBuf { + dir.join(format!("{name}.part")) +} + +/// Fetch one file from one mirror into `.part` (resuming it when its size is pinned), +/// check it and move it into place. +fn fetch_file(f: &FileSpec, url: &str, dir: &Path, opts: &Options, cancel: &AtomicBool, report: &mut dyn FnMut(u64, u64)) -> Result<(), Fail> { + let part = part_path(dir, f.name); + let limit = f.size.unwrap_or(f.max); + let disk = |e: std::io::Error| Fail::Disk(format!("{}: {e}", part.display())); + // resume only a pinned file (an unpinned one can't be checked after resuming) + let mut have = match std::fs::metadata(&part) { + Ok(m) if f.size.is_some() && m.len() <= limit => m.len(), + Ok(_) => { + std::fs::remove_file(&part).map_err(disk)?; + 0 + } + Err(_) => 0, + }; + let fetched = f.size != Some(have) || have == 0; + if fetched { + let limits = Limits { connect: opts.connect_timeout, stall: opts.stall_timeout, cancel }; + let mut url = Url::parse(url).map_err(|e| Fail::NextMirror(e.to_string()))?; + let origin = url.host.clone(); + let mut redirects = 0; + let mut resp = loop { + let mut headers = vec![]; + if have > 0 { + headers.push(("Range", format!("bytes={have}-"))); + } + // a token for the configured host only (never sent on to a redirect's host) + if let Ok(t) = std::env::var("LIGHTCRAFT_SAM3_TOKEN") + && !t.trim().is_empty() + && url.tls + && url.host == origin + { + headers.push(("Authorization", format!("Bearer {}", t.trim()))); + } + let r = http::get(&url, &headers, &limits).map_err(http_fail)?; + match r.status { + 301 | 302 | 303 | 307 | 308 => { + redirects += 1; + if redirects > MAX_REDIRECTS { + return Err(Fail::NextMirror("too many redirects".into())); + } + let loc = r.header("location").ok_or_else(|| Fail::NextMirror("redirect without a location".into()))?; + url = url.join(loc).map_err(|e| Fail::NextMirror(e.to_string()))?; + } + _ => break r, + } + }; + let mut file = match resp.status { + 200 => { + // the whole file (from the start, also when a range was asked for) + if resp.content_length().is_some_and(|n| n > limit || f.size.is_some_and(|s| s != n)) { + return Err(Fail::NextMirror(format!("the server's file has the wrong size ({:?} bytes)", resp.content_length()))); + } + have = 0; + File::create(&part).map_err(disk)? + } + 206 if have > 0 => { + let range = resp.header("content-range").unwrap_or_default(); + let (start, total) = parse_content_range(range).ok_or_else(|| Fail::NextMirror(format!("bad Content-Range `{range}`")))?; + if start != have || total.is_some_and(|t| Some(t) != f.size) { + return Err(Fail::NextMirror(format!("the server resumed at the wrong place ({range})"))); + } + OpenOptions::new().append(true).open(&part).map_err(disk)? + } + 416 => { + // the partial file doesn't fit what the server has: start over + std::fs::remove_file(&part).map_err(disk)?; + return Err(Fail::Retry("the partial download didn't match; starting over".into())); + } + 401 | 403 => return Err(Fail::NextMirror(format!("access denied (HTTP {})", resp.status))), + 404 | 410 => return Err(Fail::NextMirror(format!("not found (HTTP {})", resp.status))), + s if s >= 500 => return Err(Fail::Retry(format!("server error (HTTP {s})"))), + s => return Err(Fail::NextMirror(format!("HTTP {s}"))), + }; + let announced = f.size.unwrap_or_else(|| resp.content_length().unwrap_or(0)); + report(have, announced); + let mut buf = vec![0u8; 256 * 1024]; + loop { + let n = match resp.read(&mut buf, &limits) { + Ok(n) => n, + Err(e) => { + let _ = file.flush(); + if f.size.is_none() { + drop(file); + let _ = std::fs::remove_file(&part); + } + return Err(http_fail(e)); + } + }; + if n == 0 { + break; + } + if have.saturating_add(n as u64) > limit { + drop(file); + let _ = std::fs::remove_file(&part); + return Err(Fail::NextMirror(format!("the server sent more than the expected {limit} bytes"))); + } + file.write_all(buf.get(..n).unwrap_or_default()).map_err(disk)?; + have = have.saturating_add(n as u64); + report(have, announced); + } + file.sync_all().map_err(disk)?; + } + // checks: exact size, hash + let len = std::fs::metadata(&part).map_err(disk)?.len(); + if f.size.is_some_and(|s| s != len) { + return Err(Fail::Retry(format!("incomplete ({len} of {} bytes)", f.size.unwrap_or(0)))); + } + if len == 0 { + let _ = std::fs::remove_file(&part); + return Err(Fail::NextMirror("the server sent an empty file".into())); + } + if let Some(want) = f.sha256 { + let got = sha256_file(&part).map_err(Fail::Disk)?; + if got != want { + std::fs::remove_file(&part).map_err(disk)?; + let why = format!("the file is damaged or not the official one (SHA-256 {got})"); + // a damaged partial file from an earlier run isn't this mirror's fault + return Err(if fetched { Fail::NextMirror(why) } else { Fail::Retry(why) }); + } + } + std::fs::rename(&part, dir.join(f.name)).map_err(disk)?; + Ok(()) +} + +fn http_fail(e: HttpError) -> Fail { + match e { + HttpError::Cancelled => Fail::Cancelled, + HttpError::Stalled | HttpError::Io(_) | HttpError::Protocol(_) => Fail::Retry(e.to_string()), + HttpError::BadUrl(_) | HttpError::Connect(_) | HttpError::Tls(_) => Fail::NextMirror(e.to_string()), + } +} + +/// `bytes a-b/total` → (a, total) (`*` total → None). +fn parse_content_range(v: &str) -> Option<(u64, Option)> { + let v = v.trim().strip_prefix("bytes")?.trim(); + let (range, total) = v.split_once('/')?; + let (a, b) = range.split_once('-')?; + let (a, b) = (a.trim().parse::().ok()?, b.trim().parse::().ok()?); + if b < a { + return None; + } + let total = match total.trim() { + "*" => None, + t => Some(t.parse::().ok()?), + }; + Some((a, total)) +} + +#[cfg(test)] +mod tests; diff --git a/crates/segment/src/fetch/tests.rs b/crates/segment/src/fetch/tests.rs new file mode 100644 index 00000000..81e38f01 --- /dev/null +++ b/crates/segment/src/fetch/tests.rs @@ -0,0 +1,314 @@ +//! The downloader against a local HTTP server (std `TcpListener`; never the internet): mirror +//! fallback, redirects, resuming, chunked bodies, stalls, hash mismatches, size caps, cancel. + +use std::io::{BufRead, BufReader, Write}; +use std::net::{TcpListener, TcpStream}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; + +use super::*; + +/// The served file: 300 kB of a deterministic pattern. +fn payload() -> Vec { + (0..300_000u32).map(|i| (i.wrapping_mul(2_654_435_761) >> 13) as u8).collect() +} + +fn sha_hex(b: &[u8]) -> String { + Sha256::digest(b).iter().map(|b| format!("{b:02x}")).collect() +} + +fn leak(s: String) -> &'static str { + Box::leak(s.into_boxed_str()) +} + +fn pinned() -> FileSpec { + let p = payload(); + FileSpec { name: "model.bin", size: Some(p.len() as u64), sha256: Some(leak(sha_hex(&p))), max: p.len() as u64 } +} + +struct Server { + base: String, + hits: Arc, +} + +/// Serve on 127.0.0.1; the first path segment picks the behaviour. +fn server() -> Server { + let l = TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}", l.local_addr().unwrap()); + let hits = Arc::new(AtomicUsize::new(0)); + let h = hits.clone(); + std::thread::spawn(move || { + for s in l.incoming().flatten() { + let h = h.clone(); + std::thread::spawn(move || handle(s, &h)); + } + }); + Server { base, hits } +} + +fn handle(mut s: TcpStream, hits: &AtomicUsize) { + let n = hits.fetch_add(1, Ordering::SeqCst); + let mut r = BufReader::new(s.try_clone().unwrap()); + let mut first = String::new(); + if r.read_line(&mut first).is_err() { + return; + } + let mut range: Option = None; + loop { + let mut l = String::new(); + if r.read_line(&mut l).unwrap_or(0) == 0 || l.trim().is_empty() { + break; + } + if let Some(v) = l.to_ascii_lowercase().strip_prefix("range: bytes=") { + range = v.trim().trim_end_matches('-').parse().ok(); + } + } + let path = first.split_whitespace().nth(1).unwrap_or("/").to_string(); + let mode = path.trim_start_matches('/').split('/').next().unwrap_or("").to_string(); + let body = payload(); + let send = |s: &mut TcpStream, status: &str, headers: &str, data: &[u8]| { + let _ = s.write_all(format!("HTTP/1.1 {status}\r\n{headers}Connection: close\r\n\r\n").as_bytes()); + let _ = s.write_all(data); + }; + match mode.as_str() { + "missing" => send(&mut s, "404 Not Found", "Content-Length: 9\r\n", b"not found"), + "error" => send(&mut s, "500 Oops", "Content-Length: 25\r\n", b"server error"), + "redirect" => send(&mut s, "302 Found", &format!("Location: {}\r\nContent-Length: 0\r\n", path.replacen("/redirect", "/good", 1)), b""), + "wrong" => { + let mut b = body.clone(); + b[1000] ^= 0xff; + send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", b.len()), &b) + } + "toolong" => { + // no length announced, more bytes than the file can have + let mut b = body.clone(); + b.extend_from_slice(&[7u8; 5000]); + send(&mut s, "200 OK", "", &b) + } + "chunked" => { + let _ = s.write_all(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n"); + for c in body.chunks(70_001) { + let _ = s.write_all(format!("{:x};ext=1\r\n", c.len()).as_bytes()); + let _ = s.write_all(c); + let _ = s.write_all(b"\r\n"); + } + let _ = s.write_all(b"0\r\nX-Trailer: 1\r\n\r\n"); + } + "stall" => { + // half the file, then nothing + let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes()); + let _ = s.write_all(&body[..body.len() / 2]); + std::thread::sleep(Duration::from_secs(5)); + } + "slow" => { + let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes()); + for c in body.chunks(1000) { + if s.write_all(c).is_err() { + return; + } + std::thread::sleep(Duration::from_millis(20)); + } + } + "flaky" if n.is_multiple_of(2) && range.is_none() => { + // the first try breaks off after a third + let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes()); + let _ = s.write_all(&body[..body.len() / 3]); + } + "range416" if range.is_some() => send(&mut s, "416 Range Not Satisfiable", "Content-Length: 0\r\n", b""), + "norange" | "range416" => send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", body.len()), &body), + _ => match range { + Some(from) if from as usize >= body.len() => send(&mut s, "416 Range Not Satisfiable", "Content-Length: 0\r\n", b""), + Some(from) => send( + &mut s, + "206 Partial Content", + &format!("Content-Range: bytes {from}-{}/{}\r\nContent-Length: {}\r\n", body.len() - 1, body.len(), body.len() - from as usize), + &body[from as usize..], + ), + None => send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", body.len()), &body), + }, + } +} + +fn tmp(name: &str) -> PathBuf { + let d = std::env::temp_dir().join(format!("lc-sam3-fetch-{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&d); + d +} + +fn quick() -> Options { + Options { connect_timeout: Duration::from_secs(2), stall_timeout: Duration::from_millis(600), attempts: 2 } +} + +fn run(files: &[FileSpec], mirrors: &[String], dir: &Path, cancel: &AtomicBool) -> (Result<(), DownloadError>, Vec) { + let mut seen = Vec::new(); + let r = download(files, mirrors, dir, &quick(), cancel, &mut |p| seen.push(p.clone())); + (r, seen) +} + +#[test] +fn falls_back_to_the_next_mirror_and_verifies() { + let srv = server(); + let dir = tmp("fallback"); + let mirrors = vec![ + format!("{}/missing", srv.base), + format!("{}/error", srv.base), + "http://127.0.0.1:1/unreachable".to_string(), + format!("{}/redirect", srv.base), + ]; + let (r, seen) = run(&[pinned()], &mirrors, &dir, &AtomicBool::new(false)); + r.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + assert!(!part_path(&dir, "model.bin").exists()); + let last = seen.last().unwrap(); + assert_eq!((last.done, last.total), (300_000, 300_000)); + // an error page never ends up in the file: nothing left over from the failed mirrors + assert_eq!(std::fs::read_dir(&dir).unwrap().count(), 1); + // installed: nothing is fetched again + let hits = srv.hits.load(Ordering::SeqCst); + run(&[pinned()], &mirrors, &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(srv.hits.load(Ordering::SeqCst), hits); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn a_hash_mismatch_is_deleted_and_the_next_mirror_used() { + let srv = server(); + let dir = tmp("hash"); + let (r, _) = run(&[pinned()], &[format!("{}/wrong", srv.base)], &dir, &AtomicBool::new(false)); + let e = r.unwrap_err().to_string(); + assert!(e.contains("SHA-256"), "{e}"); + assert!(!dir.join("model.bin").exists() && !part_path(&dir, "model.bin").exists(), "a damaged file is never kept"); + run(&[pinned()], &[format!("{}/wrong", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn resumes_partial_files_and_restarts_bad_ones() { + let srv = server(); + let dir = tmp("resume"); + std::fs::create_dir_all(&dir).unwrap(); + // a third already there + std::fs::write(part_path(&dir, "model.bin"), &payload()[..100_000]).unwrap(); + run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + // a partial file with the wrong bytes: the hash fails, it's deleted, the next try starts over + std::fs::remove_file(dir.join("model.bin")).unwrap(); + std::fs::write(part_path(&dir, "model.bin"), vec![0u8; 100_000]).unwrap(); + run(&[pinned()], &[format!("{}/good", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + // a server ignoring Range: the whole file again, not appended + std::fs::remove_file(dir.join("model.bin")).unwrap(); + std::fs::write(part_path(&dir, "model.bin"), &payload()[..5]).unwrap(); + run(&[pinned()], &[format!("{}/norange", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + // a connection breaking off is resumed on the same mirror + std::fs::remove_file(dir.join("model.bin")).unwrap(); + run(&[pinned()], &[format!("{}/flaky", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn a_416_or_a_bad_complete_part_starts_over() { + let srv = server(); + let dir = tmp("416"); + std::fs::create_dir_all(&dir).unwrap(); + // the server refuses the range: the part is dropped and the file fetched whole (no loop) + std::fs::write(part_path(&dir, "model.bin"), &payload()[..100]).unwrap(); + run(&[pinned()], &[format!("{}/range416", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + // as long as the file but wrong: hash fails → deleted → downloaded again from the same mirror + std::fs::remove_file(dir.join("model.bin")).unwrap(); + std::fs::write(part_path(&dir, "model.bin"), vec![1u8; 300_000]).unwrap(); + run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn chunked_bodies_and_unpinned_files() { + let srv = server(); + let dir = tmp("chunked"); + let f = FileSpec { name: "vocab.json", size: None, sha256: None, max: 1 << 20 }; + run(std::slice::from_ref(&f), &[format!("{}/chunked", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("vocab.json")).unwrap(), payload()); + // more than the cap: refused, nothing kept + let small = FileSpec { max: 1000, ..f }; + let dir2 = tmp("cap"); + let e = run(&[small], &[format!("{}/good", srv.base)], &dir2, &AtomicBool::new(false)).0.unwrap_err().to_string(); + assert!(e.contains("wrong size") || e.contains("more than"), "{e}"); + assert!(!dir2.join("vocab.json").exists() && !part_path(&dir2, "vocab.json").exists()); + let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::remove_dir_all(&dir2); +} + +#[test] +fn a_server_sending_too_much_is_cut_off() { + let srv = server(); + let dir = tmp("toolong"); + let e = run(&[pinned()], &[format!("{}/toolong", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap_err().to_string(); + assert!(e.contains("more than"), "{e}"); + assert!(!dir.join("model.bin").exists() && !part_path(&dir, "model.bin").exists()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn a_stalled_server_times_out_and_the_next_mirror_finishes() { + let srv = server(); + let dir = tmp("stall"); + let t = Instant::now(); + let (r, _) = run(&[pinned()], &[format!("{}/stall", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)); + r.unwrap(); + assert!(t.elapsed() < Duration::from_secs(4), "{:?}", t.elapsed()); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn cancel_stops_quickly_and_keeps_the_part_for_later() { + let srv = server(); + let dir = tmp("cancel"); + let cancel = Arc::new(AtomicBool::new(false)); + let c = cancel.clone(); + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(400)); + c.store(true, Ordering::SeqCst); + }); + let t = Instant::now(); + let (r, _) = run(&[pinned()], &[format!("{}/slow", srv.base)], &dir, &cancel); + assert_eq!(r, Err(DownloadError::Cancelled)); + assert!(t.elapsed() < Duration::from_secs(3), "{:?}", t.elapsed()); + assert!(!dir.join("model.bin").exists()); + let kept = std::fs::metadata(part_path(&dir, "model.bin")).map(|m| m.len()).unwrap_or(0); + assert!(kept > 0 && kept < 300_000, "{kept}"); + // and the next run resumes it + run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn no_mirrors_is_a_clear_error_and_mirror_lists_are_parsed() { + let dir = tmp("nomirror"); + let e = run(&[pinned()], &[], &dir, &AtomicBool::new(false)).0.unwrap_err(); + assert_eq!(e, DownloadError::NoMirrors); + assert!(e.to_string().contains(MIRRORS_ENV)); + let file = std::env::temp_dir().join(format!("lc-sam3-mirrors-{}.txt", std::process::id())); + std::fs::write(&file, "# mine\nhttps://b.example/m/ # second\n\nnot a url\nhttps://a.example/x\n").unwrap(); + let m = mirrors(Some("https://a.example/x/, ftp://no"), Some(&file)); + assert_eq!(m[..2], ["https://a.example/x".to_string(), "https://b.example/m".to_string()]); + assert_eq!(m.len(), 2 + DEFAULT_MIRRORS.len()); + assert_eq!(mirrors(None, Some(Path::new("/nonexistent/mirrors.txt"))).len(), DEFAULT_MIRRORS.len()); + let _ = std::fs::remove_file(&file); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn content_ranges() { + assert_eq!(parse_content_range("bytes 100-299/300"), Some((100, Some(300)))); + assert_eq!(parse_content_range("bytes 0-0/*"), Some((0, None))); + assert_eq!(parse_content_range("bytes 9-1/300"), None); + assert_eq!(parse_content_range("items 1-2/3"), None); +} diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs new file mode 100644 index 00000000..42aed538 --- /dev/null +++ b/crates/segment/src/lib.rs @@ -0,0 +1,263 @@ +//! 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`. +//! This crate is therefore licensed under the Apache License 2.0 only (not the MIT option of +//! the rest of LightCraft). No UI dependencies (L3). +//! +//! [`fetch`] downloads the model, only when the user asks for it (from configurable mirrors, +//! verified before use). Nothing in LightCraft requires the model: without it, AI masks report +//! that it isn't installed and everything else works. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`src/transformers/models/sam3` and `sam3_tracker`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. +#![forbid(unsafe_code)] +#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] + +#[cfg(not(target_arch = "wasm32"))] +mod clip; +#[cfg(not(target_arch = "wasm32"))] +mod detector; +#[cfg(not(target_arch = "wasm32"))] +pub mod fetch; +pub mod mask; +#[cfg(not(target_arch = "wasm32"))] +mod neck; +#[cfg(not(target_arch = "wasm32"))] +mod nn; +pub mod tokenizer; +#[cfg(not(target_arch = "wasm32"))] +mod tracker; +#[cfg(not(target_arch = "wasm32"))] +mod vit; +#[cfg(not(target_arch = "wasm32"))] +mod weights; + +use std::path::{Path, PathBuf}; +#[cfg(not(target_arch = "wasm32"))] +use std::sync::Arc; + +#[cfg(not(target_arch = "wasm32"))] +use candle_core::{DType, Device, Tensor}; +#[cfg(not(target_arch = "wasm32"))] +use candle_nn::VarBuilder; + +#[cfg(not(target_arch = "wasm32"))] +pub use tracker::Click; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("SAM 3 model: {0}")] + Model(String), + #[error("SAM 3 model files not found in {0} (expected model.safetensors, vocab.json, merges.txt)")] + Missing(PathBuf), + #[cfg(not(target_arch = "wasm32"))] + #[error(transparent)] + Candle(#[from] candle_core::Error), +} + +pub type Result = std::result::Result; + +/// Side of the square model input. +pub const INPUT: usize = 1008; +/// Side of the low-resolution masks the model predicts. +pub const MASK_SIDE: usize = 288; +#[cfg(not(target_arch = "wasm32"))] +const _: () = assert!(MASK_SIDE == tracker::LOW_RES && INPUT == vit::IMAGE); + +/// The files of a model directory. +pub const WEIGHTS_FILE: &str = "model.safetensors"; + +/// Whether `dir` holds a usable checkpoint. +pub fn is_model_dir(dir: &Path) -> bool { + dir.join(WEIGHTS_FILE).is_file() && dir.join("vocab.json").is_file() && dir.join("merges.txt").is_file() +} + +#[cfg(not(target_arch = "wasm32"))] +/// The best device here: Metal on macOS (when a GPU is available), else the CPU. +pub fn best_device() -> Device { + #[cfg(target_os = "macos")] + if let Ok(d) = Device::new_metal(0) { + return d; + } + Device::Cpu +} + +#[cfg(not(target_arch = "wasm32"))] +/// A loaded model. The vision backbone loads with it; the point and text heads load on first use. +pub struct Sam3 { + dir: PathBuf, + device: Device, + weights: Arc, + backbone: vit::Vit, + tracker_neck: Option, + tracker: Option, + detector: Option, + tokenizer: Option, +} + +#[cfg(not(target_arch = "wasm32"))] +/// An image run through the backbone, with the head-specific features computed on demand. +pub struct Encoded { + backbone: Tensor, + tracker: Option, + detector: Option, +} + +/// A predicted mask: `MASK_SIDE²` logits over the model's square input (which the image was +/// stretched to), and a confidence 0..1. +#[derive(Clone, Debug)] +pub struct Prediction { + pub logits: Vec, + pub score: f32, +} + +#[cfg(not(target_arch = "wasm32"))] +impl Sam3 { + /// Load from a directory holding `model.safetensors` (the `facebook/sam3` checkpoint in + /// Hugging Face format), `vocab.json` and `merges.txt`. + pub fn load(dir: &Path) -> Result { + Self::load_on(dir, best_device()) + } + + pub fn load_on(dir: &Path, device: Device) -> Result { + if !is_model_dir(dir) { + return Err(Error::Missing(dir.to_path_buf())); + } + let weights = weights::Weights::open(&dir.join(WEIGHTS_FILE))?; + let backbone = vit::Vit::new(Self::vb(&weights, &device).pp("detector_model.vision_encoder.backbone"))?; + Ok(Self { dir: dir.to_path_buf(), device, weights, backbone, tracker_neck: None, tracker: None, detector: None, tokenizer: None }) + } + + fn vb<'a>(w: &Arc, device: &Device) -> VarBuilder<'a> { + VarBuilder::from_backend(Box::new(weights::Backend(w.clone())), DType::F32, device.clone()) + } + + pub fn device(&self) -> &Device { + &self.device + } + + /// Run the backbone on an RGB image (8-bit, row-major, `w × h`), stretched to the square + /// input like the reference preprocessing. + pub fn encode(&self, rgb: &[u8], w: usize, h: usize) -> Result { + let pixels = mask::preprocess(rgb, w, h, INPUT)?; + let pixels = Tensor::from_vec(pixels, (1, 3, INPUT, INPUT), &self.device)?; + self.encode_pixels(&pixels) + } + + /// Run the backbone on preprocessed pixels `[1, 3, 1008, 1008]` (−1..1). + pub fn encode_pixels(&self, pixels: &Tensor) -> Result { + Ok(Encoded { backbone: self.backbone.forward(&pixels.to_device(&self.device)?)?, tracker: None, detector: None }) + } + + fn tracker_parts(&mut self) -> Result<(&neck::Neck, &tracker::Tracker)> { + if self.tracker_neck.is_none() { + self.tracker_neck = Some(neck::Neck::new(Self::vb(&self.weights, &self.device).pp("tracker_neck"))?); + } + if self.tracker.is_none() { + self.tracker = Some(tracker::Tracker::new(Self::vb(&self.weights, &self.device).pp("tracker_model"))?); + } + match (&self.tracker_neck, &self.tracker) { + (Some(n), Some(t)) => Ok((n, t)), + _ => Err(Error::Model("tracker failed to load".into())), + } + } + + /// Compute the click decoder's features for `enc` now (so the first click is fast). + pub fn prepare_clicks(&mut self, enc: &mut Encoded) -> Result<()> { + let (neck, tracker) = self.tracker_parts()?; + if enc.tracker.is_none() { + let levels = neck.forward(&enc.backbone, 3)?; + enc.tracker = Some(tracker.features(&levels)?); + } + Ok(()) + } + + /// Segment one object from clicks (positions normalized to the image, 0..1). + pub fn segment_clicks(&mut self, enc: &mut Encoded, clicks: &[Click]) -> Result { + if clicks.is_empty() { + return Err(Error::Model("no clicks".into())); + } + let (neck, tracker) = self.tracker_parts()?; + if enc.tracker.is_none() { + let levels = neck.forward(&enc.backbone, 3)?; + enc.tracker = Some(tracker.features(&levels)?); + } + let f = enc.tracker.as_ref().ok_or_else(|| Error::Model("tracker features".into()))?; + // one click: SAM's three alternatives, best by predicted IoU; more: the single mask + let p = tracker.predict(f, clicks, clicks.len() == 1)?; + Ok(Prediction { logits: p.logits, score: p.iou * sigmoid(p.object_score) }) + } + + fn detector_parts(&mut self) -> Result<(&detector::Detector, &tokenizer::Tokenizer)> { + if self.detector.is_none() { + self.detector = Some(detector::Detector::new(Self::vb(&self.weights, &self.device).pp("detector_model"))?); + } + if self.tokenizer.is_none() { + self.tokenizer = Some(tokenizer::Tokenizer::load(&self.dir)?); + } + match (&self.detector, &self.tokenizer) { + (Some(d), Some(t)) => Ok((d, t)), + _ => Err(Error::Model("detector failed to load".into())), + } + } + + /// Segment every instance of the concept `text` names. Instances scoring above + /// `threshold` (0..1, 0.5 is the reference default) are merged into one mask; `None` when + /// nothing matches. + pub fn segment_text(&mut self, enc: &mut Encoded, text: &str, threshold: f32) -> Result> { + let (det, tok) = self.detector_parts()?; + if enc.detector.is_none() { + enc.detector = Some(det.features(&enc.backbone)?); + } + let f = enc.detector.as_ref().ok_or_else(|| Error::Model("detector features".into()))?; + let (ids, valid) = tok.encode(text); + let out = det.predict(f, &ids, &valid)?; + Ok(detector::merge(&out, threshold)) + } + + /// Raw detector outputs for `text` (tests and tools). + pub fn detect(&mut self, enc: &mut Encoded, text: &str) -> Result { + let (det, tok) = self.detector_parts()?; + if enc.detector.is_none() { + enc.detector = Some(det.features(&enc.backbone)?); + } + let f = enc.detector.as_ref().ok_or_else(|| Error::Model("detector features".into()))?; + let (ids, valid) = tok.encode(text); + det.predict(f, &ids, &valid) + } +} + +#[cfg(not(target_arch = "wasm32"))] +impl Encoded { + /// The backbone output `[1, 1024, 72, 72]` (tests). + pub fn backbone(&self) -> &Tensor { + &self.backbone + } +} + +#[cfg(not(target_arch = "wasm32"))] +fn sigmoid(x: f32) -> f32 { + 1.0 / (1.0 + (-x).exp()) +} + +#[cfg(not(target_arch = "wasm32"))] +pub use detector::Output as DetectorOutput; + +#[cfg(not(target_arch = "wasm32"))] +/// Wait for queued GPU work to finish (benchmarks). +pub fn sync(model: &Sam3) -> Result<()> { + Ok(model.device.synchronize()?) +} diff --git a/crates/segment/src/mask.rs b/crates/segment/src/mask.rs new file mode 100644 index 00000000..af349a26 --- /dev/null +++ b/crates/segment/src/mask.rs @@ -0,0 +1,130 @@ +//! Image preprocessing and mask post-processing (CPU, no model needed). +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/image_processing_sam3*.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. + +use crate::{Error, Result}; + +/// Separable resampling weights for `src → dst` samples with an antialiased triangle filter +/// (Pillow's BILINEAR: the support widens with the downscale factor). +fn weights(src: usize, dst: usize) -> Vec<(usize, Vec)> { + let scale = src as f64 / dst as f64; + let ss = scale.max(1.0); + let support = ss; + (0..dst) + .map(|o| { + let center = (o as f64 + 0.5) * scale; + let lo = ((center - support + 0.5).floor().max(0.0)) as usize; + let hi = ((center + support + 0.5).floor() as usize).min(src); + let mut w: Vec = (lo..hi).map(|i| (1.0 - ((i as f64 - center + 0.5) / ss).abs()).max(0.0) as f32).collect(); + let sum: f32 = w.iter().sum(); + if sum > 0.0 { + w.iter_mut().for_each(|v| *v /= sum); + } + (lo, w) + }) + .collect() +} + +/// An 8-bit RGB image (`w × h`, row-major) stretched to `side × side`, scaled to −1..1 and laid +/// out channel-first (`[3, side, side]`), as SAM 3's processor does. +pub fn preprocess(rgb: &[u8], w: usize, h: usize, side: usize) -> Result> { + let n = w.checked_mul(h).and_then(|n| n.checked_mul(3)).ok_or_else(|| Error::Model("image too large".into()))?; + if w == 0 || h == 0 || rgb.len() < n { + return Err(Error::Model(format!("image buffer of {} bytes for {w} × {h}", rgb.len()))); + } + let (wx, wy) = (weights(w, side), weights(h, side)); + // horizontal pass: h rows × side columns × 3 + let mut tmp = vec![0f32; h * side * 3]; + for y in 0..h { + for (ox, (lo, ws)) in wx.iter().enumerate() { + let mut acc = [0f32; 3]; + for (k, wt) in ws.iter().enumerate() { + let i = (y * w + lo + k) * 3; + for (c, a) in acc.iter_mut().enumerate() { + *a += wt * f32::from(rgb.get(i + c).copied().unwrap_or(0)); + } + } + if let Some(dst) = tmp.get_mut((y * side + ox) * 3..(y * side + ox) * 3 + 3) { + dst.copy_from_slice(&acc); + } + } + } + let mut out = vec![0f32; 3 * side * side]; + for (oy, (lo, ws)) in wy.iter().enumerate() { + for ox in 0..side { + let mut acc = [0f32; 3]; + for (k, wt) in ws.iter().enumerate() { + let i = ((lo + k) * side + ox) * 3; + for (c, a) in acc.iter_mut().enumerate() { + *a += wt * tmp.get(i + c).copied().unwrap_or(0.0); + } + } + for (c, a) in acc.iter().enumerate() { + // Pillow rounds to 8 bits; then /255 and (x − 0.5) / 0.5 + let v = a.round().clamp(0.0, 255.0) / 255.0; + if let Some(d) = out.get_mut((c * side + oy) * side + ox) { + *d = (v - 0.5) / 0.5; + } + } + } + } + Ok(out) +} + +/// Mask `logits` (`side × side`) resampled bilinearly (half-pixel centres) to `w × h` and +/// mapped to 0..1 coverage (sigmoid). +pub fn alpha(logits: &[f32], side: usize, w: usize, h: usize) -> Vec { + let mut out = vec![0f32; w * h]; + if side == 0 || logits.len() < side * side { + return out; + } + let sample = |x: usize, y: usize| logits.get(y.min(side - 1) * side + x.min(side - 1)).copied().unwrap_or(-20.0); + let (sx, sy) = (side as f32 / w.max(1) as f32, side as f32 / h.max(1) as f32); + for y in 0..h { + let fy = ((y as f32 + 0.5) * sy - 0.5).max(0.0); + let (y0, ty) = (fy.floor() as usize, fy.fract()); + for x in 0..w { + let fx = ((x as f32 + 0.5) * sx - 0.5).max(0.0); + let (x0, tx) = (fx.floor() as usize, fx.fract()); + let top = sample(x0, y0) * (1.0 - tx) + sample(x0 + 1, y0) * tx; + let bot = sample(x0, y0 + 1) * (1.0 - tx) + sample(x0 + 1, y0 + 1) * tx; + let l = top * (1.0 - ty) + bot * ty; + if let Some(o) = out.get_mut(y * w + x) { + *o = 1.0 / (1.0 + (-l).exp()); + } + } + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn preprocess_is_bounded_and_channel_first() { + let rgb: Vec = (0..20 * 10 * 3).map(|i| (i % 256) as u8).collect(); + let p = preprocess(&rgb, 20, 10, 8).unwrap(); + assert_eq!(p.len(), 3 * 64); + assert!(p.iter().all(|v| (-1.0..=1.0).contains(v))); + // a flat image stays flat + let flat = vec![255u8; 30 * 30 * 3]; + assert!(preprocess(&flat, 30, 30, 7).unwrap().iter().all(|v| (*v - 1.0).abs() < 1e-6)); + } + + #[test] + fn short_buffers_are_errors() { + assert!(preprocess(&[0; 5], 2, 2, 4).is_err()); + assert!(preprocess(&[], 0, 0, 4).is_err()); + } + + #[test] + fn alpha_follows_the_logits() { + let logits: Vec = (0..16).map(|i| if i % 4 < 2 { 10.0 } else { -10.0 }).collect(); + let a = alpha(&logits, 4, 8, 2); + assert!(a[0] > 0.99 && a[7] < 0.01); + assert!(alpha(&[], 4, 3, 3).iter().all(|v| *v == 0.0)); + } +} diff --git a/crates/segment/src/neck.rs b/crates/segment/src/neck.rs new file mode 100644 index 00000000..c4f716dd --- /dev/null +++ b/crates/segment/src/neck.rs @@ -0,0 +1,77 @@ +//! The feature pyramid (`Sam3VisionNeck`): the backbone's 72 × 72 map at four scales +//! (288, 144, 72 and 36 per side), each projected to 256 channels. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. + +use candle_core::{Module, Tensor}; +use candle_nn::{Conv2d, Conv2dConfig, ConvTranspose2d, ConvTranspose2dConfig, VarBuilder}; + +use crate::Result; +use crate::vit::HIDDEN; + +pub const FPN: usize = 256; + +enum Scale { + /// Two 2× transposed convolutions with GELU between (4×). + Up4(ConvTranspose2d, ConvTranspose2d), + Up2(ConvTranspose2d), + Same, + /// 2 × 2 max pooling (½×). + Down2, +} + +struct Level { + scale: Scale, + proj1: Conv2d, + proj2: Conv2d, +} + +pub struct Neck { + levels: Vec, +} + +impl Neck { + /// `vb` at `…neck` (detector) or `tracker_neck`. + pub fn new(vb: VarBuilder) -> Result { + let up = ConvTranspose2dConfig { stride: 2, ..Default::default() }; + let mut levels = Vec::new(); + for (i, factor) in [4.0f32, 2.0, 1.0, 0.5].into_iter().enumerate() { + let l = vb.pp(format!("fpn_layers.{i}")); + let s = l.pp("scale_layers"); + let (scale, inter) = if factor == 4.0 { + let a = candle_nn::conv_transpose2d(HIDDEN, HIDDEN / 2, 2, up, s.pp("0"))?; + let b = candle_nn::conv_transpose2d(HIDDEN / 2, HIDDEN / 4, 2, up, s.pp("2"))?; + (Scale::Up4(a, b), HIDDEN / 4) + } else if factor == 2.0 { + (Scale::Up2(candle_nn::conv_transpose2d(HIDDEN, HIDDEN / 2, 2, up, s.pp("0"))?), HIDDEN / 2) + } else if factor == 1.0 { + (Scale::Same, HIDDEN) + } else { + (Scale::Down2, HIDDEN) + }; + levels.push(Level { + scale, + proj1: candle_nn::conv2d(inter, FPN, 1, Conv2dConfig::default(), l.pp("proj1"))?, + proj2: candle_nn::conv2d(FPN, FPN, 3, Conv2dConfig { padding: 1, ..Default::default() }, l.pp("proj2"))?, + }); + } + Ok(Self { levels }) + } + + /// `x`: the backbone's `[1, 1024, 72, 72]`; `n`: how many levels (from the finest) to compute. + pub fn forward(&self, x: &Tensor, n: usize) -> Result> { + let mut out = Vec::with_capacity(n); + for l in self.levels.iter().take(n) { + let y = match &l.scale { + Scale::Up4(a, b) => b.forward(&a.forward(x)?.gelu_erf()?)?, + Scale::Up2(a) => a.forward(x)?, + Scale::Same => x.clone(), + Scale::Down2 => x.max_pool2d(2)?, + }; + out.push(l.proj2.forward(&l.proj1.forward(&y)?)?); + } + Ok(out) + } +} diff --git a/crates/segment/src/nn.rs b/crates/segment/src/nn.rs new file mode 100644 index 00000000..540196e5 --- /dev/null +++ b/crates/segment/src/nn.rs @@ -0,0 +1,242 @@ +//! Shared building blocks: attention, MLPs, layer norms, sine position encodings. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`, `models/sam3_tracker`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. + +use candle_core::{D, DType, Device, Module, Tensor}; +use candle_nn::{LayerNorm, Linear, VarBuilder}; + +use crate::Result; + +/// `nn.LayerNorm(d)` (PyTorch's default eps). +pub const TORCH_LN_EPS: f64 = 1e-5; + +pub fn ln(d: usize, eps: f64, vb: VarBuilder) -> Result { + Ok(candle_nn::layer_norm(d, eps, vb)?) +} + +pub fn linear(i: usize, o: usize, vb: VarBuilder) -> Result { + Ok(candle_nn::linear(i, o, vb)?) +} + +/// Largest attention-score block computed at once (elements): the 5184-token global layers +/// would otherwise need ~1.7 GB for all 16 heads. +const SCORE_BLOCK: usize = 64 << 20; + +/// Scaled dot-product attention on `[B, H, L, d]` tensors, with an optional additive bias +/// broadcastable to `[B, H, Lq, Lk]`. Heads are processed in groups to bound memory. +pub fn sdpa(q: &Tensor, k: &Tensor, v: &Tensor, bias: Option<&Tensor>, scale: f64) -> Result { + let (b, h, lq, _) = q.dims4()?; + let lk = k.dim(2)?; + let per_head = b.saturating_mul(lq).saturating_mul(lk).max(1); + let group = (SCORE_BLOCK / per_head).clamp(1, h); + let bias_heads = match bias { + Some(t) => t.dim(1)?, + None => 1, + }; + let mut outs = Vec::with_capacity(h.div_ceil(group)); + let mut start = 0; + while start < h { + let n = group.min(h - start); + let qs = q.narrow(1, start, n)?; + let ks = k.narrow(1, start, n)?; + let vs = v.narrow(1, start, n)?; + let mut att = (qs.matmul(&ks.t()?)? * scale)?; + if let Some(bias) = bias { + let bs = if bias_heads > 1 { bias.narrow(1, start, n)? } else { bias.clone() }; + att = att.broadcast_add(&bs)?; + } + let att = candle_nn::ops::softmax_last_dim(&att)?; + outs.push(att.matmul(&vs)?); + start += n; + } + Ok(Tensor::cat(&outs, 1)?) +} + +/// Multi-head attention with q/k/v/o projections (`Sam3Attention`, `Sam3TrackerAttention` +/// with `internal` = hidden / downsample rate). +pub struct Attention { + q: Linear, + k: Linear, + v: Linear, + o: Linear, + heads: usize, + head_dim: usize, +} + +impl Attention { + pub fn new(hidden: usize, internal: usize, heads: usize, vb: VarBuilder) -> Result { + Self::named(hidden, internal, heads, "o_proj", vb) + } + + /// With the output projection called `out` (CLIP's `out_proj`). + pub fn named(hidden: usize, internal: usize, heads: usize, out: &str, vb: VarBuilder) -> Result { + Ok(Self { + q: linear(hidden, internal, vb.pp("q_proj"))?, + k: linear(hidden, internal, vb.pp("k_proj"))?, + v: linear(hidden, internal, vb.pp("v_proj"))?, + o: linear(internal, hidden, vb.pp(out))?, + heads, + head_dim: internal / heads, + }) + } + + fn split(&self, x: &Tensor) -> Result { + let (b, l, _) = x.dims3()?; + Ok(x.reshape((b, l, self.heads, self.head_dim))?.transpose(1, 2)?.contiguous()?) + } + + /// The q, k and v projections of `x` (for attention variants that transform them). + pub fn qkv(&self, x: &Tensor) -> Result<(Tensor, Tensor, Tensor)> { + Ok((self.q.forward(x)?, self.k.forward(x)?, self.v.forward(x)?)) + } + + /// The output projection. + pub fn out(&self, x: &Tensor) -> Result { + Ok(self.o.forward(x)?) + } + + /// `q`, `k`, `v`: `[B, L, hidden]`; `bias` broadcastable to `[B, heads, Lq, Lk]`. + pub fn forward(&self, q: &Tensor, k: &Tensor, v: &Tensor, bias: Option<&Tensor>) -> Result { + let (b, lq, _) = q.dims3()?; + let q = self.split(&self.q.forward(q)?)?; + let k = self.split(&self.k.forward(k)?)?; + let v = self.split(&self.v.forward(v)?)?; + let o = sdpa(&q, &k, &v, bias, (self.head_dim as f64).powf(-0.5))?; + let o = o.transpose(1, 2)?.reshape((b, lq, self.heads * self.head_dim))?; + Ok(self.o.forward(&o)?) + } +} + +/// An additive key-padding bias `[1, 1, 1, L]` (0 for valid tokens, the f32 minimum for padding). +pub fn padding_bias(valid: &[bool], device: &Device) -> Result { + let v: Vec = valid.iter().map(|&ok| if ok { 0.0 } else { f32::MIN }).collect(); + Ok(Tensor::from_vec(v, (1, 1, 1, valid.len()), device)?) +} + +/// `Sam3MLP`: fc1 → activation → fc2. +pub struct Mlp { + fc1: Linear, + fc2: Linear, + gelu: bool, +} + +impl Mlp { + pub fn new(hidden: usize, inter: usize, gelu: bool, vb: VarBuilder) -> Result { + Ok(Self { fc1: linear(hidden, inter, vb.pp("fc1"))?, fc2: linear(inter, hidden, vb.pp("fc2"))?, gelu }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + let x = self.fc1.forward(x)?; + let x = if self.gelu { x.gelu_erf()? } else { x.relu()? }; + Ok(self.fc2.forward(&x)?) + } +} + +/// `Sam3DecoderMLP`: `layer1 … layerN` with ReLU between (2 or 3 layers). +pub struct DecoderMlp { + layers: Vec, +} + +impl DecoderMlp { + pub fn new(input: usize, hidden: usize, output: usize, n: usize, vb: VarBuilder) -> Result { + let mut layers = Vec::with_capacity(n); + for i in 0..n { + let (a, b) = (if i == 0 { input } else { hidden }, if i + 1 == n { output } else { hidden }); + layers.push(linear(a, b, vb.pp(format!("layer{}", i + 1)))?); + } + Ok(Self { layers }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + let mut x = x.clone(); + for (i, l) in self.layers.iter().enumerate() { + x = l.forward(&x)?; + if i + 1 < self.layers.len() { + x = x.relu()?; + } + } + Ok(x) + } +} + +/// `Sam3TrackerFeedForward`: proj_in → act → (layers.i → act)* → proj_out (→ sigmoid). +pub struct FeedForward { + proj_in: Linear, + layers: Vec, + proj_out: Linear, + sigmoid: bool, +} + +impl FeedForward { + pub fn new(input: usize, hidden: usize, output: usize, n: usize, sigmoid: bool, vb: VarBuilder) -> Result { + let mut layers = Vec::new(); + for i in 0..n.saturating_sub(2) { + layers.push(linear(hidden, hidden, vb.pp(format!("layers.{i}")))?); + } + Ok(Self { proj_in: linear(input, hidden, vb.pp("proj_in"))?, layers, proj_out: linear(hidden, output, vb.pp("proj_out"))?, sigmoid }) + } + + pub fn forward(&self, x: &Tensor) -> Result { + let mut x = self.proj_in.forward(x)?.relu()?; + for l in &self.layers { + x = l.forward(&x)?.relu()?; + } + let x = self.proj_out.forward(&x)?; + Ok(if self.sigmoid { candle_nn::ops::sigmoid(&x)? } else { x }) + } +} + +/// Layer norm over the channel axis of an NCHW tensor (`Sam3TrackerLayerNorm`, channels_first). +pub fn ln_channels(x: &Tensor, norm: &LayerNorm) -> Result { + let x = x.permute((0, 2, 3, 1))?; + Ok(norm.forward(&x)?.permute((0, 3, 1, 2))?.contiguous()?) +} + +/// The DETR sine embedding of an `h × w` grid, normalized (`Sam3SinePositionEmbedding` with +/// `normalize=True`, scale 2π): `[1, 2·feats, h, w]`, y features first. +pub fn sine_grid(h: usize, w: usize, feats: usize, device: &Device) -> Result { + let scale = 2.0 * std::f64::consts::PI; + let dim_t: Vec = (0..feats).map(|i| 10000f64.powf(2.0 * (i / 2) as f64 / feats as f64)).collect(); + let mut out = vec![0f32; 2 * feats * h * w]; + for y in 0..h { + let ye = (y + 1) as f64 / (h as f64 + 1e-6) * scale; + for x in 0..w { + let xe = (x + 1) as f64 / (w as f64 + 1e-6) * scale; + for (i, d) in dim_t.iter().enumerate() { + let (py, px) = (ye / d, xe / d); + let (vy, vx) = if i % 2 == 0 { (py.sin(), px.sin()) } else { (py.cos(), px.cos()) }; + if let Some(v) = out.get_mut((i * h + y) * w + x) { + *v = vy as f32; + } + if let Some(v) = out.get_mut(((feats + i) * h + y) * w + x) { + *v = vx as f32; + } + } + } + } + Ok(Tensor::from_vec(out, (1, 2 * feats, h, w), device)?) +} + +/// `inverse_sigmoid` with the reference's clamping (eps 1e-3). +pub fn inverse_sigmoid(x: &Tensor) -> Result { + let x = x.clamp(0f32, 1f32)?; + let x1 = x.clamp(1e-3f32, f32::MAX)?; + let x2 = (1.0 - &x)?.clamp(1e-3f32, f32::MAX)?; + Ok((x1 / x2)?.log()?) +} + +/// `[..., 4]` boxes (cx, cy, w, h) → (x0, y0, x1, y1). +pub fn cxcywh_to_xyxy(b: &Tensor) -> Result { + let cx = b.narrow(D::Minus1, 0, 1)?; + let cy = b.narrow(D::Minus1, 1, 1)?; + let w = (b.narrow(D::Minus1, 2, 1)? * 0.5)?; + let h = (b.narrow(D::Minus1, 3, 1)? * 0.5)?; + Ok(Tensor::cat(&[(&cx - &w)?, (&cy - &h)?, (&cx + &w)?, (&cy + &h)?], D::Minus1)?) +} + +/// A tensor's values as `f32` on the CPU. +pub fn to_vec(t: &Tensor) -> Result> { + Ok(t.to_dtype(DType::F32)?.flatten_all()?.to_vec1::()?) +} diff --git a/crates/segment/src/tokenizer.rs b/crates/segment/src/tokenizer.rs new file mode 100644 index 00000000..0819aa48 --- /dev/null +++ b/crates/segment/src/tokenizer.rs @@ -0,0 +1,177 @@ +//! The CLIP byte-pair-encoding tokenizer (lower-cased, `` word ends), padded to the +//! 32-token context SAM 3's text encoder takes. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (the CLIP tokenizer used by `models/sam3`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. + +use std::collections::HashMap; +use std::path::Path; + +use crate::{Error, Result}; + +pub const CONTEXT: usize = 32; +const BOS: u32 = 49406; +const EOS: u32 = 49407; + +pub struct Tokenizer { + vocab: HashMap, + ranks: HashMap<(String, String), usize>, + byte_char: [char; 256], +} + +/// GPT-2's reversible byte → printable character map. +fn bytes_to_unicode() -> [char; 256] { + let mut map = ['\0'; 256]; + let mut n = 0u32; + for b in 0..=255u8 { + let printable = (b'!'..=b'~').contains(&b) || (0xA1..=0xAC).contains(&b) || (0xAE..=0xFF).contains(&b); + let c = if printable { + u32::from(b) + } else { + n += 1; + 255 + n + }; + if let Some(slot) = map.get_mut(usize::from(b)) { + *slot = char::from_u32(c).unwrap_or('?'); + } + } + map +} + +impl Tokenizer { + /// From `vocab.json` and `merges.txt` in `dir`. + pub fn load(dir: &Path) -> Result { + let read = |name: &str| std::fs::read_to_string(dir.join(name)).map_err(|e| Error::Model(format!("{name}: {e}"))); + let vocab: HashMap = serde_json::from_str(&read("vocab.json")?).map_err(|e| Error::Model(format!("vocab.json: {e}")))?; + let merges = read("merges.txt")?; + Ok(Self::new(vocab, &merges)) + } + + pub fn new(vocab: HashMap, merges: &str) -> Self { + let ranks = merges + .lines() + .filter(|l| !l.starts_with("#version") && !l.trim().is_empty()) + .filter_map(|l| l.split_once(' ').map(|(a, b)| (a.to_string(), b.to_string()))) + .enumerate() + .map(|(i, p)| (p, i)) + .collect(); + Self { vocab, ranks, byte_char: bytes_to_unicode() } + } + + /// Token ids (start, words, end, padding) and which positions are real tokens. + pub fn encode(&self, text: &str) -> (Vec, Vec) { + let mut ids = vec![BOS]; + for word in pre_tokenize(&text.split_whitespace().collect::>().join(" ").to_lowercase()) { + let mapped: String = word.bytes().map(|b| self.byte_char.get(usize::from(b)).copied().unwrap_or('?')).collect(); + for piece in self.bpe(&mapped) { + ids.push(self.vocab.get(&piece).copied().unwrap_or(EOS)); + } + } + ids.truncate(CONTEXT - 1); + ids.push(EOS); + let n = ids.len(); + ids.resize(CONTEXT, EOS); + let valid = (0..CONTEXT).map(|i| i < n).collect(); + (ids, valid) + } + + fn bpe(&self, word: &str) -> Vec { + let mut syms: Vec = word.chars().map(String::from).collect(); + if let Some(last) = syms.last_mut() { + last.push_str(""); + } + loop { + let best = syms + .windows(2) + .enumerate() + .filter_map(|(i, w)| match w { + [a, b] => self.ranks.get(&(a.clone(), b.clone())).map(|r| (*r, i)), + _ => None, + }) + .min(); + let Some((_, i)) = best else { break }; + let (Some(a), Some(b)) = (syms.get(i).cloned(), syms.get(i + 1).cloned()) else { break }; + let merged = a.clone() + &b; + // merge every occurrence of the pair, left to right + let mut out = Vec::with_capacity(syms.len()); + let mut j = 0; + while j < syms.len() { + if j + 1 < syms.len() && syms.get(j) == Some(&a) && syms.get(j + 1) == Some(&b) { + out.push(merged.clone()); + j += 2; + } else { + out.extend(syms.get(j).cloned()); + j += 1; + } + } + syms = out; + if syms.len() == 1 { + break; + } + } + syms + } +} + +/// CLIP's pre-tokenizer: contractions, letter runs, single digits, and runs of other +/// non-space characters. +fn pre_tokenize(text: &str) -> Vec { + let chars: Vec = text.chars().collect(); + let mut out = Vec::new(); + let mut i = 0; + while let Some(&c) = chars.get(i) { + if c.is_whitespace() { + i += 1; + continue; + } + if c == '\'' { + let rest: String = chars.get(i + 1..(i + 3).min(chars.len())).map(|s| s.iter().collect()).unwrap_or_default(); + if let Some(k) = ["re", "ve", "ll"] + .iter() + .find(|k| rest.starts_with(**k)) + .map(|k| k.len()) + .or_else(|| rest.chars().next().filter(|c| matches!(c, 's' | 't' | 'm' | 'd')).map(|_| 1)) + { + out.push(chars.get(i..i + 1 + k).map(|s| s.iter().collect()).unwrap_or_default()); + i += 1 + k; + continue; + } + } + let start = i; + if c.is_alphabetic() { + while chars.get(i).is_some_and(|c| c.is_alphabetic()) { + i += 1; + } + } else if c.is_numeric() { + i += 1; + } else { + while chars.get(i).is_some_and(|c| !c.is_whitespace() && !c.is_alphabetic() && !c.is_numeric()) { + i += 1; + } + } + out.push(chars.get(start..i).map(|s| s.iter().collect()).unwrap_or_default()); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn splits_like_clip() { + assert_eq!(pre_tokenize("a dog's toy, 25 cats!"), ["a", "dog", "'s", "toy", ",", "2", "5", "cats", "!"]); + } + + #[test] + fn merges_by_rank() { + let vocab: HashMap = + [("c", 1), ("a", 2), ("t", 3), ("ca", 4), ("cat", 5)].into_iter().map(|(k, v)| (k.to_string(), v)).collect(); + let t = Tokenizer::new(vocab, "#version: 0.2\nc a\nca t\n"); + let (ids, valid) = t.encode(" CAT "); + assert_eq!(&ids[..3], &[BOS, 5, EOS]); + assert_eq!(ids.len(), CONTEXT); + assert_eq!(valid.iter().filter(|v| **v).count(), 3); + } +} diff --git a/crates/segment/src/tracker.rs b/crates/segment/src/tracker.rs new file mode 100644 index 00000000..d51219c8 --- /dev/null +++ b/crates/segment/src/tracker.rs @@ -0,0 +1,253 @@ +//! Point prompts (SAM 3's interactive "tracker" head, SAM 2 style): a prompt encoder for +//! positive/negative clicks and a two-way-transformer mask decoder over the image features. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3_tracker/modeling_sam3_tracker.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. + +use candle_core::{DType, Device, IndexOp, Module, Tensor}; +use candle_nn::{Conv2d, Conv2dConfig, ConvTranspose2d, ConvTranspose2dConfig, LayerNorm, VarBuilder}; + +use crate::nn::{Attention, FeedForward, TORCH_LN_EPS, ln, ln_channels}; +use crate::vit::{GRID, IMAGE}; +use crate::{Result, nn::to_vec}; + +const HIDDEN: usize = 256; +const MASK_TOKENS: usize = 4; +/// Low-resolution mask side (4 × the 72 grid). +pub const LOW_RES: usize = 288; + +/// The image features the decoder needs, computed once per image. +pub struct Features { + /// Level 2 (72 × 72, 256 channels). + pub image: Tensor, + /// `conv_s0`(level 0): `[1, 32, 288, 288]`. + pub s0: Tensor, + /// `conv_s1`(level 1): `[1, 64, 144, 144]`. + pub s1: Tensor, +} + +/// One click: position normalized to the image (0..1), and whether it includes or excludes. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Click { + pub x: f32, + pub y: f32, + pub positive: bool, +} + +/// A predicted mask: 288 × 288 logits over the (squared) model input, and its quality. +pub struct Prediction { + pub logits: Vec, + pub iou: f32, + pub object_score: f32, +} + +/// Random-Fourier positional encoding of points normalized to 0..1 (`Sam3TrackerPositionalEmbedding`). +struct Fourier { + pe: Vec, // [2, 128] +} + +impl Fourier { + fn encode(&self, x: f32, y: f32) -> Vec { + let half = self.pe.len() / 2; + let (cx, cy) = (2.0 * x - 1.0, 2.0 * y - 1.0); + let proj: Vec = (0..half) + .map(|i| { + let a = cx * self.pe.get(i).copied().unwrap_or(0.0) + cy * self.pe.get(half + i).copied().unwrap_or(0.0); + 2.0 * std::f32::consts::PI * a + }) + .collect(); + proj.iter().map(|a| a.sin()).chain(proj.iter().map(|a| a.cos())).collect() + } +} + +struct TwoWayBlock { + self_attn: Attention, + ln1: LayerNorm, + t2i: Attention, + ln2: LayerNorm, + mlp: FeedForward, + ln3: LayerNorm, + ln4: LayerNorm, + i2t: Attention, + skip_first_pe: bool, +} + +/// `Sam3TrackerAttention` on `[1, L, C]` (the point batch of one is folded away). +fn attn(a: &Attention, q: &Tensor, k: &Tensor, v: &Tensor) -> Result { + a.forward(q, k, v, None) +} + +impl TwoWayBlock { + fn forward(&self, queries: &Tensor, keys: &Tensor, qpe: &Tensor, kpe: &Tensor) -> Result<(Tensor, Tensor)> { + let queries = if self.skip_first_pe { + attn(&self.self_attn, queries, queries, queries)? + } else { + let q = (queries + qpe)?; + (queries + attn(&self.self_attn, &q, &q, queries)?)? + }; + let queries = self.ln1.forward(&queries)?; + let (q, k) = ((&queries + qpe)?, (keys + kpe)?); + let queries = self.ln2.forward(&(&queries + attn(&self.t2i, &q, &k, keys)?)?)?; + let queries = self.ln3.forward(&(&queries + self.mlp.forward(&queries)?)?)?; + let (q, k) = ((&queries + qpe)?, (keys + kpe)?); + let keys = self.ln4.forward(&(keys + attn(&self.i2t, &k, &q, &queries)?)?)?; + Ok((queries, keys)) + } +} + +pub struct Tracker { + device: Device, + point_pe: Fourier, + image_pe: Tensor, + point_embed: Tensor, + not_a_point: Tensor, + no_mask: Tensor, + conv_s0: Conv2d, + conv_s1: Conv2d, + output_tokens: Tensor, + blocks: Vec, + final_t2i: Attention, + ln_final: LayerNorm, + up1: ConvTranspose2d, + up_ln: LayerNorm, + up2: ConvTranspose2d, + hyper: Vec, + iou_head: FeedForward, + obj_head: FeedForward, +} + +fn read_vec(vb: &VarBuilder, shape: (usize, usize), name: &str) -> Result> { + to_vec(&vb.get(shape, name)?) +} + +impl Tracker { + /// `vb` at `tracker_model`. + pub fn new(vb: VarBuilder) -> Result { + let device = vb.device().clone(); + let pe = vb.pp("prompt_encoder"); + let md = vb.pp("mask_decoder"); + let image_fourier = Fourier { pe: read_vec(&vb, (2, HIDDEN / 2), "shared_image_embedding.positional_embedding")? }; + // image-wide positional embedding: pixel centres of the 72 × 72 grid + let mut grid = Vec::with_capacity(GRID * GRID * HIDDEN); + for r in 0..GRID { + for c in 0..GRID { + grid.extend(image_fourier.encode((c as f32 + 0.5) / GRID as f32, (r as f32 + 0.5) / GRID as f32)); + } + } + let image_pe = Tensor::from_vec(grid, (1, GRID * GRID, HIDDEN), &device)?; + let t = md.pp("transformer"); + let mut blocks = Vec::new(); + for i in 0..2 { + let b = t.pp(format!("layers.{i}")); + blocks.push(TwoWayBlock { + self_attn: Attention::new(HIDDEN, HIDDEN, 8, b.pp("self_attn"))?, + ln1: ln(HIDDEN, TORCH_LN_EPS, b.pp("layer_norm1"))?, + t2i: Attention::new(HIDDEN, HIDDEN / 2, 8, b.pp("cross_attn_token_to_image"))?, + ln2: ln(HIDDEN, TORCH_LN_EPS, b.pp("layer_norm2"))?, + mlp: FeedForward::new(HIDDEN, 2048, HIDDEN, 2, false, b.pp("mlp"))?, + ln3: ln(HIDDEN, TORCH_LN_EPS, b.pp("layer_norm3"))?, + ln4: ln(HIDDEN, TORCH_LN_EPS, b.pp("layer_norm4"))?, + i2t: Attention::new(HIDDEN, HIDDEN / 2, 8, b.pp("cross_attn_image_to_token"))?, + skip_first_pe: i == 0, + }); + } + let output_tokens = Tensor::cat( + &[ + md.get((1, HIDDEN), "obj_score_token.weight")?, + md.get((1, HIDDEN), "iou_token.weight")?, + md.get((MASK_TOKENS, HIDDEN), "mask_tokens.weight")?, + ], + 0, + )?; + let up = ConvTranspose2dConfig { stride: 2, ..Default::default() }; + let one = Conv2dConfig::default(); + Ok(Self { + point_pe: Fourier { pe: read_vec(&pe, (2, HIDDEN / 2), "shared_embedding.positional_embedding")? }, + image_pe, + point_embed: pe.get((4, HIDDEN), "point_embed.weight")?.to_device(&Device::Cpu)?, + not_a_point: pe.get((1, HIDDEN), "not_a_point_embed.weight")?.to_device(&Device::Cpu)?, + no_mask: pe.get((1, HIDDEN), "no_mask_embed.weight")?.reshape((1, HIDDEN, 1, 1))?, + conv_s0: candle_nn::conv2d(HIDDEN, HIDDEN / 8, 1, one, md.pp("conv_s0"))?, + conv_s1: candle_nn::conv2d(HIDDEN, HIDDEN / 4, 1, one, md.pp("conv_s1"))?, + output_tokens, + blocks, + final_t2i: Attention::new(HIDDEN, HIDDEN / 2, 8, t.pp("final_attn_token_to_image"))?, + ln_final: ln(HIDDEN, TORCH_LN_EPS, t.pp("layer_norm_final_attn"))?, + up1: candle_nn::conv_transpose2d(HIDDEN, HIDDEN / 4, 2, up, md.pp("upscale_conv1"))?, + up_ln: ln(HIDDEN / 4, 1e-6, md.pp("upscale_layer_norm"))?, + up2: candle_nn::conv_transpose2d(HIDDEN / 4, HIDDEN / 8, 2, up, md.pp("upscale_conv2"))?, + hyper: (0..MASK_TOKENS) + .map(|i| FeedForward::new(HIDDEN, HIDDEN, HIDDEN / 8, 3, false, md.pp(format!("output_hypernetworks_mlps.{i}")))) + .collect::>()?, + iou_head: FeedForward::new(HIDDEN, HIDDEN, MASK_TOKENS, 3, true, md.pp("iou_prediction_head"))?, + obj_head: FeedForward::new(HIDDEN, HIDDEN, 1, 3, false, md.pp("pred_obj_score_head"))?, + device, + }) + } + + /// The decoder's image features from the tracker neck's levels (288, 144, 72 per side). + pub fn features(&self, levels: &[Tensor]) -> Result { + let [l0, l1, l2, ..] = levels else { return Err(crate::Error::Model("tracker needs three feature levels".into())) }; + Ok(Features { image: l2.clone(), s0: self.conv_s0.forward(l0)?, s1: self.conv_s1.forward(l1)? }) + } + + /// Sparse prompt embeddings `[1, n + 1, 256]` (a padding point closes the list). + fn embed_clicks(&self, clicks: &[Click]) -> Result { + let point_embed = to_vec(&self.point_embed)?; + let not_a_point = to_vec(&self.not_a_point)?; + let mut rows = Vec::with_capacity((clicks.len() + 1) * HIDDEN); + for c in clicks { + // the reference shifts pixel coordinates (in the 1008² input) by half a pixel + let (x, y) = ((c.x * IMAGE as f32 + 0.5) / IMAGE as f32, (c.y * IMAGE as f32 + 0.5) / IMAGE as f32); + let e = self.point_pe.encode(x, y); + let label = usize::from(c.positive); + rows.extend(e.iter().enumerate().map(|(i, v)| v + point_embed.get(label * HIDDEN + i).copied().unwrap_or(0.0))); + } + rows.extend_from_slice(¬_a_point); + Ok(Tensor::from_vec(rows, (1, clicks.len() + 1, HIDDEN), &self.device)?) + } + + /// Predict masks for `clicks`. `multimask`: the three alternatives (best by predicted IoU is + /// returned), else the single-mask output with SAM 2's stability fallback. + pub fn predict(&self, f: &Features, clicks: &[Click], multimask: bool) -> Result { + let sparse = self.embed_clicks(clicks)?; + let tokens = Tensor::cat(&[self.output_tokens.unsqueeze(0)?, sparse], 1)?; // [1, T, 256] + let image = f.image.broadcast_add(&self.no_mask)?; // dense prompt: "no mask" + let (b, c, h, w) = image.dims4()?; + let keys = image.flatten_from(2)?.transpose(1, 2)?.contiguous()?; // [1, HW, 256] + let (mut q, mut k) = (tokens.clone(), keys); + for blk in &self.blocks { + (q, k) = blk.forward(&q, &k, &tokens, &self.image_pe)?; + } + let qf = (&q + &tokens)?; + let kf = (&k + &self.image_pe)?; + let q = self.ln_final.forward(&(&q + attn(&self.final_t2i, &qf, &kf, &k)?)?)?; + let iou_token = q.i((.., 1, ..))?; + let image = k.transpose(1, 2)?.reshape((b, c, h, w))?; + let up = (self.up1.forward(&image)? + &f.s1)?; + let up = ln_channels(&up, &self.up_ln)?.gelu_erf()?; + let up = (self.up2.forward(&up)? + &f.s0)?.gelu_erf()?; // [1, 32, 288, 288] + let (_, uc, uh, uw) = up.dims4()?; + let hyper = (0..MASK_TOKENS) + .map(|i| self.hyper.get(i).map_or(Err(crate::Error::Model("hypernetwork".into())), |m| m.forward(&q.i((.., 2 + i, ..))?))) + .collect::>>()?; + let hyper = Tensor::stack(&hyper, 1)?; // [1, 4, 32] + let masks = hyper.matmul(&up.reshape((b, uc, uh * uw))?)?.reshape((b, MASK_TOKENS, uh, uw))?; + let iou = to_vec(&self.iou_head.forward(&iou_token)?)?; + let object_score = to_vec(&self.obj_head.forward(&q.i((.., 0, ..))?)?)?.first().copied().unwrap_or(0.0); + let pick = |i: usize| -> Result { + Ok(Prediction { logits: to_vec(&masks.i((0, i))?)?, iou: iou.get(i).copied().unwrap_or(0.0), object_score }) + }; + let best_multi = || (1..MASK_TOKENS).max_by(|a, b| iou.get(*a).partial_cmp(&iou.get(*b)).unwrap_or(std::cmp::Ordering::Equal)).unwrap_or(1); + if multimask { + return pick(best_multi()); + } + // single-mask output unless it is unstable (area changes > 2 % when the threshold moves ±0.05) + let single = masks.i((0, 0))?; + let inter = single.gt(0.05)?.to_dtype(DType::F32)?.sum_all()?.to_scalar::()?; + let union = single.gt(-0.05)?.to_dtype(DType::F32)?.sum_all()?.to_scalar::()?; + let stability = if union > 0.0 { inter / union } else { 1.0 }; + if stability >= 0.98 { pick(0) } else { pick(best_multi()) } + } +} diff --git a/crates/segment/src/vit.rs b/crates/segment/src/vit.rs new file mode 100644 index 00000000..16b09a72 --- /dev/null +++ b/crates/segment/src/vit.rs @@ -0,0 +1,170 @@ +//! The SAM 3 vision backbone: a 32-layer ViT (Perception Encoder) with 2-D axial rotary +//! position encoding, windowed attention (24 × 24) in most layers and global attention in four. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. + +use candle_core::{D, Device, Module, Tensor}; +use candle_nn::{Conv2d, Conv2dConfig, LayerNorm, VarBuilder}; + +use crate::nn::{Attention, Mlp, ln, sdpa}; +use crate::{Error, Result}; + +pub const IMAGE: usize = 1008; +pub const PATCH: usize = 14; +/// Patches per side (72). +pub const GRID: usize = IMAGE / PATCH; +pub const HIDDEN: usize = 1024; +const HEADS: usize = 16; +const HEAD_DIM: usize = HIDDEN / HEADS; +const LAYERS: usize = 32; +const WINDOW: usize = 24; +const GLOBAL: [usize; 4] = [7, 15, 23, 31]; +const PRETRAIN_GRID: usize = 336 / PATCH; +const EPS: f64 = 1e-6; + +/// Cosine and sine tables `[N, head_dim]` for a `side × side` grid whose coordinates are scaled +/// by `scale` (`Sam3ViTRotaryEmbedding`: x frequencies then y, each duplicated pairwise). +fn rope_tables(side: usize, scale: f64, device: &Device) -> Result<(Tensor, Tensor)> { + let spatial = HEAD_DIM / 2; + let inv: Vec = (0..spatial).step_by(2).map(|i| 1.0 / 10000f64.powf(i as f64 / spatial as f64)).collect(); + let n = side * side; + let (mut cos, mut sin) = (Vec::with_capacity(n * HEAD_DIM), Vec::with_capacity(n * HEAD_DIM)); + for r in 0..side { + for c in 0..side { + let (x, y) = (c as f64 * scale, r as f64 * scale); + for pos in [x, y] { + for f in &inv { + let a = (pos * f) as f32 as f64; + let (s, co) = (a.sin() as f32, a.cos() as f32); + cos.extend([co, co]); + sin.extend([s, s]); + } + } + } + } + Ok((Tensor::from_vec(cos, (n, HEAD_DIM), device)?, Tensor::from_vec(sin, (n, HEAD_DIM), device)?)) +} + +/// `x·cos + rotate_pairwise(x)·sin`, where rotate_pairwise maps (a, b) pairs to (−b, a). +fn apply_rope(x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result { + let dims = x.dims().to_vec(); + let last = *dims.last().ok_or_else(|| Error::Model("rope on a scalar".into()))?; + let mut paired = dims.clone(); + paired.pop(); + paired.extend([last / 2, 2]); + let p = x.reshape(paired.as_slice())?; + let a = p.narrow(D::Minus1, 0, 1)?; + let b = p.narrow(D::Minus1, 1, 1)?; + let rot = Tensor::cat(&[b.neg()?, a], D::Minus1)?.reshape(dims.as_slice())?; + Ok((x.broadcast_mul(cos)? + rot.broadcast_mul(sin)?)?) +} + +struct Layer { + ln1: LayerNorm, + attn: Attention, + ln2: LayerNorm, + mlp: Mlp, + window: usize, +} + +impl Layer { + /// `x`: `[B, H, W, C]`. + fn forward(&self, x: &Tensor, rope: &(Tensor, Tensor)) -> Result { + let (b, h, w, c) = x.dims4()?; + let y = self.ln1.forward(x)?; + let y = match (h.checked_div(self.window), w.checked_div(self.window)) { + (Some(nh), Some(nw)) => { + // non-overlapping windows (72 is a multiple of 24: no padding) + let ws = self.window; + let y = y.reshape((b, nh, ws, nw, ws, c))?.permute((0, 1, 3, 2, 4, 5))?.reshape((b * nh * nw, ws, ws, c))?; + let y = self.attend(&y, rope)?; + y.reshape((b, nh, nw, ws, ws, c))?.permute((0, 1, 3, 2, 4, 5))?.reshape((b, h, w, c))? + } + // global attention + _ => self.attend(&y, rope)?, + }; + let x = (x + y)?; + let y = self.mlp.forward(&self.ln2.forward(&x)?)?; + Ok((x + y)?) + } + + /// RoPE self-attention over `[B, H, W, C]`. + fn attend(&self, x: &Tensor, rope: &(Tensor, Tensor)) -> Result { + let (b, h, w, c) = x.dims4()?; + let seq = x.reshape((b, h * w, c))?; + Ok(self.attn_rope(&seq, rope)?.reshape((b, h, w, c))?) + } + + fn attn_rope(&self, x: &Tensor, rope: &(Tensor, Tensor)) -> Result { + let (b, n, _) = x.dims3()?; + let split = |t: Tensor| -> Result { Ok(t.reshape((b, n, HEADS, HEAD_DIM))?.transpose(1, 2)?.contiguous()?) }; + let (q, k, v) = self.attn.qkv(x)?; + let q = apply_rope(&split(q)?, &rope.0, &rope.1)?; + let k = apply_rope(&split(k)?, &rope.0, &rope.1)?; + let v = split(v)?; + let o = sdpa(&q, &k, &v, None, (HEAD_DIM as f64).powf(-0.5))?; + let o = o.transpose(1, 2)?.reshape((b, n, HIDDEN))?; + self.attn.out(&o) + } +} + +pub struct Vit { + patch: Conv2d, + pos: Tensor, + ln_pre: LayerNorm, + layers: Vec, + rope_window: (Tensor, Tensor), + rope_global: (Tensor, Tensor), +} + +impl Vit { + /// `vb` at `detector_model.vision_encoder.backbone`. + pub fn new(vb: VarBuilder) -> Result { + let device = vb.device().clone(); + let cfg = Conv2dConfig { stride: PATCH, ..Default::default() }; + let patch = candle_nn::conv2d_no_bias(3, HIDDEN, PATCH, cfg, vb.pp("embeddings.patch_embeddings.projection"))?; + let pos = vb.get((1, PRETRAIN_GRID * PRETRAIN_GRID, HIDDEN), "embeddings.position_embeddings")?; + // tile (not interpolate) the 24 × 24 pretraining table over the 72 × 72 grid + let reps = GRID.div_ceil(PRETRAIN_GRID); + let pos = pos + .reshape((1, PRETRAIN_GRID, PRETRAIN_GRID, HIDDEN))? + .repeat((1, reps, reps, 1))? + .narrow(1, 0, GRID)? + .narrow(2, 0, GRID)? + .reshape((1, GRID * GRID, HIDDEN))?; + let mut layers = Vec::with_capacity(LAYERS); + for i in 0..LAYERS { + let l = vb.pp(format!("layers.{i}")); + layers.push(Layer { + ln1: ln(HIDDEN, EPS, l.pp("layer_norm1"))?, + attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("attention"))?, + ln2: ln(HIDDEN, EPS, l.pp("layer_norm2"))?, + mlp: Mlp::new(HIDDEN, 4736, true, l.pp("mlp"))?, + window: if GLOBAL.contains(&i) { 0 } else { WINDOW }, + }); + } + Ok(Self { + patch, + pos, + ln_pre: ln(HIDDEN, EPS, vb.pp("layer_norm"))?, + layers, + rope_window: rope_tables(WINDOW, 1.0, &device)?, + rope_global: rope_tables(GRID, WINDOW as f64 / GRID as f64, &device)?, + }) + } + + /// `pixels`: `[1, 3, 1008, 1008]` normalized to −1..1. Returns `[1, 1024, 72, 72]`. + pub fn forward(&self, pixels: &Tensor) -> Result { + let x = self.patch.forward(pixels)?; // [1, C, 72, 72] + let b = x.dim(0)?; + let x = x.flatten_from(2)?.transpose(1, 2)?.broadcast_add(&self.pos)?; + let mut x = self.ln_pre.forward(&x.reshape((b, GRID, GRID, HIDDEN))?)?; + for l in &self.layers { + let rope = if l.window > 0 { &self.rope_window } else { &self.rope_global }; + x = l.forward(&x, rope)?; + } + Ok(x.permute((0, 3, 1, 2))?.contiguous()?) + } +} diff --git a/crates/segment/src/weights.rs b/crates/segment/src/weights.rs new file mode 100644 index 00000000..70ac92b2 --- /dev/null +++ b/crates/segment/src/weights.rs @@ -0,0 +1,151 @@ +//! Safetensors weights read on demand. +//! +//! Tensors are read from the file when a module asks for them (positional reads, no memory +//! map: `unsafe` stays in `lightcraft-sysmem`), so building only the point-prompt model never +//! touches the text encoder's bytes. + +use std::collections::HashMap; +use std::fs::File; +use std::io::{Read, Seek, SeekFrom}; +use std::path::Path; +use std::sync::{Arc, Mutex}; + +use candle_core::{DType, Device, Shape, Tensor}; +use candle_nn::var_builder::SimpleBackend; + +use crate::{Error, Result}; + +struct Entry { + shape: Vec, + start: u64, + end: u64, +} + +/// An open safetensors file: its header, and the file to read tensor bytes from. +pub struct Weights { + file: Mutex, + entries: HashMap, + data_start: u64, +} + +impl Weights { + pub fn open(path: &Path) -> Result> { + let mut file = File::open(path).map_err(|e| Error::Model(format!("{}: {e}", path.display())))?; + let file_len = file.metadata().map_err(|e| Error::Model(format!("{}: {e}", path.display())))?.len(); + let mut len = [0u8; 8]; + file.read_exact(&mut len).map_err(|e| Error::Model(format!("{}: {e}", path.display())))?; + let n = u64::from_le_bytes(len); + if n > 100 << 20 || n.saturating_add(8) > file_len { + return Err(Error::Model(format!("{}: not a safetensors file (header of {n} bytes)", path.display()))); + } + let mut header = vec![0u8; n as usize]; + file.read_exact(&mut header).map_err(|e| Error::Model(format!("{}: {e}", path.display())))?; + let header: serde_json::Map = + serde_json::from_slice(&header).map_err(|e| Error::Model(format!("{}: bad header: {e}", path.display())))?; + let mut entries = HashMap::new(); + for (name, v) in header { + if name == "__metadata__" { + continue; + } + let dtype = match v.get("dtype").and_then(|d| d.as_str()) { + Some("F32") => DType::F32, + other => return Err(Error::Model(format!("{name}: unsupported dtype {other:?} (the SAM 3 checkpoint is F32)"))), + }; + let shape: Vec = v + .get("shape") + .and_then(|s| s.as_array()) + .map(|a| a.iter().filter_map(|x| x.as_u64().map(|x| x as usize)).collect()) + .unwrap_or_default(); + let offsets: Vec = + v.get("data_offsets").and_then(|s| s.as_array()).map(|a| a.iter().filter_map(|x| x.as_u64()).collect()).unwrap_or_default(); + let [start, end] = offsets[..] else { return Err(Error::Model(format!("{name}: bad data_offsets"))) }; + let want = shape.iter().try_fold(dtype.size_in_bytes(), |a, d| a.checked_mul(*d)); + if end.checked_sub(start) != want.map(|w| w as u64) { + return Err(Error::Model(format!("{name}: size does not match its shape"))); + } + // the bytes must be in the file (a damaged or truncated download fails here, not later) + if (8 + n).checked_add(end).is_none_or(|e| e > file_len) { + return Err(Error::Model(format!("{}: tensor {name} lies outside the file (incomplete or damaged download?)", path.display()))); + } + entries.insert(name, Entry { shape, start, end }); + } + Ok(Arc::new(Self { file: Mutex::new(file), entries, data_start: 8 + n })) + } + + pub fn contains(&self, name: &str) -> bool { + self.entries.contains_key(name) + } + + /// Tensor `name` as stored, on the CPU. + pub fn read(&self, name: &str) -> Result { + let e = self.entries.get(name).ok_or_else(|| Error::Model(format!("missing tensor `{name}`")))?; + // checked against the file's length in `open` + let len = usize::try_from(e.end.saturating_sub(e.start)).map_err(|_| Error::Model(format!("{name}: too large")))?; + let mut bytes = vec![0u8; len]; + { + let mut f = self.file.lock().unwrap_or_else(std::sync::PoisonError::into_inner); + f.seek(SeekFrom::Start(self.data_start + e.start)).map_err(|err| Error::Model(format!("{name}: {err}")))?; + f.read_exact(&mut bytes).map_err(|err| Error::Model(format!("{name}: {err}")))?; + } + let v: Vec = bytes.as_chunks::<4>().0.iter().map(|c| f32::from_le_bytes(*c)).collect(); + Ok(Tensor::from_vec(v, e.shape.as_slice(), &Device::Cpu)?) + } +} + +/// [`Weights`] as a candle `VarBuilder` backend. +pub struct Backend(pub Arc); + +impl SimpleBackend for Backend { + fn get(&self, s: Shape, name: &str, _h: candle_nn::Init, dtype: DType, dev: &Device) -> candle_core::Result { + let t = self.get_unchecked(name, dtype, dev)?; + if t.shape() != &s { + candle_core::bail!("shape mismatch for {name}: expected {s:?}, got {:?}", t.shape()) + } + Ok(t) + } + + fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> candle_core::Result { + let t = self.0.read(name).map_err(|e| candle_core::Error::Msg(e.to_string()))?; + t.to_dtype(dtype)?.to_device(dev) + } + + fn contains_tensor(&self, name: &str) -> bool { + self.0.contains(name) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn write(name: &str, header: &str, data_len: usize) -> std::path::PathBuf { + let p = std::env::temp_dir().join(format!("lc-weights-{name}-{}.safetensors", std::process::id())); + let mut b = (header.len() as u64).to_le_bytes().to_vec(); + b.extend_from_slice(header.as_bytes()); + b.extend(std::iter::repeat_n(0u8, data_len)); + std::fs::write(&p, b).unwrap(); + p + } + + #[test] + fn tensors_must_lie_inside_the_file() { + let ok = write("ok", r#"{"a":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]}}"#, 16); + let w = Weights::open(&ok).unwrap(); + assert_eq!(w.read("a").unwrap().dims(), &[2, 2]); + // a truncated download: the header promises bytes that aren't there + let short = write("short", r#"{"a":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]}}"#, 8); + assert!(Weights::open(&short).err().unwrap().to_string().contains("outside the file")); + // hostile offsets: huge, or beyond u64 + let huge = write("huge", r#"{"a":{"dtype":"F32","shape":[1073741824,1073741824],"data_offsets":[0,4611686018427387904]}}"#, 16); + assert!(Weights::open(&huge).is_err()); + let wrap = write("wrap", r#"{"a":{"dtype":"F32","shape":[1],"data_offsets":[18446744073709551612,18446744073709551616]}}"#, 16); + assert!(Weights::open(&wrap).is_err()); + // a header longer than the file + let p = std::env::temp_dir().join(format!("lc-weights-hdr-{}.safetensors", std::process::id())); + std::fs::write(&p, 1000u64.to_le_bytes()).unwrap(); + assert!(Weights::open(&p).is_err()); + for f in [ok, short, huge, wrap, p] { + let _ = std::fs::remove_file(f); + } + } +} diff --git a/crates/segment/tests/reference.rs b/crates/segment/tests/reference.rs new file mode 100644 index 00000000..a73fcaa7 --- /dev/null +++ b/crates/segment/tests/reference.rs @@ -0,0 +1,97 @@ +//! The port against the reference implementation (Hugging Face `transformers`, fp32). +//! +//! Needs the checkpoint and a reference file made by `tools/sam3_reference.py`: +//! `LIGHTCRAFT_SAM3_DIR= LIGHTCRAFT_SAM3_REF= cargo test -p lightcraft-segment --release -- --nocapture`. +//! Skipped (passes) when either is missing. + +use std::collections::HashMap; +use std::path::PathBuf; + +use candle_core::{Device, IndexOp, Tensor}; +use lightcraft_segment::{Click, Sam3, tokenizer::Tokenizer}; + +fn setup() -> Option<(PathBuf, HashMap)> { + let dir = PathBuf::from(std::env::var_os("LIGHTCRAFT_SAM3_DIR")?); + let reference = PathBuf::from(std::env::var_os("LIGHTCRAFT_SAM3_REF")?); + let r = candle_core::safetensors::load(&reference, &Device::Cpu).ok()?; + Some((dir, r)) +} + +fn vec(t: &Tensor) -> Vec { + t.to_device(&Device::Cpu).unwrap().flatten_all().unwrap().to_vec1::().unwrap() +} + +/// Max |a − b| relative to the reference's max |b|. +fn rel_err(name: &str, a: &[f32], b: &[f32]) -> f32 { + assert_eq!(a.len(), b.len(), "{name}: length"); + let scale = b.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6); + let err = a.iter().zip(b).fold(0f32, |m, (x, y)| m.max((x - y).abs())); + eprintln!("{name:>18}: max abs err {err:.2e}, rel {:.2e} (scale {scale:.2e})", err / scale); + err / scale +} + +/// Fraction of pixels where the two logit maps disagree on the mask (> 0). +fn mask_disagreement(a: &[f32], b: &[f32]) -> f32 { + a.iter().zip(b).filter(|(x, y)| (**x > 0.0) != (**y > 0.0)).count() as f32 / a.len() as f32 +} + +#[test] +fn matches_reference() { + let Some((dir, r)) = setup() else { + eprintln!("skipped: set LIGHTCRAFT_SAM3_DIR and LIGHTCRAFT_SAM3_REF"); + return; + }; + let mut model = Sam3::load(&dir).unwrap(); + eprintln!("device: {:?}", model.device()); + let t = std::time::Instant::now(); + let mut enc = model.encode_pixels(&r["pixel_values"]).unwrap(); + let backbone = enc.backbone().permute((0, 2, 3, 1)).unwrap(); + assert!(rel_err("backbone", &vec(&backbone), &vec(&r["backbone"])) < 1e-3); + eprintln!("backbone: {:?}", t.elapsed()); + + // points: two clicks (single-mask output), one click (best of three) + let clicks = [Click { x: 0.5, y: 0.55, positive: true }, Click { x: 0.2, y: 0.2, positive: false }]; + let t = std::time::Instant::now(); + let p = model.segment_clicks(&mut enc, &clicks).unwrap(); + eprintln!("clicks (incl. tracker load + features): {:?}", t.elapsed()); + let want = vec(&r["trk_masks"]); + assert!(rel_err("2-click logits", &p.logits, &want) < 1e-2); + assert!(mask_disagreement(&p.logits, &want) < 1e-3); + let t = std::time::Instant::now(); + let p = model.segment_clicks(&mut enc, &clicks[..1]).unwrap(); + eprintln!("one more click: {:?}", t.elapsed()); + let iou = vec(&r["trk_multi_iou"]); + let best = (0..3).max_by(|a, b| iou[*a].partial_cmp(&iou[*b]).unwrap()).unwrap(); + let want = vec(&r["trk_multi_masks"].i((0, 0, best)).unwrap()); + assert!(rel_err("1-click logits", &p.logits, &want) < 1e-2); + assert!(mask_disagreement(&p.logits, &want) < 1e-3); + + // tokenizer + let tok = Tokenizer::load(&dir).unwrap(); + let cases = ["person", "Red Car", "a dog's toy", "sky, clouds!", "café 2 people", "tree-line"]; + for (i, c) in cases.iter().enumerate() { + let (ids, _) = tok.encode(c); + let want: Vec = r[&format!("tok{i}")].flatten_all().unwrap().to_vec1().unwrap(); + assert_eq!(ids.iter().map(|v| *v as i64).collect::>(), want, "tokens of {c:?}"); + } + + // text prompts + for (i, text) in ["moon", "sky", "bird"].iter().enumerate() { + let t = std::time::Instant::now(); + let out = model.detect(&mut enc, text).unwrap(); + eprintln!("detect {text:?}: {:?}", t.elapsed()); + assert!(rel_err("text features", &vec(&out.text), &vec(&r[&format!("det{i}_text")])) < 1e-3); + assert!(rel_err("query logits", &out.logits, &vec(&r[&format!("det{i}_logits")])) < 1e-2); + let presence = vec(&r[&format!("det{i}_presence")])[0]; + assert!((out.presence - presence).abs() < 1e-2, "presence {} vs {presence}", out.presence); + let boxes: Vec = out.boxes.iter().flatten().copied().collect(); + assert!(rel_err("boxes", &boxes, &vec(&r[&format!("det{i}_boxes")])) < 1e-2); + // the best instance's mask + let logits = vec(&r[&format!("det{i}_logits")]); + let q = (0..logits.len()).max_by(|a, b| logits[*a].partial_cmp(&logits[*b]).unwrap()).unwrap(); + let ours = vec(&out.masks.i(q).unwrap()); + let want = vec(&r[&format!("det{i}_masks")].i((0, q)).unwrap()); + assert!(rel_err("best mask", &ours, &want) < 2e-2); + assert!(mask_disagreement(&ours, &want) < 2e-3); + } +} diff --git a/crates/ui-egui/src/control.rs b/crates/ui-egui/src/control.rs index b592557a..1b9d22bb 100644 --- a/crates/ui-egui/src/control.rs +++ b/crates/ui-egui/src/control.rs @@ -278,7 +278,9 @@ pub fn handle(app: &mut LightcraftApp, ctx: &egui::Context, req: &ControlRequest Some(d) => { let r = crate::panels::dialogs::confirm_dialog(app, &d); // the import review stays open on an error, as with its button - if r.is_err() && matches!(d, crate::state::Dialog::Import { .. }) { + if (r.is_err() && matches!(d, crate::state::Dialog::Import { .. } | crate::state::Dialog::SamModel { .. })) + || (r.is_ok() && crate::panels::dialogs::keeps_open(app, &d)) + { app.ui.dialog = Some(d); } wrap(r) diff --git a/crates/ui-egui/src/lib.rs b/crates/ui-egui/src/lib.rs index bdf6f4f4..a4addbf3 100644 --- a/crates/ui-egui/src/lib.rs +++ b/crates/ui-egui/src/lib.rs @@ -203,7 +203,9 @@ pub struct LightcraftApp { } impl LightcraftApp { - pub fn new(session: Session, services: Services) -> Self { + pub fn new(mut session: Session, services: Services) -> Self { + // AI mask requests run on the model's worker; frames apply their results (never wait) + session.segmenter.background = true; Self { session, ui: UiState::default(), @@ -264,6 +266,7 @@ impl LightcraftApp { } let r = self.session.execute(id, ¶ms).map_err(|e| e.to_string()); if let Err(e) = &r { + log::warn!("{id}: {e}"); self.ui.status = e.clone(); } r @@ -357,6 +360,84 @@ impl LightcraftApp { self.ui.toast = Some((text.into(), t + secs)); } + /// AI masks: apply finished background requests (clicks, descriptions, detail passes) and + /// show their errors; start a zoomed-in detail pass once the clicking stops + /// (`ui.detail_due`); watch the model download. Never waits for the model. + fn ai_mask_detail(&mut self, ctx: &egui::Context) { + let polled = self.session.segment_poll(); + let now = ctx.input(|i| i.time); + if polled.changed { + ctx.request_repaint(); + } + if let Some(mask) = polled.refine { + self.ui.detail_due = Some((now + 0.3, mask)); + } + if let Some(e) = polled.messages.into_iter().last() { + self.ai_error(ctx, e, None); + } + let seg = &self.session.segmenter; + if seg.busy() || seg.pending_clicks().is_some() { + ctx.request_repaint_after(std::time::Duration::from_millis(50)); + } + let download = seg.download_status(); + if download.running { + ctx.request_repaint_after(std::time::Duration::from_millis(250)); + } + // a download finishing while its dialog is closed: say so once + if self.ui.sam_downloading && !download.running { + self.ui.sam_downloading = false; + let open = matches!(self.ui.dialog, Some(state::Dialog::SamModel { .. })); + match (&download.error, open) { + (_, true) => {} + (Some(e), false) if e.contains("cancelled") => { + self.toast(ctx, crate::i18n::tr("SAM 3 download stopped: it resumes where it left off next time.")) + } + (Some(e), false) => self.toast_error(ctx, format!("The SAM 3 download failed: {e}")), + (None, false) if download.finished => { + self.toast_error(ctx, crate::i18n::tr("The SAM 3 model is installed: Object and Describe masks are ready.")) + } + (None, false) => {} + } + } + if let Some((due, mask)) = self.ui.detail_due { + if now >= due && !self.session.segmenter.busy() && !self.session.segmenter.detail_busy() { + self.ui.detail_due = None; + 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)); + } + } + + /// An AI mask error: when the model isn't installed, the dialog that offers to download it + /// (`then`: the AI mask to start afterwards); otherwise a toast. + pub fn ai_error(&mut self, ctx: &egui::Context, e: impl Into, then: Option<(&str, &str)>) { + let e = e.into(); + if lightcraft_engine::segment::Segmenter::AVAILABLE && e.starts_with(lightcraft_engine::segment::NOT_INSTALLED) { + self.offer_sam_download(then); + } else { + self.toast_error(ctx, e); + } + } + + /// Open the dialog offering the SAM 3 download (never downloads by itself). + pub fn offer_sam_download(&mut self, then: Option<(&str, &str)>) { + if self.ui.dialog.is_none() || matches!(self.ui.dialog, Some(state::Dialog::SamModel { .. })) { + self.ui.dialog = Some(state::Dialog::SamModel { then: then.map(|(k, o)| (k.to_string(), o.to_string())), error: None }); + } + } + + /// A toast for an error the user has to read and act on (stays 6 s). + pub fn toast_error(&mut self, ctx: &egui::Context, text: impl Into) { + let t = ctx.input(|i| i.time); + 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() { @@ -688,6 +769,7 @@ impl LightcraftApp { } // panels set it again this frame while the pointer rests on a preset or profile self.hover_preview = None; + self.ai_mask_detail(&ctx); if self.ui.fullscreen { // full-screen preview: the photo alone on black egui::CentralPanel::default().frame(egui::Frame::NONE.fill(egui::Color32::BLACK)).show(ui, |ui| panels::detail::show(self, ui)); diff --git a/crates/ui-egui/src/panels/detail.rs b/crates/ui-egui/src/panels/detail.rs index cde3efdb..3b3622be 100644 --- a/crates/ui-egui/src/panels/detail.rs +++ b/crates/ui-egui/src/panels/detail.rs @@ -1172,6 +1172,47 @@ fn mask_overlay(app: &mut LightcraftApp, ui: &mut egui::Ui, resp: &egui::Respons } } } + // Object selections (SAM 3) of the selected mask: each click, green to include, red to exclude + // (clicks still on their way to the model included) + let pending = app.session.segmenter.pending_clicks().cloned(); + for m in d.masks.iter().filter(|m| Some(m.id) == active) { + for (k, c) in m.components.iter().enumerate() { + if let MaskShape::Object { hint, exclude, .. } = &c.shape { + let (hint, exclude) = match pending.as_ref().filter(|q| (q.mask, q.comp) == (m.id, k)) { + Some(q) => (&q.hint, &q.exclude), + None => (hint, exclude), + }; + for (pts, col) in [(hint, Color32::from_rgb(40, 200, 90)), (exclude, Color32::from_rgb(230, 60, 60))] { + for q in pts { + let q = map.screen(*q); + 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.ai_error(ui.ctx(), e, Some(("object", "new"))), + } + } + return; + } // brush tool if app.ui.tool == "brush" { let r = (app.ui.brush_size as f64 * long) as f32; diff --git a/crates/ui-egui/src/panels/dialogs.rs b/crates/ui-egui/src/panels/dialogs.rs index a9ecdcfe..bc23383b 100644 --- a/crates/ui-egui/src/panels/dialogs.rs +++ b/crates/ui-egui/src/panels/dialogs.rs @@ -94,6 +94,7 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) { Dialog::Merge { opts } => opts.title(), Dialog::Settings { .. } => "Settings", Dialog::ConfirmDelete { .. } => "Delete Photos", + Dialog::SamModel { .. } => "Download the SAM 3 Model?", Dialog::About => "About LightCraft", Dialog::Shortcuts => "Keyboard Shortcuts", } @@ -700,6 +701,7 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) { Dialog::Merge { opts } => crate::merge::body(app, ui, opts), Dialog::Import { opts } => crate::import::body(app, ui, opts), Dialog::Settings { tab } => crate::panels::settings::body(app, ui, tab), + Dialog::SamModel { error, .. } => sam_model_body(app, ui, error.as_deref()), Dialog::ConfirmDelete { count } => { let what = if *count == 1 { crate::i18n::tr("this photo").to_string() } else { crate::i18n::tr_format!("these {count} photos", count = count) }; ui.label(crate::i18n::tr_format!("Move {what} to Recently Deleted?", what = what)); @@ -767,8 +769,21 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) { ui.add_space(4.0); ui.horizontal(|ui| { let informational = matches!(dlg, Dialog::About | Dialog::Shortcuts | Dialog::Settings { .. }); - if !informational && ui.button(crate::i18n::tr("Cancel")).clicked() { - close = true; + let sam = &app.session.segmenter; + let (sam_installed, sam_running, sam_failed) = (sam.installed(), sam.download_status().running, sam.download_status().error.is_some()); + // no download location in this build: nothing to offer but the manual install + let sam_nowhere = !sam_installed && !sam_running && sam.mirrors().is_empty(); + let cancel = match &dlg { + Dialog::SamModel { .. } if sam_running || sam_installed || sam_nowhere => "Close", + Dialog::SamModel { .. } => "Not Now", + _ => "Cancel", + }; + if !informational { + let r = ui.button(crate::i18n::tr(cancel)); + crate::widgets::register(ui.ctx(), "button:dialogCancel", r.rect); + if r.clicked() { + close = true; + } } let add_label; let ok = match &dlg { @@ -780,10 +795,18 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) { } Dialog::Merge { .. } => "Merge", Dialog::ConfirmDelete { .. } => "Delete", + Dialog::SamModel { then: Some(_), .. } if sam_installed => "Continue", + Dialog::SamModel { .. } if sam_installed || sam_running || sam_nowhere => "", + Dialog::SamModel { error, .. } if error.is_some() || sam_failed => "Try Again", + Dialog::SamModel { .. } => "Download", _ if informational => "Close", _ => "OK", }; - if ui.button(crate::i18n::tr(ok)).clicked() { + let r = (!ok.is_empty()).then(|| ui.button(crate::i18n::tr(ok))); + if let Some(r) = &r { + crate::widgets::register(ui.ctx(), "button:dialogOk", r.rect); + } + if r.is_some_and(|r| r.clicked()) { if informational { close = true; } else { @@ -803,6 +826,17 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) { match confirm_dialog(app, &dlg) { // the import review stays open on an error (e.g. an unusable folder template) Err(e) if matches!(dlg, Dialog::Import { .. }) => app.toast(ctx, e), + // the SAM 3 dialog stays open to show the download (or why it can't start) + Err(e) if matches!(dlg, Dialog::SamModel { .. }) => { + if let Dialog::SamModel { error, .. } = &mut dlg { + *error = Some(e); + } + } + Ok(_) if keeps_open(app, &dlg) => { + if let Dialog::SamModel { error, .. } = &mut dlg { + *error = None; + } + } _ => close = true, } } @@ -823,8 +857,86 @@ pub fn fmt_gap(v: f64) -> String { } } +/// Whether a dialog stays open after its action succeeded (the SAM 3 dialog while the model +/// downloads). +pub fn keeps_open(app: &LightcraftApp, dlg: &Dialog) -> bool { + matches!(dlg, Dialog::SamModel { .. }) && !app.session.segmenter.installed() +} + +/// The SAM 3 dialog: what the model is, its size and licence, and the download's progress. +fn sam_model_body(app: &mut LightcraftApp, ui: &mut egui::Ui, error: Option<&str>) { + use lightcraft_engine::segment::{LICENSE_NAME, LICENSE_URL, MODEL_BYTES}; + let t = Tokens::get(ui.ctx()); + let seg = &app.session.segmenter; + let d = seg.download_status(); + if seg.installed() { + ui.label(crate::i18n::tr("The SAM 3 model is installed: Object and Describe masks are ready.")); + return; + } + let gb = |b: u64| b as f64 / 1e9; + ui.label(crate::i18n::tr( + "Object and Describe masks use SAM 3, Meta's segmentation model. It isn't part of LightCraft, and everything else works without it.", + )); + let dir = seg.dir.as_ref().map(|d| d.display().to_string()).unwrap_or_default(); + ui.label(format!("{} {:.1} GB, {} {dir}", crate::i18n::tr("A one-time download of about"), gb(MODEL_BYTES), crate::i18n::tr("saved in"))); + ui.label( + egui::RichText::new(format!( + "{} {LICENSE_NAME} — {}", + crate::i18n::tr("Licence:"), + crate::i18n::tr("Meta's terms, not LightCraft's. Downloading it means accepting them.") + )) + .color(t.text_label), + ); + let r = ui.link(crate::i18n::tr("Read the SAM License")).on_hover_text(LICENSE_URL); + crate::widgets::register(ui.ctx(), "link:samLicense", r.rect); + if r.clicked() { + let _ = crate::links::open(app, LICENSE_URL); + } + let seg = &app.session.segmenter; + if d.running { + ui.add_space(4.0); + let frac = if d.total > 0 { d.done as f64 / d.total as f64 } else { 0.0 }; + let text = format!("{:.2} / {:.2} GB · {}", gb(d.done), gb(d.total), d.file); + let r = ui.add(egui::ProgressBar::new(frac as f32).text(text)); + crate::widgets::register(ui.ctx(), "progress:samDownload", r.rect); + let r = ui.button(crate::i18n::tr("Cancel Download")); + crate::widgets::register(ui.ctx(), "button:samCancel", r.rect); + if r.clicked() { + app.session.segmenter.cancel_download(); + } + ui.label(egui::RichText::new(crate::i18n::tr("You can close this: the download continues, and resumes if interrupted.")).color(t.text_dim)); + ui.ctx().request_repaint_after(std::time::Duration::from_millis(250)); + return; + } + if seg.mirrors().is_empty() { + ui.label( + egui::RichText::new(crate::i18n::tr( + "This build has no download location for the model yet: put the files in the folder above yourself (see docs/ai-masks.md).", + )) + .color(t.text_dim), + ); + } + if let Some(e) = error.map(str::to_string).or(d.error) { + ui.label(egui::RichText::new(format!("{} {e}", crate::i18n::tr("The download didn't work:"))).color(egui::Color32::from_rgb(230, 90, 80))); + } +} + pub fn confirm_dialog(app: &mut LightcraftApp, dlg: &Dialog) -> Result { match dlg { + Dialog::SamModel { then, .. } => { + if app.session.segmenter.installed() { + // installed: start what the user was doing + if let Some((kind, op)) = then { + crate::panels::masking::begin_ai(app, kind, op)?; + } + return Ok(serde_json::Value::Null); + } + let r = app.run("segment.model.download", json!({"acknowledged": true})); + if r.is_ok() { + app.ui.sam_downloading = true; + } + r + } Dialog::NewAlbum { name, folder } => app.run("album.create", json!({"name": name, "folder": folder, "addSelected": !folder})), Dialog::RenameAlbum { id, name } => app.run("album.rename", json!({"id": id, "name": name})), Dialog::TextPrompt { value, command, params, key, .. } => { diff --git a/crates/ui-egui/src/panels/masking.rs b/crates/ui-egui/src/panels/masking.rs index f938dc8d..2e150a60 100644 --- a/crates/ui-egui/src/panels/masking.rs +++ b/crates/ui-egui/src/panels/masking.rs @@ -53,6 +53,7 @@ fn kind_label(s: &MaskShape) -> (&'static str, Icon) { MaskShape::Sky => ("Sky", Icon::Sky), MaskShape::Background => ("Background", Icon::Subject), MaskShape::Object { .. } => ("Object", Icon::Subject), + MaskShape::Prompt { .. } => ("Describe", Icon::Subject), MaskShape::People { .. } => ("People", Icon::Subject), MaskShape::Landscape { .. } => ("Landscape", Icon::Sky), } @@ -75,7 +76,9 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { egui::Frame::NONE.inner_margin(egui::Margin { left: 24, right: 22, top: 0, bottom: 10 }).show(ui, |ui| { ui.label(egui::RichText::new(crate::i18n::tr("Create New Mask")).color(t.text_dim)); ui.add_space(6.0); - let tiles: [(&str, &str, Icon); 8] = [ + let tiles: [(&str, &str, Icon); 10] = [ + ("object", "Object", Icon::Subject), + ("prompt", "Describe", Icon::Subject), ("subject", "Subject", Icon::Subject), ("sky", "Sky", Icon::Sky), ("background", "Background", Icon::Subject), @@ -102,6 +105,8 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { app.ui.tool = "colorRange".into(); app.toast(ui.ctx(), "Click the photo to pick a colour · ⇧-click adds more"); } + "object" => start_object(app, ui.ctx(), "new"), + "prompt" => start_describe(app, "new"), "brush" | "linear" | "radial" => { app.ui.tool = kind.to_string(); if *kind != "brush" { @@ -118,6 +123,37 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) { } } }); + describe_field(app, ui, true); + let seg = &app.session.segmenter; + let status = if seg.analyzing() { + Some("Analyzing the photo for AI masks…") + } else if seg.busy() { + Some("Selecting…") + } else if seg.detail_busy() || app.ui.detail_due.is_some() { + Some("Refining the mask's detail…") + } else { + None + }; + 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)); + } + // the model download, while its dialog is closed + let download = seg.download_status(); + if download.running && app.ui.dialog.is_none() { + ui.add_space(4.0); + ui.horizontal(|ui| { + let pct = if download.total > 0 { download.done as f64 / download.total as f64 } else { 0.0 }; + let r = ui.add( + egui::ProgressBar::new(pct as f32).desired_width(ui.available_width() - 70.0).text(format!("SAM 3: {} %", (pct * 100.0).floor())), + ); + register(ui.ctx(), "maskSamProgress", r.rect); + if text_button(ui, "maskSamDetails", crate::i18n::tr("Details"), false).clicked() { + app.offer_sam_download(None); + } + }); + } }); divider(ui); // mask list @@ -201,6 +237,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 +261,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); @@ -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,121 @@ fn component_row_menu(app: &mut LightcraftApp, ui: &mut egui::Ui, mask: u32, k: } } +/// Whether the SAM 3 model is missing in a build that could use it: then the download is +/// offered (with `then` to start afterwards) instead of starting an AI mask. +fn needs_model(app: &mut LightcraftApp, kind: &str, op: &str) -> bool { + let seg = &app.session.segmenter; + let missing = lightcraft_engine::segment::Segmenter::AVAILABLE && seg.dir.is_some() && !seg.installed(); + if missing { + app.offer_sam_download(Some((kind, op))); + } + missing +} + +/// Start an Object selection (SAM 3 clicks): a new mask, or a component of the selected one +/// combined by `op`; the photo is analyzed meanwhile (in the background). +pub(crate) fn start_object(app: &mut LightcraftApp, ctx: &egui::Context, op: &str) { + if needs_model(app, "object", op) { + return; + } + let r = + if op == "new" { app.run("mask.add", json!({"kind": "object"})) } else { app.run("mask.addComponent", json!({"op": op, "kind": "object"})) }; + match r { + Ok(_) => { + app.ui.tool = "object".into(); + app.toast(ctx, "Click the object to select it · ⌥-click leaves a part out"); + } + Err(e) => app.ai_error(ctx, e, Some(("object", op))), + } +} + +/// Start an AI mask of `kind` (object|prompt) combined by `op` (new|add|subtract|intersect), +/// after the model was installed. +pub(crate) fn begin_ai(app: &mut LightcraftApp, kind: &str, op: &str) -> Result { + match kind { + "object" => { + let r = if op == "new" { + app.run("mask.add", json!({"kind": "object"})) + } else { + app.run("mask.addComponent", json!({"op": op, "kind": "object"})) + }?; + app.ui.tool = "object".into(); + Ok(r) + } + _ => { + app.ui.describe = Some((op.to_string(), String::new())); + Ok(serde_json::Value::Null) + } + } +} + +/// Open the Describe field (a new mask, or a component combined by `op`). +pub(crate) fn start_describe(app: &mut LightcraftApp, op: &str) { + if needs_model(app, "prompt", op) { + return; + } + app.ui.describe = Some((op.to_string(), String::new())); +} + +/// The Describe field: type what to select ("sky", "the red car") and press Return. +fn describe_field(app: &mut LightcraftApp, ui: &mut egui::Ui, new_mask: bool) { + let Some((op, mut text)) = app.ui.describe.clone() else { return }; + // 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 { + // the mask appears when the model has found it (a detail pass follows) + Ok(_) => app.ui.describe = None, + Err(e) => app.ai_error(ui.ctx(), e, Some(("prompt", op.as_str()))), + } + } +} + 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() { + start_describe(app, op); + } for (kind, label) in [ ("brush", "Brush"), ("linear", "Linear Gradient"), diff --git a/crates/ui-egui/src/state.rs b/crates/ui-egui/src/state.rs index 151c3324..bdbbcf29 100644 --- a/crates/ui-egui/src/state.rs +++ b/crates/ui-egui/src/state.rs @@ -287,6 +287,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)>, + /// A SAM 3 download was started from the app (to report its end once). + #[serde(skip)] + pub sam_downloading: bool, + /// When to start the zoomed-in detail pass of an AI mask (app time) and which mask: set by + /// each click or description, so the pass runs once the clicking stops. + #[serde(skip)] + 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). @@ -506,6 +517,15 @@ pub enum Dialog { Settings { tab: String, }, + /// Object and Describe masks need the SAM 3 model, which isn't installed: offer to download + /// it (size, licence, progress). `then`: the AI mask to start once it is there (`kind` + /// object|prompt, `op` new|add|subtract|intersect). + SamModel { + then: Option<(String, String)>, + /// Why the download couldn't start (shown in the dialog). + #[serde(default)] + error: Option, + }, /// Confirm moving photos to Recently Deleted. ConfirmDelete { count: usize, @@ -561,6 +581,9 @@ impl Default for UiState { search: String::new(), focus_search: false, renaming_mask: None, + describe: None, + detail_due: None, + sam_downloading: false, renaming_component: None, quit: false, dragging_photos: None, diff --git a/crates/ui-egui/src/tests_masking.rs b/crates/ui-egui/src/tests_masking.rs index 1e4c0d14..2abc141f 100644 --- a/crates/ui-egui/src/tests_masking.rs +++ b/crates/ui-egui/src/tests_masking.rs @@ -620,3 +620,58 @@ fn soft_proofing_flags_out_of_gamut_colours_and_makes_proof_copies() { assert!(!h.app.ui.soft_proof); assert_ne!(h.app.session.catalog.stack_of(copy.id).unwrap().collapsed, stack.collapsed, "S toggled the stack"); } + +/// Object / Describe without the SAM 3 model: the download is offered (never started without a +/// yes), the dialog says why it can't start when no mirror is configured, nothing freezes and no +/// empty mask is left behind. +#[test] +fn ai_masks_without_the_model_offer_the_download() { + use crate::state::Dialog; + use lightcraft_engine::segment::Segmenter; + let mut h = detail("panel.masking"); + let dir = std::env::temp_dir().join(format!("lc-ui-no-sam3-{}", std::process::id())); + h.app.session.segmenter.dir = Some(dir.clone()); + h.app.session.segmenter.mirrors_file = Some(dir.join("none.txt")); + let t = std::time::Instant::now(); + let r = h.request("ui.clickWidget", json!({"id": "maskNew:object"}), T); + assert_eq!(r["ok"], true, "{r}"); + assert!(develop(&h).masks.is_empty()); + if !Segmenter::AVAILABLE { + // a build without AI masks says so (a toast), no dialog + assert_eq!(h.app.ui.dialog, None); + return; + } + assert_eq!(h.app.ui.dialog, Some(Dialog::SamModel { then: Some(("object".into(), "new".into())), error: None })); + assert!(!h.app.session.segmenter.download_status().running, "nothing downloads without a yes"); + if h.app.session.segmenter.mirrors().is_empty() { + // no location configured in this build: no Download button (only Close), and an agent + // confirming anyway gets the reason; the dialog stays + h.step(); + let r = h.request("ui.clickWidget", json!({"id": "button:dialogOk"}), T); + assert_eq!(r["ok"], false, "{r}"); + let r = h.request("ui.dialog.confirm", json!({}), T); + assert_eq!(r["ok"], false, "{r}"); + assert!(r["error"].as_str().unwrap_or_default().contains("LIGHTCRAFT_SAM3_MIRRORS"), "{r}"); + assert!(matches!(h.app.ui.dialog, Some(Dialog::SamModel { .. })), "stays open"); + } else { + // with a mirror: Download starts it in the background (here it fails: nothing listens) + let r = h.request("ui.clickWidget", json!({"id": "button:dialogOk"}), T); + assert_eq!(r["ok"], true, "{r}"); + } + // (a frame with the message laid out, so the buttons are where they are drawn) + h.step(); + h.step(); + let r = h.request("ui.clickWidget", json!({"id": "button:dialogCancel"}), T); + assert_eq!(r["ok"], true, "{r}"); + assert_eq!(h.app.ui.dialog, None); + // Describe asks too + let r = h.request("ui.clickWidget", json!({"id": "maskNew:prompt"}), T); + assert_eq!(r["ok"], true, "{r}"); + assert_eq!(h.app.ui.dialog, Some(Dialog::SamModel { then: Some(("prompt".into(), "new".into())), error: None })); + assert_eq!(h.app.ui.describe, None); + // a click command from an agent gets the not-installed error at once, and the app offers the model + let e = h.app.run("mask.add", json!({"kind": "prompt", "text": "sky"})).unwrap_err(); + assert!(e.starts_with(lightcraft_engine::segment::NOT_INSTALLED), "{e}"); + assert!(develop(&h).masks.is_empty()); + assert!(t.elapsed() < SETTLE, "{:?}", t.elapsed()); +} diff --git a/docs/ai-masks.md b/docs/ai-masks.md new file mode 100644 index 00000000..26d4f717 --- /dev/null +++ b/docs/ai-masks.md @@ -0,0 +1,198 @@ +# 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). + +**The model is optional.** It is not part of LightCraft, and nothing requires it: without it, +Object and Describe offer to download it, every other feature (and every existing mask, see +below) works, and no render, export or command ever waits for it. + +## Using them + +In the Masking panel (M): + +- **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). +- **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 (at most 4 per selection), + so small objects get 5–10× finer edges than one pass over the whole photo gives. + +The first Object or Describe on a photo analyzes it (the image encoder runs once per photo and +look: ~3 s on an Apple M4 Pro, plus a one-time ~6 s for compiling GPU kernels and loading the +model after launch). That starts in the background as soon as you pick Object; after it, each +click takes ~35 ms and each description ~0.6 s. Everything the model does runs on its own worker +thread: the window stays responsive, a click shows its dot at once and the selection follows +when it's computed ("Selecting…" / "Analyzing the photo…" under the mask tiles). The model is +unloaded after 10 minutes without use (it takes several GB of memory) and loaded again when +needed. + +The result is stored with the mask: the model's 288 × 288 mask logits over the uncropped photo +(quantized, compressed, ~10–30 KB). Renders and exports sample them at any resolution and never +need the model, so masks survive moving the library to a machine without it, and the web build +renders them. Editing the photo's look later does not recompute a mask (click again to update). + +## Getting the model + +### In the app (recommended) + +The first time you pick Object or Describe without the model, LightCraft asks: + +> **Download the SAM 3 model?** Object and Describe masks use SAM 3, Meta's segmentation model. +> It isn't part of LightCraft, and everything else works without it. Download it once (about +> 3.4 GB) · *folder*. Licence: SAM License (Meta) — Meta's terms, not LightCraft's. Downloading +> it means accepting them. [Read the SAM License] — **Not Now** / **Download** + +Nothing is downloaded until you choose **Download**. The download runs in the background (close +the dialog and keep working; the Masking panel shows its progress), can be cancelled, resumes +where it stopped if interrupted (also after quitting), and tries the next download location when +one fails or stalls. Each file is written as `.part` and moved into the model folder only +when complete and verified — `model.safetensors` must have the exact size and SHA-256 of Meta's +official checkpoint; a file that doesn't match is deleted, never used. If no location works, the +dialog says why and LightCraft carries on without the model. + +Agents and scripts use the same commands (control channel, `lightcraft-cli mcp --connect`): +`segment.model.status`, `segment.model.download {"acknowledged": true}` (refused without the +explicit `acknowledged`: pass it only after the user agreed to the download and the licence), +`segment.model.cancel`. + +### Download locations + +The model is fetched from an ordered list of mirrors: base URLs where `/model.safetensors`, +`/vocab.json` and `/merges.txt` live. The list is, in order: + +1. `LIGHTCRAFT_SAM3_MIRRORS` (environment variable; URLs separated by commas or spaces); +2. the file `models/sam3-mirrors.txt` in LightCraft's settings folder (one URL per line, `#` + comments) — `~/Library/Application Support/LightCraft/` (macOS), `%APPDATA%\LightCraft\` + (Windows), `~/.config/lightcraft/` (Linux); +3. the built-in list, `DEFAULT_MIRRORS` in `crates/segment/src/fetch/mod.rs`. + +> **Maintainers:** the built-in list is **empty** until LightCraft's own CDN locations exist +> (see the `TODO(maintainer)` there): add them in order of preference, host the three files +> unchanged, and pin `vocab.json` / `merges.txt` (size + SHA-256) in `SAM3_FILES` at the same +> time. Until then the in-app download needs a user-configured mirror, and the dialog says so. + +Hugging Face's `facebook/sam3` can't be a default: it is **gated** (each person must accept the +licence there, wait for approval and download with their own token). Someone with access can +use it as their mirror: + +```sh +LIGHTCRAFT_SAM3_MIRRORS=https://huggingface.co/facebook/sam3/resolve/main \ +LIGHTCRAFT_SAM3_TOKEN=hf_... lightcraft +``` + +(`LIGHTCRAFT_SAM3_TOKEN` is sent as a bearer token only over https, and only to the mirror's own +host — not to the storage host it redirects to.) + +The download is plain HTTP/1.1 over TLS in pure Rust (rustls with the RustCrypto provider and +the Mozilla root certificates; no OpenSSL, `ring` or `aws-lc`, nothing compiled from C), with a +15 s connect timeout and a 30 s stall timeout; it follows redirects but doesn't use HTTP proxies. + +### By hand + +Put `model.safetensors`, `vocab.json` and `merges.txt` from `facebook/sam3` in the model folder: +`models/sam3/` in the settings folder above, or the folder `LIGHTCRAFT_SAM3_DIR` names. Developers +can use the installer scripts, which download from Hugging Face with your token and verify the +weights: + +```sh +HF_TOKEN=hf_... tools/install-sam3.sh # macOS, Linux +tools/install-sam3.sh --check # verify an installation +``` + +```powershell +$env:HF_TOKEN = "hf_..."; .\tools\install-sam3.ps1 # Windows +``` + +## Licence of the model + +LightCraft's code is MIT OR Apache-2.0; `crates/segment` (a port of Apache-2.0 code) is +Apache-2.0. **The SAM 3 weights are neither**: they are Meta's, released under the **SAM +License** (), a custom licence that +is not an OSI open-source licence. It grants a non-exclusive, worldwide, non-transferable, +royalty-free licence to use, reproduce, distribute, copy, modify and create derivative works of +the "SAM Materials", subject to its terms, including: anyone redistributing them must provide a +copy of the licence with them, and derivative works are distributed under its terms; uses +subject to ITAR, military or warfare purposes, nuclear industries, espionage, guns or illegal +weapons, and reverse engineering are prohibited, and trade controls (sanctions, export rules) +apply; the materials come "as is", without warranty; Meta may terminate the licence on a breach +(then the materials must be deleted); California law governs. Read the licence itself — this +summary is not legal advice. + +Therefore: + +- **The weights are never part of this repository, its releases or installers**, and must not be + committed. +- They reach a user's computer only when that user asks for them, after the dialog names the + licence (or when they install them by hand). +- Whoever hosts a mirror redistributes the SAM Materials and must follow the licence (ship it + next to the files). + +## Commands (control channel, MCP) + +| Command | Params | | +|---|---|---| +| `mask.add` | `{kind: "object", points?: [[x,y],…], exclude?: [[x,y],…], seg?}` | an Object mask (empty until clicked) | +| `mask.add` | `{kind: "prompt", text, seg?}` | a Describe mask | +| `mask.addComponent` | `{op: add\|subtract\|intersect, kind: "object"\|"prompt", …}` | the same as a component | +| `mask.objectPoint` | `{x, y, exclude?: bool, id?}` | one click on the selected mask's Object selection (≤ 64 clicks) | +| `mask.refineDetail` | `{id?, component?}` | the zoomed-in detail pass → `{started}` | +| `segment.prepare` | `{}` | load the model and analyze the active photo → `{busy}` | +| `segment.model.status` | `{}` | `{available, installed, dir, loaded, busy, analyzing, sizeBytes, license, licenseUrl, mirrors, download}` | +| `segment.model.download` | `{acknowledged: true}` | start the download in the background → `{started}` | +| `segment.model.cancel` | `{}` | stop it (it resumes next time) → `{cancelled}` | + +Coordinates are normalized to the uncropped, oriented photo, like every mask shape. `seg` passes +a stored segmentation (as `mask.list` / the develop settings hold it): the mask is made without +the model. Without the model, the AI kinds fail with an error starting "The SAM 3 model is not +installed" that says how to get it. + +In the desktop app the commands don't wait for the model: `mask.objectPoint` returns +`{pending: true}` and the selection is applied when it's computed (one undo step per settled +selection); `mask.add {kind: "prompt"}` returns `{pending: true}` and the mask appears when +something is found (or a message says nothing was). `lightcraft-cli` and the headless MCP server +wait and return the result (when built with the `sam` feature; the default CLI has no AI masks). + +## Implementation notes + +- `crates/segment` ports the Hugging Face `transformers` implementation (Apache-2.0; see + `NOTICE` — the ported files say so and that they were modified): the 32-layer ViT backbone + with 2-D RoPE and windowed attention, the feature pyramids, the SAM 2-style prompt encoder and + two-way mask decoder for clicks, and the CLIP text encoder, DETR encoder/decoder (box + relative-position bias, presence token) and pixel decoder for text. The CLIP tokenizer is a + small BPE in `tokenizer.rs`. `fetch/` is the downloader (original code). +- candle is pinned at 0.9.2: later releases make `candle-core` depend on `tokenizers` with the + Oniguruma C library, and the product is pure Rust. Weights are read with positional reads (no + memory map: `unsafe` stays in `lightcraft-sysmem`), only the tensors a path needs, and every + tensor's byte range is checked against the file's length when it is opened (a truncated or + damaged file is an error, never a crash). +- The engine (`crates/engine/src/segment/`) owns one worker thread with the model and the last + encoded photo; the session sends it jobs and receives results over channels, never holding a + lock the model holds. Each job runs under `catch_unwind`: a panic inside candle becomes an + error and the model is reloaded on the next request; busy counters are reset by drop guards. +- Hostile documents: a selection keeps at most 4 detail patches (when read and when rendered), + each at most 1024² logits; more than 64 clicks are never sent to the model. +- Accuracy: `crates/segment/tests/reference.rs` compares every stage against the Python + reference (`tools/sam3_reference.py`); on Metal the backbone matches to 7e-5 relative error + and masks to ~5e-5. f16 was tried and rejected (overflow, ~7 % faster). +- The engine feature `sam` (on in the desktop app, off for web and CLI) gates the model; without + it the commands report that the build has no AI masks. + +**Developers:** the accuracy test against the Python reference needs `torch` and +`transformers` (`tools/requirements-sam3-reference.txt`): + +```sh +python3 -m venv .venv-sam3 && .venv-sam3/bin/pip install -r tools/requirements-sam3-reference.txt +.venv-sam3/bin/python tools/sam3_reference.py photo.jpg ref.safetensors +LIGHTCRAFT_SAM3_DIR= LIGHTCRAFT_SAM3_REF=ref.safetensors cargo test -p lightcraft-segment --release -- --nocapture +``` diff --git a/docs/parity.md b/docs/parity.md index 38d768af..82729971 100644 --- a/docs/parity.md +++ b/docs/parity.md @@ -31,7 +31,7 @@ and prints the summary; `cargo xtask parity --write` refreshes the summary table | H. Crop & rotate (CROP) | 9 | 0 | 0 | 1 | 6/6 (100%) | 3/3 (100%) | | I. Remove / healing (REM) | 7 | 1 | 2 | 2 | 4/4 (100%) | 2/3 (67%) | | J. Red eye (EYE) | 2 | 0 | 0 | 0 | — | 1/1 (100%) | -| K. Masking (MASK) | 14 | 3 | 6 | 0 | 8/8 (100%) | 5/5 (100%) | +| K. Masking (MASK) | 14 | 4 | 5 | 0 | 8/8 (100%) | 5/5 (100%) | | L. Presets (PRE) | 6 | 0 | 1 | 1 | 2/2 (100%) | 2/2 (100%) | | M. Versions & history (VER) | 5 | 0 | 0 | 0 | 1/1 (100%) | 3/3 (100%) | | N. Copy / paste / sync (SYNC) | 5 | 0 | 0 | 0 | 3/3 (100%) | 1/1 (100%) | @@ -48,9 +48,9 @@ and prints the summary; `cargo xtask parity --write` refreshes the summary table | Y. Menus | 81 | 0 | 4 | 8 | 47/47 (100%) | 23/23 (100%) | | Z. Keyboard shortcuts (desktop) | 74 | 3 | 3 | 1 | 49/52 (94%) | 22/23 (96%) | | Lightroom Classic extras | 31 | 10 | 39 | 9 | — | 21/22 (95%) | -| **Total** | 388 | 33 | 87 | 37 | 194/200 (97%) | 139/149 (93%) | +| **Total** | 388 | 34 | 86 | 37 | 194/200 (97%) | 139/149 (93%) | -Weighted completion (✅ = 1, 🟡 = ½, 🚫 left out): **79.6%** of 508 in-scope rows — P0 98.5% of 200 · P1 95.6% of 149 · P2 40.9% of 159. +Weighted completion (✅ = 1, 🟡 = ½, 🚫 left out): **79.7%** of 508 in-scope rows — P0 98.5% of 200 · P1 95.6% of 149 · P2 41.2% of 159. ## Top gaps @@ -69,8 +69,10 @@ Take the first one nobody is working on. then tune tone, highlights, texture/clarity/dehaze, NR and sharpening against it. 5. **LR-EDIT-OPTICS-PROFILE** (P1): a lens-profile database of our own (embedded DNG/maker corrections work today). 6. **AI masks and Enhance** (LR-MASK-SUBJECT / SKY / PEOPLE / OBJECTS, LR-EDIT-DETAIL-DENOISE, SUPERRES, LENSBLUR): - 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`); the weights (SAM License) are never bundled: + the app offers a consented, verified download, but **LightCraft's CDN mirrors are not configured yet** (the + built-in list is empty, so today users need their own mirror or a manual install). Subject / Sky / People could + use the same model with fixed prompts; denoise / super-resolution models remain a maintainer decision. 7. **HDR** (Q. HDR, LR-EXP-HDR), **video** (R. Video), **Classic output modules** (Map view, Book, Slideshow module, Print, publish): large, well understood, lower priority than 1–5. @@ -328,7 +330,7 @@ under the UI command that wraps it (a no-op error); the UI command now wins. | LR-MASK-SUBJECT | Select subject | P2 | 🟡 | `cmd:mask.add` (`subject`), `crates/pipeline/src/masks.rs` | saliency heuristic, no segmentation model | | LR-MASK-SKY | Select sky | P2 | 🟡 | `cmd:mask.add` (`sky`) | heuristic | | LR-MASK-BACKGROUND | Select background | P2 | 🟡 | `cmd:mask.add` (`background`) | inverse of the subject heuristic | -| LR-MASK-OBJECTS | Object selection | P2 | ⬜ | | 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:mask.refineDetail`, `cmd:segment.prepare`, `cmd:segment.model.status`, `cmd:segment.model.download`, `cmd:segment.model.cancel`, `crates/segment`, `crates/engine/src/segment/mod.rs` | SAM 3 in pure Rust (candle; Metal on macOS, CPU elsewhere): Object tile → click to include, ⌥-click to leave out; Describe tile → a text prompt selects every instance ("sky", "the red car"); both also as Add/Subtract/Intersect components; + / − under the selected mask; comma lists (`car, road`); a zoomed-in detail pass for 5–10× finer edges on small objects; per-selection Edge (hard ↔ soft). The model runs on its own worker thread (the UI never waits; panics become errors), is unloaded after 10 min idle, and is optional: the segmentation is stored with the mask (288² logits), so renders and exports never need it. Weights are never bundled (SAM License): without them the app offers a consented background download (mirrors, resume, timeouts, SHA-256) — **but no default download location is configured yet** (users set `LIGHTCRAFT_SAM3_MIRRORS` or install by hand; see docs/ai-masks.md). Not verified against Lightroom's Select Object; no brush/box object mode; the first click on a photo waits for its analysis (~4 s on an M4 Pro, much longer on CPU) | | LR-MASK-PEOPLE | People parts | P2 | ⬜ | | | | LR-MASK-LANDSCAPE | Landscape classes | P2 | ⬜ | | shape exists, evaluates empty | | LR-MASK-BRUSH | Brush mask | P0 | ✅ | `cmd:tool.brush`, `cmd:mask.brushStroke` (`autoMask`), `crates/pipeline/src/masks.rs` | size/feather/flow/density/erase; Auto Mask: dabs weighted by similarity to the colour under the dab centre, refined by a guided filter on luminance (CPU + GPU); no A/B brushes, no pressure | diff --git a/docs/whats-new.md b/docs/whats-new.md index 12984b9a..00480132 100644 --- a/docs/whats-new.md +++ b/docs/whats-new.md @@ -131,6 +131,12 @@ lost everything at quit; a temporary session shows a banner the whole time and never writes to your library. ### Editing +- AI masks with SAM 3 (Object and Describe in the Masking panel): click an object to select it (⌥-click leaves a + part out), or type what to select ("sky", "the red car", "car, road"); both combine with other masks, have an + Edge setting, and get a sharper zoomed-in pass in the background. The model runs inside LightCraft in pure Rust + and never freezes the window. It is optional and not part of LightCraft (Meta's SAM License, about 3.4 GB): the + first time you use an AI mask, LightCraft asks before downloading it, shows the progress, can cancel and resume, + and checks the file before using it. Masks keep their selection, so they render and export without the model. - Auto Sync: edits apply to every selected photo. Auto B&W mix. Automatic versions. - Colour-range masks: click the photo to sample. Luminance ranges: range bar, smoothness, luminance map. - ⌘-drag to straighten, ⇧G Guided Upright, a grid while transforming. diff --git a/tools/install-sam3.ps1 b/tools/install-sam3.ps1 new file mode 100644 index 00000000..0036c561 --- /dev/null +++ b/tools/install-sam3.ps1 @@ -0,0 +1,60 @@ +# 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). The desktop app +# offers to download the model itself when an AI mask first needs it; this is for developers. +# +# facebook/sam3 is gated: accept the SAM License at https://huggingface.co/facebook/sam3, wait +# for approval, then: $env:HF_TOKEN = "hf_..." ; .\tools\install-sam3.ps1 +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" + # into a .part file, moved into place only when complete (and, for the weights, verified): + # a failed or damaged download never ends up under the real name + $final = Join-Path $Dir $f + $part = "$final.part" + Remove-Item -Force -ErrorAction SilentlyContinue $part + try { + Invoke-WebRequest -Uri "$Base/$f" -Headers $Headers -OutFile $part + } catch { + Remove-Item -Force -ErrorAction SilentlyContinue $part + throw "download of $Repo/$f failed: $($_.Exception.Message)" + } + if ($f -eq "model.safetensors") { + Write-Host "checking model.safetensors (SHA-256 of 3.4 GB)..." + $got = (Get-FileHash $part -Algorithm SHA256).Hash.ToLower() + if ($got -ne $Sha256) { + Remove-Item -Force $part + throw "the download is damaged or not the official checkpoint (got $got); deleted it" + } + } + Move-Item -Force $part $final +} +Test-Install +Write-Host "Restart LightCraft; Object and Describe in the Masking panel now use it." diff --git a/tools/install-sam3.sh b/tools/install-sam3.sh new file mode 100755 index 00000000..4ec9d0a9 --- /dev/null +++ b/tools/install-sam3.sh @@ -0,0 +1,130 @@ +#!/bin/sh +# Install the SAM 3 model for LightCraft's Object and Describe masks (see docs/ai-masks.md). +# +# The desktop app offers to download the model itself when an AI mask first needs it; this script +# is for developers (reference tests, benchmarks) and for installing from Hugging Face directly. +# +# The weights are not part of LightCraft: they are Meta's facebook/sam3 checkpoint on Hugging +# Face, under the SAM License. Access is gated: open https://huggingface.co/facebook/sam3, accept +# the license, wait for approval, then run this with a token that can read it: +# +# 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 FILE DEST [resume]: HTTP errors never reach DEST (--fail writes no error page), and a +# resume the server refuses (416) starts over instead of counting as done. +fetch() { + set -- "$1" "$2" "${3:-}" + if [ -n "$TOKEN" ]; then auth="Authorization: Bearer $TOKEN"; else auth="X-No-Auth: 1"; fi + code=$(curl -fsSL ${3:+-C -} -H "$auth" -w '%{http_code}' -o "$2" "$BASE/$1") || true + case "$code" in + 200|206) return 0 ;; + 416) + echo "the partial download of $1 doesn't match the server's; starting over" >&2 + rm -f "$2" + code=$(curl -fsSL -H "$auth" -w '%{http_code}' -o "$2" "$BASE/$1") || true + [ "$code" = 200 ] && return 0 + echo "download of $REPO/$1 failed ($code)" >&2; return 1 ;; + 401|403) echo "access denied to $REPO/$1 (has your request for facebook/sam3 been approved?)" >&2; return 1 ;; + *) echo "download of $REPO/$1 failed (${code:-no response})" >&2; return 1 ;; + esac +} + +for f in $FILES; do + echo "fetching $f" + # small files: whole, into place only when complete + fetch "$f" "$DIR/$f.part" && mv -f "$DIR/$f.part" "$DIR/$f" || { rm -f "$DIR/$f.part"; exit 1; } +done + +if [ -s "$DIR/model.safetensors" ] && [ "$(sha256 "$DIR/model.safetensors")" = "$SHA256" ]; then + echo "model.safetensors already installed" +else + echo "fetching model.safetensors (3.4 GB; resumes if interrupted)" + part="$DIR/model.safetensors.part" + # a wrong model.safetensors is never resumed from: it goes + rm -f "$DIR/model.safetensors" + if [ -f "$part" ] && [ "$(wc -c < "$part" | tr -d ' ')" -ge "$SIZE" ]; then rm -f "$part"; fi + fetch model.safetensors "$part" resume || exit 1 + have=$(wc -c < "$part" | tr -d ' ') + if [ "$have" -lt "$SIZE" ]; then + echo "incomplete download ($have of $SIZE bytes); run the script again to resume" >&2 + exit 1 + fi + echo "checking model.safetensors (SHA-256 of 3.4 GB, takes a few seconds)…" + if [ "$have" != "$SIZE" ] || [ "$(sha256 "$part")" != "$SHA256" ]; then + rm -f "$part" + echo "the download is damaged or not the official checkpoint; deleted it, run the script again" >&2 + exit 1 + fi + mv -f "$part" "$DIR/model.safetensors" +fi + +verify +echo "Restart LightCraft; Object and Describe in the Masking panel now use it." diff --git a/tools/requirements-sam3-reference.txt b/tools/requirements-sam3-reference.txt new file mode 100644 index 00000000..cf8f263a --- /dev/null +++ b/tools/requirements-sam3-reference.txt @@ -0,0 +1,8 @@ +# Python packages for tools/sam3_reference.py (reference tensors for crates/segment/tests/reference.rs). +# python3 -m venv .venv-sam3 && .venv-sam3/bin/pip install -r tools/requirements-sam3-reference.txt +torch>=2.6 +torchvision>=0.21 +transformers>=4.57 +pillow +safetensors +numpy diff --git a/tools/sam3_reference.py b/tools/sam3_reference.py new file mode 100644 index 00000000..66d9def9 --- /dev/null +++ b/tools/sam3_reference.py @@ -0,0 +1,71 @@ +"""Reference tensors for the Rust SAM 3 port (crates/segment/tests/reference.rs). + +Runs the Hugging Face transformers implementation (fp32, CPU) on one photo and writes the +shared pixel input plus intermediate and final outputs of the click (tracker) and text +(detector) paths. Usage (needs torch, torchvision, transformers, pillow, safetensors): + LIGHTCRAFT_SAM3_DIR= python tools/sam3_reference.py photo.jpg ref.safetensors +The text prompts below suit a photo of the moon in a blue sky; change them for other photos. +""" +import os, sys, time +import numpy as np +import torch +from PIL import Image +from safetensors.torch import save_file +from transformers import Sam3TrackerModel, Sam3Model, CLIPTokenizer + +MODEL = os.environ.get("LIGHTCRAFT_SAM3_DIR", os.path.expanduser("~/Library/Application Support/LightCraft/models/sam3")) +img_path, out_path = sys.argv[1], sys.argv[2] +torch.set_grad_enabled(False) + +img = Image.open(img_path).convert("RGB") +W, H = img.size +# Same preprocessing as the Rust side: bilinear resize to 1008², /255, (x-0.5)/0.5. +r = img.resize((1008, 1008), Image.BILINEAR) +px = torch.from_numpy(np.asarray(r, dtype=np.float32) / 255.0).permute(2, 0, 1)[None] +px = (px - 0.5) / 0.5 +out = {"pixel_values": px.contiguous()} + +t = time.time() +trk = Sam3TrackerModel.from_pretrained(MODEL).eval() +vis = trk.vision_encoder(px) +out["backbone"] = vis.last_hidden_state.contiguous() +for i, f in enumerate(vis.fpn_hidden_states): + out[f"trk_fpn{i}"] = f.contiguous() +# points in 1008 space: one positive, one negative +pts = torch.tensor([[[[0.5 * 1008, 0.55 * 1008], [0.2 * 1008, 0.2 * 1008]]]], dtype=torch.float32) +lbl = torch.tensor([[[1, 0]]], dtype=torch.int32) +o = trk(pixel_values=px, input_points=pts, input_labels=lbl, multimask_output=False) +out["trk_points"] = pts.contiguous() +out["trk_masks"] = o.pred_masks.contiguous() +out["trk_iou"] = o.iou_scores.contiguous() +o3 = trk(pixel_values=px, input_points=pts[:, :, :1], input_labels=lbl[:, :, :1], multimask_output=True) +out["trk_multi_masks"] = o3.pred_masks.contiguous() +out["trk_multi_iou"] = o3.iou_scores.contiguous() +print("tracker", time.time() - t, o.pred_masks.shape, o.iou_scores, file=sys.stderr) +del trk + +t = time.time() +det = Sam3Model.from_pretrained(MODEL).eval() +tok = CLIPTokenizer.from_pretrained(MODEL) +for i, text in enumerate(["moon", "sky", "bird"]): + enc = tok(text, padding="max_length", max_length=32, return_tensors="pt") + out[f"det{i}_ids"] = enc.input_ids.to(torch.int64).contiguous() + out[f"det{i}_mask"] = enc.attention_mask.to(torch.int64).contiguous() + tf = det.get_text_features(input_ids=enc.input_ids, attention_mask=enc.attention_mask) + out[f"det{i}_text"] = tf.pooler_output.contiguous() + o = det(pixel_values=px, input_ids=enc.input_ids, attention_mask=enc.attention_mask) + out[f"det{i}_logits"] = o.pred_logits.contiguous() + out[f"det{i}_presence"] = o.presence_logits.contiguous() + out[f"det{i}_boxes"] = o.pred_boxes.contiguous() + out[f"det{i}_masks"] = o.pred_masks.contiguous() + sc = (o.pred_logits.sigmoid() * o.presence_logits.sigmoid())[0] + print(text, "presence", o.presence_logits.sigmoid().item(), "top", sc.topk(3).values.tolist(), file=sys.stderr) +for i, f in enumerate(det.vision_encoder(px).fpn_hidden_states): + out[f"det_fpn{i}"] = f.contiguous() +print("detector", time.time() - t, file=sys.stderr) +# tokenizer cases for the Rust BPE +cases = ["person", "Red Car", "a dog's toy", "sky, clouds!", "café 2 people", "tree-line"] +for i, c in enumerate(cases): + out[f"tok{i}"] = tok(c, padding="max_length", max_length=32, return_tensors="pt").input_ids.to(torch.int64).contiguous() +save_file(out, out_path, metadata={"image": img_path, "W": str(W), "H": str(H), "cases": "|".join(cases)}) +print("saved", out_path, file=sys.stderr) diff --git a/xtask/src/layers.rs b/xtask/src/layers.rs index b5c63142..28c99462 100644 --- a/xtask/src/layers.rs +++ b/xtask/src/layers.rs @@ -49,6 +49,7 @@ pub const TABLE: &[(&str, Class)] = &[ ("preview", Class::Layer(3)), ("export", Class::Layer(3)), ("merge", Class::Layer(3)), + ("segment", Class::Layer(3)), ("engine", Class::Layer(4)), ("ui-egui", Class::Layer(5)), ("mcp", Class::Layer(5)),