mirror of
https://github.com/storytold/effectcraft.git
synced 2026-10-31 07:38:24 +00:00
S2.11: Hiera backbone matches PyTorch block by block
Weights loader with shape/finite validation (S2.10), patch embed, bicubic positional embedding, window/global attention, query pooling. Reference test (skips without weights): max abs error 2.6e-5 over all 12 blocks. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
Kristian MyramandClaude Sonnet 5.5 committed
1 parent
36a22733a2
commit
92b3a34732
7 files changed
+916
-10
No files matched your search
Generated
+331
-10
@@ -582,7 +582,7 @@ dependencies = [
|
||||
"oneshot",
|
||||
"pastey",
|
||||
"portable-atomic",
|
||||
"rand",
|
||||
"rand 0.10.3",
|
||||
"regress",
|
||||
"rustc-hash 2.1.3",
|
||||
"ryu-js",
|
||||
@@ -723,6 +723,12 @@ dependencies = [
|
||||
"syn 3.0.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "byteorder"
|
||||
version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
|
||||
|
||||
[[package]]
|
||||
name = "byteorder-lite"
|
||||
version = "0.1.0"
|
||||
@@ -786,6 +792,45 @@ dependencies = [
|
||||
"wayland-client",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "candle-core"
|
||||
version = "0.9.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c15b675b80d994b2eadb20a4bbe434eabeb454eac3ee5e2b4cf6f147ee9be091"
|
||||
dependencies = [
|
||||
"byteorder",
|
||||
"float8",
|
||||
"gemm",
|
||||
"half",
|
||||
"libm",
|
||||
"memmap2",
|
||||
"num-traits",
|
||||
"num_cpus",
|
||||
"rand 0.9.5",
|
||||
"rand_distr",
|
||||
"rayon",
|
||||
"safetensors",
|
||||
"thiserror 2.0.21",
|
||||
"yoke",
|
||||
"zip 7.4.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "candle-nn"
|
||||
version = "0.9.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3045fa9e7aef8567d209a27d56b692f60b96f4d0569f4c3011f8ca6715c65e03"
|
||||
dependencies = [
|
||||
"candle-core",
|
||||
"half",
|
||||
"libc",
|
||||
"num-traits",
|
||||
"rayon",
|
||||
"safetensors",
|
||||
"serde",
|
||||
"thiserror 2.0.21",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cast"
|
||||
version = "0.3.0"
|
||||
@@ -839,7 +884,7 @@ checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cpufeatures 0.3.1",
|
||||
"rand_core",
|
||||
"rand_core 0.10.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1141,7 +1186,7 @@ dependencies = [
|
||||
"verifysign",
|
||||
"winreg",
|
||||
"winsafe",
|
||||
"zip",
|
||||
"zip 9.1.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1310,6 +1355,22 @@ version = "1.0.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813"
|
||||
|
||||
[[package]]
|
||||
name = "dyn-stack"
|
||||
version = "0.13.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1c4713e43e2886ba72b8271aa66c93d722116acf7a75555cce11dcde84388fe8"
|
||||
dependencies = [
|
||||
"bytemuck",
|
||||
"dyn-stack-macros",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "dyn-stack-macros"
|
||||
version = "0.1.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e1d926b4d407d372f141f93bb444696142c29d32962ccbd3531117cf3aa0bfa9"
|
||||
|
||||
[[package]]
|
||||
name = "dynify"
|
||||
version = "0.1.2"
|
||||
@@ -1767,7 +1828,11 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "effectcraft-sam2"
|
||||
version = "0.9.2"
|
||||
version = "0.18.0"
|
||||
dependencies = [
|
||||
"candle-core",
|
||||
"candle-nn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "effectcraft-script"
|
||||
@@ -2065,6 +2130,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"
|
||||
@@ -2590,6 +2667,18 @@ dependencies = [
|
||||
"zlib-rs",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "float8"
|
||||
version = "0.6.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "719a903cc23e4a89e87962c2a80fdb45cdaad0983a89bd150bb57b4c8571a7d5"
|
||||
dependencies = [
|
||||
"half",
|
||||
"num-traits",
|
||||
"rand 0.9.5",
|
||||
"rand_distr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fmv-macos-events"
|
||||
version = "0.1.0"
|
||||
@@ -2732,6 +2821,125 @@ dependencies = [
|
||||
"slab",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "gemm"
|
||||
version = "0.19.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "aa0673db364b12263d103b68337a68fbecc541d6f6b61ba72fe438654709eacb"
|
||||
dependencies = [
|
||||
"dyn-stack",
|
||||
"gemm-c32",
|
||||
"gemm-c64",
|
||||
"gemm-common",
|
||||
"gemm-f16",
|
||||
"gemm-f32",
|
||||
"gemm-f64",
|
||||
"num-complex",
|
||||
"num-traits",
|
||||
"paste",
|
||||
"raw-cpuid",
|
||||
"seq-macro",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "gemm-c32"
|
||||
version = "0.19.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "086936dbdcb99e37aad81d320f98f670e53c1e55a98bee70573e83f95beb128c"
|
||||
dependencies = [
|
||||
"dyn-stack",
|
||||
"gemm-common",
|
||||
"num-complex",
|
||||
"num-traits",
|
||||
"paste",
|
||||
"raw-cpuid",
|
||||
"seq-macro",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "gemm-c64"
|
||||
version = "0.19.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "20c8aeeeec425959bda4d9827664029ba1501a90a0d1e6228e48bef741db3a3f"
|
||||
dependencies = [
|
||||
"dyn-stack",
|
||||
"gemm-common",
|
||||
"num-complex",
|
||||
"num-traits",
|
||||
"paste",
|
||||
"raw-cpuid",
|
||||
"seq-macro",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "gemm-common"
|
||||
version = "0.19.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "88027625910cc9b1085aaaa1c4bc46bb3a36aad323452b33c25b5e4e7c8e2a3e"
|
||||
dependencies = [
|
||||
"bytemuck",
|
||||
"dyn-stack",
|
||||
"half",
|
||||
"libm",
|
||||
"num-complex",
|
||||
"num-traits",
|
||||
"once_cell",
|
||||
"paste",
|
||||
"pulp",
|
||||
"raw-cpuid",
|
||||
"rayon",
|
||||
"seq-macro",
|
||||
"sysctl",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "gemm-f16"
|
||||
version = "0.19.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e3df7a55202e6cd6739d82ae3399c8e0c7e1402859b30e4cb780e61525d9486e"
|
||||
dependencies = [
|
||||
"dyn-stack",
|
||||
"gemm-common",
|
||||
"gemm-f32",
|
||||
"half",
|
||||
"num-complex",
|
||||
"num-traits",
|
||||
"paste",
|
||||
"raw-cpuid",
|
||||
"rayon",
|
||||
"seq-macro",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "gemm-f32"
|
||||
version = "0.19.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "02e0b8c9da1fbec6e3e3ab2ce6bc259ef18eb5f6f0d3e4edf54b75f9fd41a81c"
|
||||
dependencies = [
|
||||
"dyn-stack",
|
||||
"gemm-common",
|
||||
"num-complex",
|
||||
"num-traits",
|
||||
"paste",
|
||||
"raw-cpuid",
|
||||
"seq-macro",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "gemm-f64"
|
||||
version = "0.19.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "056131e8f2a521bfab322f804ccd652520c79700d81209e9d9275bbdecaadc6a"
|
||||
dependencies = [
|
||||
"dyn-stack",
|
||||
"gemm-common",
|
||||
"num-complex",
|
||||
"num-traits",
|
||||
"paste",
|
||||
"raw-cpuid",
|
||||
"seq-macro",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "generic-array"
|
||||
version = "0.14.7"
|
||||
@@ -2785,7 +2993,7 @@ dependencies = [
|
||||
"js-sys",
|
||||
"libc",
|
||||
"r-efi 6.0.0",
|
||||
"rand_core",
|
||||
"rand_core 0.10.1",
|
||||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
@@ -2933,9 +3141,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",
|
||||
]
|
||||
|
||||
@@ -2975,6 +3186,8 @@ dependencies = [
|
||||
"allocator-api2",
|
||||
"equivalent",
|
||||
"foldhash 0.2.0",
|
||||
"serde",
|
||||
"serde_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3610,6 +3823,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"stable_deref_trait",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3861,6 +4075,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"
|
||||
@@ -4300,7 +4524,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7d8fae84b431384b68627d0f9b3b1245fcf9f46f6c0e3dc902e9dce64edd1967"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys 0.45.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4569,6 +4793,15 @@ version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391"
|
||||
|
||||
[[package]]
|
||||
name = "ppv-lite86"
|
||||
version = "0.2.21"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9"
|
||||
dependencies = [
|
||||
"zerocopy",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "presser"
|
||||
version = "0.3.1"
|
||||
@@ -4611,7 +4844,7 @@ dependencies = [
|
||||
"chacha20",
|
||||
"core_detect",
|
||||
"num-traits",
|
||||
"rand",
|
||||
"rand 0.10.3",
|
||||
"rand_xorshift",
|
||||
"regex-syntax",
|
||||
"rusty-fork",
|
||||
@@ -4699,6 +4932,16 @@ version = "6.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
|
||||
|
||||
[[package]]
|
||||
name = "rand"
|
||||
version = "0.9.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41"
|
||||
dependencies = [
|
||||
"rand_chacha",
|
||||
"rand_core 0.9.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand"
|
||||
version = "0.10.3"
|
||||
@@ -4707,7 +4950,26 @@ checksum = "65c9fb96cbc91e3478eaae79a69fcd3f1ae4ad052e471fe6732fff548984b4af"
|
||||
dependencies = [
|
||||
"chacha20",
|
||||
"getrandom 0.4.3",
|
||||
"rand_core",
|
||||
"rand_core 0.10.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand_chacha"
|
||||
version = "0.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
|
||||
dependencies = [
|
||||
"ppv-lite86",
|
||||
"rand_core 0.9.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand_core"
|
||||
version = "0.9.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c"
|
||||
dependencies = [
|
||||
"getrandom 0.3.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4716,13 +4978,23 @@ version = "0.10.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69"
|
||||
|
||||
[[package]]
|
||||
name = "rand_distr"
|
||||
version = "0.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463"
|
||||
dependencies = [
|
||||
"num-traits",
|
||||
"rand 0.9.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand_xorshift"
|
||||
version = "0.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "60aa6af80be32871323012e02e6e65f8a7cc7890931ae421d217ad8fe0df2ccf"
|
||||
dependencies = [
|
||||
"rand_core",
|
||||
"rand_core 0.10.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -5044,6 +5316,17 @@ version = "1.0.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "04d056b875a9d2e6cb9a61d127afee9ac5999b9f87bcb32079d1318e505be714"
|
||||
|
||||
[[package]]
|
||||
name = "safetensors"
|
||||
version = "0.7.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "675656c1eabb620b921efea4f9199f97fc86e36dd6ffd1fbbe48d0f59a4987f5"
|
||||
dependencies = [
|
||||
"hashbrown 0.16.1",
|
||||
"serde",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "same-file"
|
||||
version = "1.0.6"
|
||||
@@ -5119,6 +5402,12 @@ version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cd0b0ec5f1c1ca621c432a25813d8d60c88abe6d3e08a3eb9cf37d97a0fe3d73"
|
||||
|
||||
[[package]]
|
||||
name = "seq-macro"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc"
|
||||
|
||||
[[package]]
|
||||
name = "serde"
|
||||
version = "1.0.229"
|
||||
@@ -5646,6 +5935,20 @@ dependencies = [
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sysctl"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "01198a2debb237c62b6826ec7081082d951f46dbb64b0e8c7649a452230d1dfc"
|
||||
dependencies = [
|
||||
"bitflags 2.13.2",
|
||||
"byteorder",
|
||||
"enum-as-inner",
|
||||
"libc",
|
||||
"thiserror 1.0.69",
|
||||
"walkdir",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tag_ptr"
|
||||
version = "0.1.0"
|
||||
@@ -5937,6 +6240,12 @@ dependencies = [
|
||||
"rustc-hash 2.1.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "typed-path"
|
||||
version = "0.12.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "54ec2755506471f54baf5e4bb5499d8e510cc1caf271270e778a736bd2d7336c"
|
||||
|
||||
[[package]]
|
||||
name = "typed-path"
|
||||
version = "0.13.0"
|
||||
@@ -7475,6 +7784,18 @@ dependencies = [
|
||||
"syn 3.0.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zip"
|
||||
version = "7.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cc12baa6db2b15a140161ce53d72209dacea594230798c24774139b54ecaa980"
|
||||
dependencies = [
|
||||
"crc32fast",
|
||||
"indexmap",
|
||||
"memchr",
|
||||
"typed-path 0.12.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zip"
|
||||
version = "9.1.3"
|
||||
@@ -7487,7 +7808,7 @@ dependencies = [
|
||||
"indexmap",
|
||||
"memchr",
|
||||
"nix",
|
||||
"typed-path",
|
||||
"typed-path 0.13.0",
|
||||
"windows-sys 0.61.2",
|
||||
"zopfli",
|
||||
]
|
||||
|
||||
@@ -10,3 +10,8 @@ rust-version.workspace = true
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
|
||||
# Native only (D1): candle is not built for wasm32, where this crate is empty.
|
||||
[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
|
||||
candle-core = "=0.9.2"
|
||||
candle-nn = "=0.9.2"
|
||||
@@ -0,0 +1,359 @@
|
||||
// Modified work: ported from Hugging Face `transformers` `models/sam2/modeling_sam2.py`
|
||||
// (Sam2PatchEmbeddings, Sam2MultiScaleAttention, Sam2MultiScaleBlock, Sam2HieraDetModel;
|
||||
// Apache-2.0, Copyright The HuggingFace Inc. team) to Rust on candle. See NOTICE.
|
||||
//! The Hiera backbone: patch embedding, windowed positional embedding, and blocks with window or
|
||||
//! global attention and query pooling at the stage changes. Activations are NHWC like the reference.
|
||||
|
||||
use candle_core::{D, Tensor};
|
||||
use candle_nn::{LayerNorm, Linear, Module, ops::softmax};
|
||||
|
||||
use crate::nn::{Mlp, layer_norm, linear};
|
||||
use crate::weights::Weights;
|
||||
use crate::{Error, Result};
|
||||
|
||||
/// Hiera hyper-parameters (the released SAM 2.1 sizes differ only here).
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct HieraConfig {
|
||||
pub embed_dims: Vec<usize>,
|
||||
pub blocks: Vec<usize>,
|
||||
pub heads: Vec<usize>,
|
||||
pub windows: Vec<usize>,
|
||||
pub global_blocks: Vec<usize>,
|
||||
pub query_stride: usize,
|
||||
pub query_pool_stages: usize,
|
||||
pub mlp_ratio: usize,
|
||||
pub eps: f64,
|
||||
pub pos_background: usize,
|
||||
}
|
||||
|
||||
impl HieraConfig {
|
||||
/// `facebook/sam2.1-hiera-tiny`.
|
||||
pub fn tiny() -> HieraConfig {
|
||||
HieraConfig {
|
||||
embed_dims: vec![96, 192, 384, 768],
|
||||
blocks: vec![1, 2, 7, 2],
|
||||
heads: vec![1, 2, 4, 8],
|
||||
windows: vec![8, 4, 14, 7],
|
||||
global_blocks: vec![5, 7, 9],
|
||||
query_stride: 2,
|
||||
query_pool_stages: 3,
|
||||
mlp_ratio: 4,
|
||||
eps: 1e-6,
|
||||
pos_background: 7,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn at<T: Copy>(v: &[T], i: usize, what: &str) -> Result<T> {
|
||||
v.get(i).copied().ok_or_else(|| Error::Weights(format!("Hiera config has no {what} for stage {i}")))
|
||||
}
|
||||
|
||||
/// Max pool a `[B, H, W, C]` tensor by `stride` (kernel = stride, no padding, floor).
|
||||
fn do_pool(x: &Tensor, stride: usize) -> Result<Tensor> {
|
||||
let x = x.permute((0, 3, 1, 2))?.contiguous()?.max_pool2d_with_stride(stride, stride)?;
|
||||
Ok(x.permute((0, 2, 3, 1))?.contiguous()?)
|
||||
}
|
||||
|
||||
/// Split `[B, H, W, C]` into `[B * nH * nW, ws, ws, C]` windows, zero padding the bottom/right.
|
||||
fn window_partition(x: &Tensor, ws: usize) -> Result<(Tensor, (usize, usize))> {
|
||||
let (b, h, w, c) = x.dims4()?;
|
||||
if ws == 0 {
|
||||
return Err(Error::Input("window size 0".into()));
|
||||
}
|
||||
let (ph, pw) = ((ws - h % ws) % ws, (ws - w % ws) % ws);
|
||||
let x = x.pad_with_zeros(2, 0, pw)?.pad_with_zeros(1, 0, ph)?;
|
||||
let (nh, nw) = ((h + ph) / ws, (w + pw) / ws);
|
||||
let x = x.reshape((b, nh, ws, nw, ws, c))?.permute((0, 1, 3, 2, 4, 5))?.contiguous()?.reshape((b * nh * nw, ws, ws, c))?;
|
||||
Ok((x, (h + ph, w + pw)))
|
||||
}
|
||||
|
||||
fn window_unpartition(win: &Tensor, ws: usize, padded: (usize, usize), orig: (usize, usize)) -> Result<Tensor> {
|
||||
if ws == 0 {
|
||||
return Err(Error::Input("window size 0".into()));
|
||||
}
|
||||
let (nh, nw) = (padded.0 / ws, padded.1 / ws);
|
||||
let (n, _, _, c) = win.dims4()?;
|
||||
let b = n / (nh * nw).max(1);
|
||||
let x = win.reshape((b, nh, nw, ws, ws, c))?.permute((0, 1, 3, 2, 4, 5))?.contiguous()?.reshape((b, padded.0, padded.1, c))?;
|
||||
Ok(x.narrow(1, 0, orig.0)?.narrow(2, 0, orig.1)?.contiguous()?)
|
||||
}
|
||||
|
||||
struct Attention {
|
||||
qkv: Linear,
|
||||
proj: Linear,
|
||||
heads: usize,
|
||||
scale: f64,
|
||||
query_stride: Option<usize>,
|
||||
}
|
||||
|
||||
impl Attention {
|
||||
fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let (b, h, w, _) = x.dims4()?;
|
||||
let qkv = self.qkv.forward(x)?.reshape((b, h * w, 3, self.heads, ()))?;
|
||||
let part = |i: usize| -> Result<Tensor> { Ok(qkv.narrow(2, i, 1)?.squeeze(2)?) };
|
||||
let (mut q, k, v) = (part(0)?, part(1)?, part(2)?);
|
||||
let (mut oh, mut ow) = (h, w);
|
||||
if let Some(s) = self.query_stride {
|
||||
let pooled = do_pool(&q.reshape((b, h, w, ()))?, s)?;
|
||||
(oh, ow) = (pooled.dim(1)?, pooled.dim(2)?);
|
||||
q = pooled.reshape((b, oh * ow, self.heads, ()))?;
|
||||
}
|
||||
let q = q.transpose(1, 2)?.contiguous()?;
|
||||
let k = k.transpose(1, 2)?.contiguous()?;
|
||||
let v = v.transpose(1, 2)?.contiguous()?;
|
||||
let att = (q.matmul(&k.transpose(2, 3)?.contiguous()?)? * self.scale)?;
|
||||
let att = softmax(&att, D::Minus1)?;
|
||||
let out = att.matmul(&v)?.transpose(1, 2)?.contiguous()?.reshape((b, oh, ow, ()))?;
|
||||
Ok(self.proj.forward(&out)?)
|
||||
}
|
||||
}
|
||||
|
||||
/// One Hiera block: pre-norm attention (window or global, optionally pooling the queries) and MLP.
|
||||
pub struct Block {
|
||||
norm1: LayerNorm,
|
||||
attn: Attention,
|
||||
norm2: LayerNorm,
|
||||
mlp: Mlp,
|
||||
proj: Option<Linear>,
|
||||
window: usize,
|
||||
query_stride: Option<usize>,
|
||||
}
|
||||
|
||||
impl Block {
|
||||
fn load(w: &Weights, cfg: &HieraConfig, stage: usize, block_in_stage: usize, total: usize) -> Result<Block> {
|
||||
let prefix = format!("vision_encoder.backbone.blocks.{total}");
|
||||
let first = stage > 0 && block_in_stage == 0;
|
||||
let dim = at(&cfg.embed_dims, if first { stage - 1 } else { stage }, "embed dim")?;
|
||||
let dim_out = at(&cfg.embed_dims, stage, "embed dim")?;
|
||||
let window = if first { at(&cfg.windows, stage - 1, "window")? } else { at(&cfg.windows, stage, "window")? };
|
||||
let window = if cfg.global_blocks.contains(&total) { 0 } else { window };
|
||||
let query_stride = (stage > 0 && stage <= cfg.query_pool_stages && block_in_stage == 0).then_some(cfg.query_stride);
|
||||
let heads = at(&cfg.heads, stage, "heads")?;
|
||||
Ok(Block {
|
||||
norm1: layer_norm(w, &format!("{prefix}.layer_norm1"), dim, cfg.eps)?,
|
||||
attn: Attention {
|
||||
qkv: linear(w, &format!("{prefix}.attn.qkv"), dim, dim_out * 3)?,
|
||||
proj: linear(w, &format!("{prefix}.attn.proj"), dim_out, dim_out)?,
|
||||
heads,
|
||||
scale: ((dim_out / heads.max(1)) as f64).powf(-0.5),
|
||||
query_stride,
|
||||
},
|
||||
norm2: layer_norm(w, &format!("{prefix}.layer_norm2"), dim_out, cfg.eps)?,
|
||||
mlp: Mlp::load(w, &format!("{prefix}.mlp"), dim_out, dim_out * cfg.mlp_ratio)?,
|
||||
proj: (dim != dim_out).then(|| linear(w, &format!("{prefix}.proj"), dim, dim_out)).transpose()?,
|
||||
window,
|
||||
query_stride,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let h = self.norm1.forward(x)?;
|
||||
let residual = match (&self.proj, self.query_stride) {
|
||||
(Some(p), Some(s)) => do_pool(&p.forward(&h)?, s)?,
|
||||
(Some(p), None) => p.forward(&h)?,
|
||||
(None, _) => x.clone(),
|
||||
};
|
||||
let (hh, ww) = (h.dim(1)?, h.dim(2)?);
|
||||
let mut ws = self.window;
|
||||
let mut padded = (hh, ww);
|
||||
let h = if ws > 0 {
|
||||
let (win, pad) = window_partition(&h, ws)?;
|
||||
padded = pad;
|
||||
win
|
||||
} else {
|
||||
h
|
||||
};
|
||||
let mut h = self.attn.forward(&h)?;
|
||||
let mut orig = (hh, ww);
|
||||
if let Some(s) = self.query_stride {
|
||||
ws /= s.max(1);
|
||||
orig = (residual.dim(1)?, residual.dim(2)?);
|
||||
if ws > 0 {
|
||||
padded = (orig.0 + (ws - orig.0 % ws) % ws, orig.1 + (ws - orig.1 % ws) % ws);
|
||||
}
|
||||
}
|
||||
if ws > 0 {
|
||||
h = window_unpartition(&h, ws, padded, orig)?;
|
||||
}
|
||||
let h = (residual + h)?;
|
||||
let m = self.mlp.forward(&self.norm2.forward(&h)?)?;
|
||||
Ok((h + m)?)
|
||||
}
|
||||
}
|
||||
|
||||
/// PyTorch `F.interpolate(mode="bicubic", align_corners=False)` of `[C, in_h, in_w]` to `[C, out_h, out_w]`.
|
||||
fn bicubic_resize(src: &[f32], c: usize, in_hw: (usize, usize), out_hw: (usize, usize)) -> Vec<f32> {
|
||||
const A: f32 = -0.75;
|
||||
let near = |x: f32| ((A + 2.0) * x - (A + 3.0)) * x * x + 1.0;
|
||||
let far = |x: f32| ((A * x - 5.0 * A) * x + 8.0 * A) * x - 4.0 * A;
|
||||
let axis = |n_in: usize, n_out: usize| -> Vec<([usize; 4], [f32; 4])> {
|
||||
let scale = n_in as f32 / n_out.max(1) as f32;
|
||||
(0..n_out)
|
||||
.map(|o| {
|
||||
let s = (o as f32 + 0.5) * scale - 0.5;
|
||||
let i0 = s.floor();
|
||||
let t = s - i0;
|
||||
let w = [far(t + 1.0), near(t), near(1.0 - t), far(2.0 - t)];
|
||||
let last = n_in as i64 - 1;
|
||||
let idx = [-1i64, 0, 1, 2].map(|d| (i0 as i64 + d).clamp(0, last.max(0)) as usize);
|
||||
(idx, w)
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
let (ys, xs) = (axis(in_hw.0, out_hw.0), axis(in_hw.1, out_hw.1));
|
||||
let mut out = vec![0f32; c * out_hw.0 * out_hw.1];
|
||||
for ch in 0..c {
|
||||
let plane = src.get(ch * in_hw.0 * in_hw.1..(ch + 1) * in_hw.0 * in_hw.1).unwrap_or(&[]);
|
||||
for (oy, (yi, yw)) in ys.iter().enumerate() {
|
||||
for (ox, (xi, xw)) in xs.iter().enumerate() {
|
||||
let mut acc = 0f32;
|
||||
for (a, wy) in yi.iter().zip(yw) {
|
||||
for (b, wx) in xi.iter().zip(xw) {
|
||||
acc += plane.get(a * in_hw.1 + b).copied().unwrap_or(0.0) * wy * wx;
|
||||
}
|
||||
}
|
||||
if let Some(o) = out.get_mut((ch * out_hw.0 + oy) * out_hw.1 + ox) {
|
||||
*o = acc;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
pub struct Hiera {
|
||||
cfg: HieraConfig,
|
||||
patch_w: Tensor,
|
||||
patch_b: Tensor,
|
||||
pos_embed: Vec<f32>,
|
||||
pos_embed_window: Vec<f32>,
|
||||
window0: usize,
|
||||
hidden: usize,
|
||||
pub blocks: Vec<Block>,
|
||||
stage_ends: Vec<usize>,
|
||||
}
|
||||
|
||||
impl Hiera {
|
||||
pub fn load(w: &Weights, cfg: &HieraConfig) -> Result<Hiera> {
|
||||
let hidden = at(&cfg.embed_dims, 0, "embed dim")?;
|
||||
let window0 = at(&cfg.windows, 0, "window")?;
|
||||
let bg = cfg.pos_background;
|
||||
let mut blocks = vec![];
|
||||
let mut stage_ends = vec![];
|
||||
for (stage, n) in cfg.blocks.iter().enumerate() {
|
||||
for b in 0..*n {
|
||||
blocks.push(Block::load(w, cfg, stage, b, blocks.len())?);
|
||||
}
|
||||
stage_ends.push(blocks.len().saturating_sub(1));
|
||||
}
|
||||
let flat = |t: Tensor| -> Result<Vec<f32>> { Ok(t.flatten_all()?.to_vec1::<f32>()?) };
|
||||
Ok(Hiera {
|
||||
cfg: cfg.clone(),
|
||||
patch_w: w.get("vision_encoder.backbone.patch_embed.projection.weight", &[hidden, 3, 7, 7])?,
|
||||
patch_b: w.get("vision_encoder.backbone.patch_embed.projection.bias", &[hidden])?,
|
||||
pos_embed: flat(w.get("vision_encoder.backbone.pos_embed", &[1, hidden, bg, bg])?)?,
|
||||
pos_embed_window: flat(w.get("vision_encoder.backbone.pos_embed_window", &[1, hidden, window0, window0])?)?,
|
||||
window0,
|
||||
hidden,
|
||||
blocks,
|
||||
stage_ends,
|
||||
})
|
||||
}
|
||||
|
||||
/// Patch embedding of `[B, 3, H, W]` pixels: `[B, H/4, W/4, C]` (before the positional embedding).
|
||||
pub fn patch_embed(&self, pixels: &Tensor) -> Result<Tensor> {
|
||||
let x = pixels.conv2d(&self.patch_w, 3, 4, 1, 1)?;
|
||||
let x = x.broadcast_add(&self.patch_b.reshape((1, self.hidden, 1, 1))?)?;
|
||||
Ok(x.permute((0, 2, 3, 1))?.contiguous()?)
|
||||
}
|
||||
|
||||
/// Background embedding resized to the grid plus the tiled window embedding, `[1, H, W, C]`.
|
||||
fn pos(&self, h: usize, w: usize) -> Result<Tensor> {
|
||||
let bg = self.cfg.pos_background;
|
||||
let big = bicubic_resize(&self.pos_embed, self.hidden, (bg, bg), (h, w));
|
||||
let ws = self.window0.max(1);
|
||||
let mut out = vec![0f32; h * w * self.hidden];
|
||||
for c in 0..self.hidden {
|
||||
for y in 0..h {
|
||||
for x in 0..w {
|
||||
let tile = self.pos_embed_window.get((c * ws + y % ws) * ws + x % ws).copied().unwrap_or(0.0);
|
||||
let base = big.get((c * h + y) * w + x).copied().unwrap_or(0.0);
|
||||
if let Some(o) = out.get_mut((y * w + x) * self.hidden + c) {
|
||||
*o = base + tile;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Tensor::from_vec(out, (1, h, w, self.hidden), &candle_core::Device::Cpu)?)
|
||||
}
|
||||
|
||||
/// Run all blocks. Returns the output of the last block of each stage (NHWC).
|
||||
pub fn forward(&self, pixels: &Tensor) -> Result<Vec<Tensor>> {
|
||||
let (_, _, ph, pw) = pixels.dims4()?;
|
||||
if ph == 0 || pw == 0 || ph % 4 != 0 || pw % 4 != 0 {
|
||||
return Err(Error::Input(format!("image size {pw}x{ph} must be a positive multiple of 4")));
|
||||
}
|
||||
let mut x = self.patch_embed(pixels)?;
|
||||
let (h, w) = (x.dim(1)?, x.dim(2)?);
|
||||
if h % self.window0 != 0 || w % self.window0 != 0 {
|
||||
return Err(Error::Input(format!("patch grid {w}x{h} must be a multiple of {}", self.window0)));
|
||||
}
|
||||
x = x.broadcast_add(&self.pos(h, w)?)?;
|
||||
let mut stages = vec![];
|
||||
for (i, b) in self.blocks.iter().enumerate() {
|
||||
x = b.forward(&x)?;
|
||||
if self.stage_ends.contains(&i) {
|
||||
stages.push(x.clone());
|
||||
}
|
||||
}
|
||||
Ok(stages)
|
||||
}
|
||||
|
||||
/// Run the blocks, returning every block's output (for the stage-by-stage reference test).
|
||||
pub fn forward_all(&self, pixels: &Tensor) -> Result<(Tensor, Vec<Tensor>)> {
|
||||
let emb = self.patch_embed(pixels)?;
|
||||
let (h, w) = (emb.dim(1)?, emb.dim(2)?);
|
||||
let mut x = emb.broadcast_add(&self.pos(h, w)?)?;
|
||||
let mut outs = vec![];
|
||||
for b in &self.blocks {
|
||||
x = b.forward(&x)?;
|
||||
outs.push(x.clone());
|
||||
}
|
||||
Ok((emb, outs))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use candle_core::Device;
|
||||
|
||||
#[test]
|
||||
fn bicubic_of_a_constant_is_that_constant() {
|
||||
let out = bicubic_resize(&[2.5; 49], 1, (7, 7), (32, 32));
|
||||
assert!(out.iter().all(|v| (v - 2.5).abs() < 1e-5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bicubic_identity_size_is_identity() {
|
||||
let src: Vec<f32> = (0..49).map(|i| (i as f32 * 0.37).sin()).collect();
|
||||
let out = bicubic_resize(&src, 1, (7, 7), (7, 7));
|
||||
assert!(src.iter().zip(&out).all(|(a, b)| (a - b).abs() < 1e-5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn windows_round_trip_with_padding() {
|
||||
let x = Tensor::arange(0f32, 60.0, &Device::Cpu).unwrap().reshape((1, 5, 6, 2)).unwrap();
|
||||
let (win, pad) = window_partition(&x, 4).unwrap();
|
||||
assert_eq!((win.dims(), pad), (&[4, 4, 4, 2][..], (8, 8)));
|
||||
let back = window_unpartition(&win, 4, pad, (5, 6)).unwrap();
|
||||
let d = (back - x).unwrap().abs().unwrap().sum_all().unwrap().to_scalar::<f32>().unwrap();
|
||||
assert_eq!(d, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_window_is_an_error_not_a_panic() {
|
||||
let x = Tensor::zeros((1, 4, 4, 2), candle_core::DType::F32, &Device::Cpu).unwrap();
|
||||
assert!(window_partition(&x, 0).is_err());
|
||||
}
|
||||
}
|
||||
@@ -11,6 +11,10 @@
|
||||
|
||||
use std::fmt;
|
||||
|
||||
pub mod hiera;
|
||||
pub mod nn;
|
||||
pub mod weights;
|
||||
|
||||
/// Everything that can go wrong in this crate. Never a panic.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum Error {
|
||||
@@ -22,6 +26,8 @@ pub enum Error {
|
||||
Cancelled,
|
||||
/// A time or memory budget was exceeded.
|
||||
Budget(String),
|
||||
/// A tensor operation failed (shape mismatch, out of memory, device error).
|
||||
Compute(String),
|
||||
}
|
||||
|
||||
impl fmt::Display for Error {
|
||||
@@ -31,12 +37,19 @@ impl fmt::Display for Error {
|
||||
Error::Input(m) => write!(f, "SAM 2 input: {m}"),
|
||||
Error::Cancelled => write!(f, "SAM 2: cancelled"),
|
||||
Error::Budget(m) => write!(f, "SAM 2 budget: {m}"),
|
||||
Error::Compute(m) => write!(f, "SAM 2 compute: {m}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for Error {}
|
||||
|
||||
impl From<candle_core::Error> for Error {
|
||||
fn from(e: candle_core::Error) -> Error {
|
||||
Error::Compute(e.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
pub type Result<T> = std::result::Result<T, Error>;
|
||||
|
||||
/// The released SAM 2.1 sizes. Only Tiny is wired up first; the others are configuration.
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
// Modified work: small layers used by the ported SAM 2 modules (Apache-2.0, see NOTICE).
|
||||
//! Linear, layer norm and a few tensor helpers on candle.
|
||||
|
||||
use candle_core::Tensor;
|
||||
use candle_nn::{LayerNorm, Linear, Module};
|
||||
|
||||
use crate::Result;
|
||||
use crate::weights::Weights;
|
||||
|
||||
pub fn linear(w: &Weights, prefix: &str, input: usize, output: usize) -> Result<Linear> {
|
||||
Ok(Linear::new(w.get(&format!("{prefix}.weight"), &[output, input])?, Some(w.get(&format!("{prefix}.bias"), &[output])?)))
|
||||
}
|
||||
|
||||
pub fn layer_norm(w: &Weights, prefix: &str, dim: usize, eps: f64) -> Result<LayerNorm> {
|
||||
Ok(LayerNorm::new(w.get(&format!("{prefix}.weight"), &[dim])?, w.get(&format!("{prefix}.bias"), &[dim])?, eps))
|
||||
}
|
||||
|
||||
/// Two-layer MLP with exact (erf) GELU, as `Sam2FeedForward` with `num_layers = 2`.
|
||||
pub struct Mlp {
|
||||
pub proj_in: Linear,
|
||||
pub proj_out: Linear,
|
||||
}
|
||||
|
||||
impl Mlp {
|
||||
pub fn load(w: &Weights, prefix: &str, dim: usize, hidden: usize) -> Result<Mlp> {
|
||||
Ok(Mlp { proj_in: linear(w, &format!("{prefix}.proj_in"), dim, hidden)?, proj_out: linear(w, &format!("{prefix}.proj_out"), hidden, dim)? })
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
Ok(self.proj_out.forward(&self.proj_in.forward(x)?.gelu_erf()?)?)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
// Modified work: loads the safetensors checkpoint of Hugging Face `transformers` SAM 2
|
||||
// (Apache-2.0, Copyright The HuggingFace Inc. team). See NOTICE.
|
||||
//! Checkpoint loading with validation: every tensor is f32 and finite, every lookup checks its
|
||||
//! shape, so a truncated, mislabelled or poisoned file is an error, not a crash or garbage.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use candle_core::{DType, Device, Tensor};
|
||||
|
||||
use crate::{Error, Result};
|
||||
|
||||
/// Named tensors of one checkpoint.
|
||||
pub struct Weights {
|
||||
map: HashMap<String, Tensor>,
|
||||
}
|
||||
|
||||
impl Weights {
|
||||
/// Read a `.safetensors` file; rejects non-finite values and unsupported dtypes.
|
||||
pub fn open(path: &Path) -> Result<Weights> {
|
||||
let map = candle_core::safetensors::load(path, &Device::Cpu).map_err(|e| Error::Weights(format!("{}: {e}", path.display())))?;
|
||||
Weights::from_map(map)
|
||||
}
|
||||
|
||||
/// Validate an already-loaded set of tensors.
|
||||
pub fn from_map(map: HashMap<String, Tensor>) -> Result<Weights> {
|
||||
let mut out = HashMap::with_capacity(map.len());
|
||||
for (name, t) in map {
|
||||
let t = match t.dtype() {
|
||||
DType::F32 => t,
|
||||
DType::F16 | DType::BF16 | DType::F64 => t.to_dtype(DType::F32)?,
|
||||
other => return Err(Error::Weights(format!("`{name}` has unsupported dtype {other:?}"))),
|
||||
};
|
||||
let flat = t.flatten_all()?.to_vec1::<f32>()?;
|
||||
if flat.iter().any(|v| !v.is_finite()) {
|
||||
return Err(Error::Weights(format!("`{name}` contains NaN or infinity")));
|
||||
}
|
||||
out.insert(name, t);
|
||||
}
|
||||
Ok(Weights { map: out })
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.map.len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.map.is_empty()
|
||||
}
|
||||
|
||||
/// The tensor `name`, which must have exactly `shape`.
|
||||
pub fn get(&self, name: &str, shape: &[usize]) -> Result<Tensor> {
|
||||
let t = self.map.get(name).ok_or_else(|| Error::Weights(format!("missing tensor `{name}`")))?;
|
||||
if t.dims() != shape {
|
||||
return Err(Error::Weights(format!("`{name}` has shape {:?}, expected {shape:?}", t.dims())));
|
||||
}
|
||||
Ok(t.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn one(name: &str, t: Tensor) -> HashMap<String, Tensor> {
|
||||
HashMap::from([(name.to_string(), t)])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_missing_wrong_shape_and_non_finite() {
|
||||
let ok = Tensor::zeros((2, 3), DType::F32, &Device::Cpu).unwrap();
|
||||
let w = Weights::from_map(one("a", ok)).unwrap();
|
||||
assert!(w.get("a", &[2, 3]).is_ok());
|
||||
assert!(matches!(w.get("a", &[3, 2]), Err(Error::Weights(_))));
|
||||
assert!(matches!(w.get("b", &[1]), Err(Error::Weights(_))));
|
||||
let bad = Tensor::new(&[1.0f32, f32::NAN], &Device::Cpu).unwrap();
|
||||
assert!(matches!(Weights::from_map(one("n", bad)), Err(Error::Weights(_))));
|
||||
let inf = Tensor::new(&[f32::INFINITY], &Device::Cpu).unwrap();
|
||||
assert!(Weights::from_map(one("i", inf)).is_err());
|
||||
let ints = Tensor::new(&[1u32], &Device::Cpu).unwrap();
|
||||
assert!(Weights::from_map(one("u", ints)).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_truncated_file_is_an_error() {
|
||||
let dir = std::env::temp_dir().join(format!("sam2-weights-{}", std::process::id()));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let p = dir.join("w.safetensors");
|
||||
let t = Tensor::zeros((4, 4), DType::F32, &Device::Cpu).unwrap();
|
||||
candle_core::safetensors::save(&one("x", t), &p).unwrap();
|
||||
let bytes = std::fs::read(&p).unwrap();
|
||||
std::fs::write(&p, &bytes[..bytes.len() - 8]).unwrap();
|
||||
assert!(matches!(Weights::open(&p), Err(Error::Weights(_))));
|
||||
assert!(matches!(Weights::open(&dir.join("absent")), Err(Error::Weights(_))));
|
||||
std::fs::remove_dir_all(&dir).ok();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
//! Stage-by-stage comparison with the PyTorch / transformers reference (`tools/sam2_reference.py`).
|
||||
//!
|
||||
//! Needs the `facebook/sam2.1-hiera-tiny` weights (default `<repo>/sam2-tiny`, or the folder in
|
||||
//! `EFFECTCRAFT_SAM2_DIR`) and the fixture tensors (`<repo>/plan/sam2-ref/image.safetensors`, or
|
||||
//! `EFFECTCRAFT_SAM2_REF`). Without them every test here skips, so CI stays green.
|
||||
#![cfg(not(target_arch = "wasm32"))]
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use candle_core::{Device, Tensor};
|
||||
use effectcraft_sam2::hiera::{Hiera, HieraConfig};
|
||||
use effectcraft_sam2::weights::Weights;
|
||||
|
||||
fn repo() -> PathBuf {
|
||||
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../..")
|
||||
}
|
||||
|
||||
fn weights_file() -> PathBuf {
|
||||
std::env::var_os("EFFECTCRAFT_SAM2_DIR").map_or_else(|| repo().join("sam2-tiny"), PathBuf::from).join("model.safetensors")
|
||||
}
|
||||
|
||||
fn fixture_file() -> PathBuf {
|
||||
std::env::var_os("EFFECTCRAFT_SAM2_REF").map_or_else(|| repo().join("plan/sam2-ref/image.safetensors"), PathBuf::from)
|
||||
}
|
||||
|
||||
/// The weights and the reference tensors, or `None` (the test then skips).
|
||||
fn load() -> Option<(Weights, HashMap<String, Tensor>)> {
|
||||
let (wf, rf) = (weights_file(), fixture_file());
|
||||
if !wf.exists() || !rf.exists() {
|
||||
eprintln!("skipped: SAM 2 weights or reference tensors not found ({} / {})", wf.display(), rf.display());
|
||||
return None;
|
||||
}
|
||||
let w = Weights::open(&wf).unwrap();
|
||||
let r = candle_core::safetensors::load(&rf, &Device::Cpu).unwrap();
|
||||
Some((w, r))
|
||||
}
|
||||
|
||||
fn max_abs(a: &Tensor, b: &Tensor) -> f32 {
|
||||
assert_eq!(a.dims(), b.dims(), "shape mismatch");
|
||||
(a - b).unwrap().abs().unwrap().flatten_all().unwrap().max(0).unwrap().to_scalar::<f32>().unwrap()
|
||||
}
|
||||
|
||||
/// f32 activations differ from PyTorch by accumulation order only; the Hiera blocks stay well
|
||||
/// inside this bound (measured max 2.6e-5, 2026-10-15).
|
||||
const HIERA_TOL: f32 = 1e-4;
|
||||
|
||||
#[test]
|
||||
fn hiera_matches_pytorch_block_by_block() {
|
||||
let Some((w, r)) = load() else { return };
|
||||
let hiera = Hiera::load(&w, &HieraConfig::tiny()).unwrap();
|
||||
let pixels = r.get("input.pixel_values").unwrap();
|
||||
let (emb, blocks) = hiera.forward_all(pixels).unwrap();
|
||||
let e = max_abs(&emb, r.get("vision_encoder.backbone.patch_embed.0").unwrap());
|
||||
eprintln!("patch_embed max abs error {e:e}");
|
||||
assert!(e < 1e-4, "patch embed error {e}");
|
||||
assert_eq!(blocks.len(), 12);
|
||||
for (i, b) in blocks.iter().enumerate() {
|
||||
let want = r.get(&format!("vision_encoder.backbone.blocks.{i}.0")).unwrap();
|
||||
let e = max_abs(b, want);
|
||||
eprintln!("block {i}: max abs error {e:e}");
|
||||
assert!(e < HIERA_TOL, "block {i} error {e} (tolerance {HIERA_TOL})");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hiera_is_deterministic_and_rejects_bad_sizes() {
|
||||
let Some((w, r)) = load() else { return };
|
||||
let hiera = Hiera::load(&w, &HieraConfig::tiny()).unwrap();
|
||||
let pixels = r.get("input.pixel_values").unwrap();
|
||||
let a = hiera.forward(pixels).unwrap();
|
||||
let b = hiera.forward(pixels).unwrap();
|
||||
assert_eq!(a.len(), 4);
|
||||
for (x, y) in a.iter().zip(&b) {
|
||||
assert_eq!(max_abs(x, y), 0.0);
|
||||
}
|
||||
let odd = pixels.narrow(2, 0, 1022).unwrap();
|
||||
assert!(hiera.forward(&odd).is_err());
|
||||
}
|
||||
Reference in new issue
Block a user