From 92b3a34732c03f7a9cc44df39568c3f228245a3d Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:15:51 +0200 Subject: [PATCH] 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 --- Cargo.lock | 341 ++++++++++++++++++++++++++++++- crates/sam2/Cargo.toml | 5 + crates/sam2/src/hiera.rs | 359 +++++++++++++++++++++++++++++++++ crates/sam2/src/lib.rs | 13 ++ crates/sam2/src/nn.rs | 32 +++ crates/sam2/src/weights.rs | 97 +++++++++ crates/sam2/tests/reference.rs | 79 ++++++++ 7 files changed, 916 insertions(+), 10 deletions(-) create mode 100644 crates/sam2/src/hiera.rs create mode 100644 crates/sam2/src/nn.rs create mode 100644 crates/sam2/src/weights.rs create mode 100644 crates/sam2/tests/reference.rs diff --git a/Cargo.lock b/Cargo.lock index abf1f898b..afb66fd53 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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", ] diff --git a/crates/sam2/Cargo.toml b/crates/sam2/Cargo.toml index 24c317bc1..f32a4e4e6 100644 --- a/crates/sam2/Cargo.toml +++ b/crates/sam2/Cargo.toml @@ -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" diff --git a/crates/sam2/src/hiera.rs b/crates/sam2/src/hiera.rs new file mode 100644 index 000000000..5091d331b --- /dev/null +++ b/crates/sam2/src/hiera.rs @@ -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, + pub blocks: Vec, + pub heads: Vec, + pub windows: Vec, + pub global_blocks: Vec, + 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(v: &[T], i: usize, what: &str) -> Result { + 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 { + 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 { + 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, +} + +impl Attention { + fn forward(&self, x: &Tensor) -> Result { + let (b, h, w, _) = x.dims4()?; + let qkv = self.qkv.forward(x)?.reshape((b, h * w, 3, self.heads, ()))?; + let part = |i: usize| -> Result { 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, + window: usize, + query_stride: Option, +} + +impl Block { + fn load(w: &Weights, cfg: &HieraConfig, stage: usize, block_in_stage: usize, total: usize) -> Result { + 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 { + 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 { + 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, + pos_embed_window: Vec, + window0: usize, + hidden: usize, + pub blocks: Vec, + stage_ends: Vec, +} + +impl Hiera { + pub fn load(w: &Weights, cfg: &HieraConfig) -> Result { + 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> { Ok(t.flatten_all()?.to_vec1::()?) }; + 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 { + 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 { + 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> { + 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)> { + 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 = (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::().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()); + } +} diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index d41de59fa..246739e45 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -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 for Error { + fn from(e: candle_core::Error) -> Error { + Error::Compute(e.to_string()) + } +} + pub type Result = std::result::Result; /// The released SAM 2.1 sizes. Only Tiny is wired up first; the others are configuration. diff --git a/crates/sam2/src/nn.rs b/crates/sam2/src/nn.rs new file mode 100644 index 000000000..ef334432a --- /dev/null +++ b/crates/sam2/src/nn.rs @@ -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 { + 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 { + 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 { + 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 { + Ok(self.proj_out.forward(&self.proj_in.forward(x)?.gelu_erf()?)?) + } +} diff --git a/crates/sam2/src/weights.rs b/crates/sam2/src/weights.rs new file mode 100644 index 000000000..c8b79cc70 --- /dev/null +++ b/crates/sam2/src/weights.rs @@ -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, +} + +impl Weights { + /// Read a `.safetensors` file; rejects non-finite values and unsupported dtypes. + pub fn open(path: &Path) -> Result { + 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) -> Result { + 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::()?; + 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 { + 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 { + 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(); + } +} diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs new file mode 100644 index 000000000..e5b7ae596 --- /dev/null +++ b/crates/sam2/tests/reference.rs @@ -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 `/sam2-tiny`, or the folder in +//! `EFFECTCRAFT_SAM2_DIR`) and the fixture tensors (`/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)> { + 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::().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()); +}