| --- |
| license: apache-2.0 |
| base_model: |
| - facebook/sam2.1-hiera-tiny |
| tags: |
| - sam2 |
| - video |
| - onnx |
| - webgpu |
| - fp16 |
| --- |
| |
| # SAM 2.1 Hiera-Tiny video tracking, fp16, 512 input (ONNX) |
|
|
| The full SAM 2.1 video tracker, including the memory encoder, memory attention and |
| object pointers, as five fixed-shape ONNX graphs for ONNX Runtime on WebGPU. It is used by |
| [Diffusion Studio](https://github.com/diffusionstudio)'s object mask tool. The graphs were |
| exported from the `transformers` port of `facebook/sam2.1-hiera-tiny` |
| (`Sam2VideoModel`, transformers 5.17), following the graph layout of |
| [square-zero-labs/sam2.1-tiny-video-onnx](https://huggingface.co/square-zero-labs/sam2.1-tiny-video-onnx). |
|
|
| It is built for speed on integrated GPUs: |
|
|
| - **512×512 input** instead of 1024. Image features are 32×32 instead of 64×64. |
| - **fp16 weights and compute**, with float32 inputs and outputs (a cast at each graph |
| boundary). Input-independent position encodings are computed in float32 at export time. |
|
|
| Halving the input makes memory attention about 16 times less work and the image encoder |
| about 4 times less. The memory bank is SAM 2's own: the prompted frame, the 6 most recent |
| tracked frames and 16 object pointers. |
|
|
| The trade-off is detail. Small objects and thin structures are resolved less finely, and a |
| small object's identity is weaker: after a scene cut, the tracker is more likely to latch |
| onto a similar-looking object than to report the object as gone. |
|
|
| ## Speed |
|
|
| Measured through the full tracking loop on ONNX Runtime Web 1.30 with WebGPU on an Apple M1 |
| (8-core GPU), plugged in: |
|
|
| | | 1024 input, fp16 | 512 input, fp16 | |
| |---|---|---| |
| | per tracked frame | 1.3 s | 0.25 s | |
|
|
| The export script is `packages/sam2/scripts/export.py` in the Diffusion Studio repository. |
| It takes the input size and the number of memories as parameters. |
|
|
| ## Graphs |
|
|
| `F` = 32 (input / 16), `R` = 7·F² + 64 (frame memories, then 16 object pointers as 4 tokens each). |
|
|
| | file | inputs | outputs | |
| |---|---|---| |
| | `onnx/vision_encoder.onnx` | `pixel_values [1,3,512,512]` | `feats0 [1,32,128,128]`, `feats1 [1,64,64,64]`, `feats2 [1,256,F,F]`, `feats2_no_mem` (prompted-frame variant), `vision_pos_embed [1,256,F,F]` | |
| | `onnx/mask_decoder.onnx` | `feats0`, `feats1`, `feats2_cond`, `input_points [1,1,N,2]` (pixels of the 512 input), `input_labels [1,1,N] int32` | `low_res_mask [1,1,256,256]` (upsampled logits), `high_res_mask [1,1,512,512]`, `iou [1,1]`, `object_score_logits [1,1,1]`, `object_pointer [1,1,256]` | |
| | `onnx/memory_encoder.onnx` | `feats2`, `high_res_mask`, `object_score_logits [1,1]`, `binarize` (1 for prompted frames) | `memory_tokens [F²,1,64]`, `memory_pos [F²,1,64]` | |
| | `onnx/memory_attention.onnx` | `current_vision_features [F²,1,256]`, `current_vision_position_embeddings [F²,1,256]`, `memory [R,1,64]`, `memory_pos [R,1,64]` | `conditioned_feats [1,256,F,F]` | |
| | `onnx/pointer_tpos.onnx` | `normalized_diffs [16]` | `pointer_pos [16,64]` | |
|
|
| The mask decoder chooses its output as the video predictor does. It considers several |
| candidate masks when there is at most one real point (every tracked frame, and a single |
| click), and otherwise uses a single mask with the stability fallback. Masks and object |
| pointers are suppressed in-graph when `object_score_logits ≤ 0`. |
|
|
| `constants.json` carries the input size, the number of memories, the normalization |
| constants and the 7×64 temporal positional encoding table. The prompted frame's memory uses |
| row 6, and the memory `k` frames back uses row `k − 1`. |
|
|
