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