SAM 2.1 Hiera Tiny for ExecuTorch WebGPU

This repository contains an ExecuTorch conversion of 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).

Base model: 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
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.

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) 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), 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:

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