--- license: apache-2.0 base_model: - facebook/sam2.1-hiera-tiny library_name: executorch pipeline_tag: mask-generation inference: false tags: - executorch - webgpu - on-device - browser --- # SAM 2.1 Hiera Tiny for ExecuTorch WebGPU This repository contains an [ExecuTorch](https://github.com/pytorch/executorch) conversion of [facebook/sam2.1-hiera-tiny](https://huggingface.co/facebook/sam2.1-hiera-tiny) for the ExecuTorch **WebGPU** backend, so the model runs entirely in a web browser on the visitor's GPU. It powers the corresponding demo in the ExecuTorch WebGPU Playground (launching soon on [executorch.ai](https://executorch.ai)). - **Base model:** [facebook/sam2.1-hiera-tiny](https://huggingface.co/facebook/sam2.1-hiera-tiny) - **Format:** ExecuTorch `.pte` programs delegated to the WebGPU backend - **Quantization:** fp32 (no quantization) - **Total size:** 128.1 MB ## About the original model The tiny Hiera variant of SAM 2.1, the improved checkpoint release of SAM 2, a foundation model for promptable visual segmentation in images and videos. - **Developed by:** Meta FAIR - **Paper:** [SAM 2: Segment Anything in Images and Videos](https://arxiv.org/abs/2408.00714) - 39.0M parameters - Segments any object from point or box prompts For training data, evaluation results and full details, see the original model card: [facebook/sam2.1-hiera-tiny](https://huggingface.co/facebook/sam2.1-hiera-tiny). ## Intended use and limitations Interactive segmentation of objects in images from clicks or boxes. This conversion covers single-image segmentation only (no video memory). Masks can be imprecise on thin structures, low-contrast boundaries or heavily occluded objects. ## What changed in this conversion - Image encoder and mask decoder exported as two separate programs; the prompt encoder runs on the host from `prompt_*.bin` tables generated from the upstream weights - Fixed 512x512 input (upstream: 1024x1024), so masks are 128x128 (upstream: 256x256) - The encoder adds `no_memory_embedding` to the image embedding, as upstream does for single images - Three prompt tokens: a click is [point, padding, padding] and a box is [top-left, bottom-right, padding]. The decoder ignores a padding token that follows another padding token, so every prompt type matches upstream - No mask-prompt input. All weights fp32. Attention uses ExecuTorch's fused `et_vk.sdpa` op (within 5e-6 of PyTorch attention), and Hiera's windowed max-pooling runs as one batched pool (exact) ## Inputs and outputs - `sam2_encoder.pte`: `forward(pixels float32 [1,3,512,512])`, an RGB image resized to 512x512 (aspect ratio not kept), scaled to [0, 1] and ImageNet-normalized (mean 0.485/0.456/0.406, std 0.229/0.224/0.225) -> `image_embeddings [1,256,32,32]`, `hr0 [1,256,128,128]`, `hr1 [1,256,64,64]` (high-resolution features) - `sam2_decoder.pte`: `forward(image_embeddings, image_pe [1,256,32,32], sparse [1,1,3,256], dense [1,256,32,32], hr0, hr1)` -> `masks [1,3,128,128]` (logits of 3 candidate masks), `iou_scores [1,1,3]` (predicted IoU, after sigmoid), `object_score_logits [1,1,1]` - `image_pe` is `prompt_img_pos.bin` and `dense` is `prompt_dense.bin` (no mask prompt). Each sparse token is `[sin(2πcG), cos(2πcG)] + e`, where `G` [2,128] is `prompt_gaussian.bin`, `c` is the point mapped to the 512 grid (+0.5 pixel) and normalized to [-1, 1], and `e` is `prompt_point_embed_fg.bin` / `prompt_point_embed_bg.bin` / `prompt_box_tl.bin` / `prompt_box_br.bin`; a padding token is `prompt_not_a_point.bin` alone; fill unused slots with padding tokens - All tensors are float32, contiguous NCHW; `.bin` files are raw little-endian float32 ## Files | File | Size | SHA-256 | |---|---|---| | `prompt_box_br.bin` | 1.0 KB | `a0df40a29af4a9b6…` | | `prompt_box_tl.bin` | 1.0 KB | `96c208c44e95f704…` | | `prompt_dense.bin` | 1.0 MB | `52cece57b549dc81…` | | `prompt_gaussian.bin` | 1.0 KB | `c37a6a9361886286…` | | `prompt_img_pos.bin` | 1.0 MB | `89dc4182ee4ac6a3…` | | `prompt_not_a_point.bin` | 1.0 KB | `19771c067b922d9d…` | | `prompt_point_embed_bg.bin` | 1.0 KB | `91c432d10fcef8e3…` | | `prompt_point_embed_fg.bin` | 1.0 KB | `2e037dab5bee002d…` | | `sam2_decoder.pte` | 16.2 MB | `57b91034acc71be2…` | | `sam2_encoder.pte` | 109.9 MB | `30f9fb469b779558…` | Full checksums are in `manifest.json`. The playground verifies them after download. ## Running it These programs were exported with the ExecuTorch WebGPU backend ([backends/webgpu](https://github.com/pytorch/executorch/tree/main/backends/webgpu)) and checked against the upstream transformers model at 512x512 (click and box prompts give masks with IoU 1.000), and run in a browser on WebGPU with the WebAssembly runtime build used by the ExecuTorch WebGPU Playground (launching soon on [executorch.ai](https://executorch.ai)), where the masks match the CPU run with IoU 1.000. They have not yet been tested on a hardware GPU. That runtime build is not published yet; until it is, the playground is the supported way to run them. It downloads these files once, caches them in the browser, and runs everything locally. Requirements: a browser with WebGPU enabled. ## Provenance - **Upstream:** `facebook/sam2.1-hiera-tiny` at revision `de431c4043854a71d8101e17995dfe596bf101a5` (official) - **Checked against the official checkpoint:** `model.safetensors` (sha256 `48c14467e5cf9e51870511feb72c89688e82dd74523142c0538b663e193ac2a7`) has 469/469 tensors bit-identical to the authors' original checkpoint `sam2.1_hiera_tiny.pt` (sha256 `7402e0d864fa82708a20fbd15bc84245c2f26dff0eb43a4b5b93452deb34be69`) in the same repo at the same revision; `python export/_sam2.py verify-official` repeats the check - **Exported with:** executorch 1.5.1+cpu, torch 2.14.1+cpu, torchao 0.18.0+cpu, transformers 5.19.0 (executorch git `3b6068392324`) - **Reproduce:** `python recipes/sam2-hiera-tiny.py --out out/sam2.1-hiera-tiny` (recipe in `export/`) - Re-running the recipe produces byte-identical files. ## Citation If you use this model, please cite the original work: ```bibtex @article{ravi2024sam2, title={SAM 2: Segment Anything in Images and Videos}, author={Ravi, Nikhila and Gabeur, Valentin and Hu, Yuan-Ting and Hu, Ronghang and Ryali, Chaitanya and Ma, Tengyu and Khedr, Haitham and R{\"a}dle, Roman and Rolland, Chloe and Gustafson, Laura and Mintun, Eric and Pan, Junting and Alwala, Kalyan Vasudev and Carion, Nicolas and Wu, Chao-Yuan and Girshick, Ross and Doll{\'a}r, Piotr and Feichtenhofer, Christoph}, journal={arXiv preprint arXiv:2408.00714}, url={https://arxiv.org/abs/2408.00714}, year={2024} } ``` ## License This is a derivative of the upstream model and is distributed under the same terms (`apache-2.0`); see `LICENSE`. Modifications (conversion to the ExecuTorch format) by Meta Platforms, Inc.