README.md

---
license: apache-2.0
library_name: litert
pipeline_tag: mask-generation
tags:
  - litert
  - tflite
  - sam2
  - segment-anything
  - mask-decoder
  - interactive-segmentation
  - on-device
  - gpu
base_model: facebook/sam2.1-hiera-tiny
---

# SAM 2.1 (Hiera-Tiny) mask decoder — LiteRT GPU

On-device **LiteRT / TFLite** conversion of the **prompt-conditioned mask decoder** of
[**SAM 2.1 Hiera-Tiny**](https://huggingface.co/facebook/sam2.1-hiera-tiny) (Meta, Apache-2.0),
running **fully on the mobile GPU** via the LiteRT `CompiledModel` API (ML Drift / `LITERT_CL` delegate).

![SAM 2.1 tap-to-segment running on-device (LiteRT)](demo.gif)

This is the lightweight, per-click half of the SAM 2 image path. Pair it with the
[**SAM 2.1 Hiera-Tiny image encoder**](https://huggingface.co/litert-community/SAM2.1-Hiera-Tiny-Image-Encoder)
(run **once** per image, ~7 ms): the encoder produces the multi-scale feature pyramid, and this decoder
turns a **point prompt** into segmentation masks **per tap** (a few ms each) — interactive "tap to segment".

| | |
|---|---|
| Task | Mask decoder for promptable segmentation (SAM 2 image path) |
| Architecture | 2-layer two-way transformer (token↔image cross-attention) + mask up-sampler |
| Inputs | `image_embeddings [1,256,64,64]`, `sparse_prompt [1,2,256]`, `feat_s1 [1,64,128,128]`, `feat_s0 [1,32,256,256]` |
| Outputs | `pred_masks [1,3,256,256]` (logits, 3 multimask candidates), `iou_scores [1,3]` |
| Precision / size | FP16, **17 MB** |
| File | **`sam2_tiny_mask_decoder_v2_fp16.tflite`** (recommended). `sam2_tiny_mask_decoder_fp16.tflite` is the earlier build, kept for reference — see the note below. |
| Device | Pixel 8a — **fully GPU** (LITERT_CL), correct masks, ~7 ms/tap |
| Op set | `banned ops = NONE`, `>4-D tensors = 0` (BATCH_MATMUL ×15, SOFTMAX ×7, GELU ×2, CONV_2D ×2) |

> **⚠ Residency ≠ correctness — and why v2 exists.** The first build (`sam2_tiny_mask_decoder_fp16.tflite`)
> fully delegated to the GPU (358/358 LITERT_CL nodes, `banned ops = NONE`, `>4-D = 0`, desktop parity
> corr 1.0) yet returned **silently wrong masks on the Pixel 8a GPU** (corr 0.265 vs CPU; a face tap at
> IoU ≈ 0.62 on CPU collapsed to ≈ 0.10 with the mask on the background).
>
> The cause was found by device A/B bisection: its attention was written with the **batch dim collapsed**
> (q/k/v shaped `[heads, N, d]`, rank 3). The GPU delegate mis-computes that form. It is **not** an fp16
> problem (forcing fp32 GPU compute still gives corr 0.473) and **not** LayerNorm (plain and
> overflow-safe LN give the same wrong result). The mask head's rank-2 matmul is innocent.
>
> **`v2` keeps the leading batch dim (rank-4 SDPA, `[1, heads, N, d]`).** Host numerics are identical
> (eager cos 0.999999); on the Pixel 8a GPU it restores **corr 0.9998 / binary-IoU 0.999** vs CPU and is
> **~20 % faster** (6.8 ms vs 8.5 ms). Inputs and outputs are unchanged, so v2 is a drop-in replacement.
> Note the companion encoder's rank-3 SDPA *is* GPU-correct — a healthy sibling graph proves nothing;
> only a numeric GPU-vs-CPU check on device catches this.

## Pipeline (how the inputs are produced)

```
RGB image ──> image encoder (run once) ──> image_embeddings[1,256,64,64], feat_s1[1,64,128,128], feat_s0[1,32,256,256]
tap (x,y in 1024-space) ──> prompt encode (host-side, see below) ──> sparse_prompt[1,2,256]
                                              │
       image_embeddings + feat_s0/s1 + sparse_prompt ──> THIS decoder ──> 3 masks + 3 IoU
       pick argmax(IoU) ──> upsample 256×256 logits to image size ──> threshold > 0 ──> overlay
```

The decoder uses the **encoder variant that already folds `conv_s0` / `conv_s1` + `no_memory`** so its
outputs are directly decoder-ready (no host reshaping between the two models).

### Host-side prompt encoding (single positive point)

The tiny point→token step (a sin/cos positional encoding) is done on the host to keep the GPU graph
sin/cos-free. For a positive click `(x, y)` in 1024×1024 model space, with the bundled constants
`posmat [2,128]`, `point_embed[1] [256]`, `not_a_point [256]`:

```
c      = (([x, y]) + 0.5) / 1024          # normalize, half-pixel shift
c      = 2*c - 1
coord  = 2*pi * (c @ posmat)              # [128]
token0 = concat(sin(coord), cos(coord)) + point_embed[1]   # the positive point
token1 = not_a_point                      # the padding point
sparse_prompt = [[token0, token1]]        # [1, 2, 256]
```

This matches the upstream `Sam2PromptEncoder` to ~3.7e-7.

## GPU-clean conversion (what was re-authored)

Converted with `litert-torch`, **model-side rewrites only — no converter patch**, each weights-faithful:

1. **Two-way attention (×7)**: re-expressed as **3-D batched SDPA** `[heads, N, d]` (a 4-D SDPA makes the
   delegate emit a `BROADCAST_TO`).
2. **Mask up-sampler `ConvTranspose2d` (×2)**: replaced with the exact **zero-stuff + `Conv2d`** identity
   (`TRANSPOSE_CONV` is rejected on Pixel 8a; this is numerically identical, not a bilinear approximation).
3. **Mask head**: the `hyper_in @ upscaled` mask projection is kept **≤4-D** (the upstream
   `[1,1,4,256,256]` 5-D tensor is collapsed; batch/point-batch are 1).
4. **LayerNorm (×9)**: scale-before-square **SafeLayerNorm** (fp16-overflow-safe, mathematically identical).
5. **Constants baked**: `image_positional_embeddings` and the no-mask dense prompt are baked as buffers.
6. **Multimask path**: static slice `[1:]` of the 3 candidate masks — no dynamic-stability
   `argmax` / `gather` / `where`.

## Fidelity (honest)

Eager re-authoring is **numerically exact** (`cos = 1.000`). End-to-end through the two FP16 tflite models
(encoder → host prompt-encode → decoder) vs the PyTorch reference, for a center click:

| Metric | value |
|---|---|
| mask logits cosine | **0.999999** |
| binary mask IoU (threshold 0) | **0.99964** |
| IoU-score head | ref `[0.936, 0.022, 0.399]` vs got `[0.936, 0.022, 0.399]` |

The deepest `64×64` image embedding drifts slightly on the GPU (true-fp16 deep attention; see the encoder
card). Mask boundaries are carried by the near-exact high-resolution features, so mask quality holds.

## Minimal usage

**Android (Kotlin, CompiledModel GPU)**

```kotlin
// once per image - encoder on GPU (decoder-ready v2 variant from the companion repo)
val enc = CompiledModel.create(context.assets, "sam2_tiny_image_encoder_v2_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
// per tap - decoder on GPU (v2: rank-4 attention, GPU-correct; see the note above)
val dec = CompiledModel.create(context.assets, "sam2_tiny_mask_decoder_v2_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
// dec inputs: 0 image_embeddings[1,256,64,64], 1 sparse[1,2,256],
//             2 feat_s1[1,64,128,128], 3 feat_s0[1,32,256,256]
// dec outputs: pred_masks[1,3,256,256] logits, iou_scores[1,3] -> argmax(iou), threshold 0
```

**Python (desktop verification)**

```python
MEAN = np.array([0.485, 0.456, 0.406], np.float32)
STD  = np.array([0.229, 0.224, 0.225], np.float32)
import numpy as np
from PIL import Image
from ai_edge_litert.interpreter import Interpreter

# 1) encode once (decoder-ready v2 encoder from the companion encoder repo)
img = Image.open("photo.jpg").convert("RGB").resize((1024, 1024))
x = ((np.asarray(img, np.float32) / 255 - MEAN) / STD).transpose(2, 0, 1)[None]
enc = Interpreter(model_path="sam2_tiny_image_encoder_v2_fp16.tflite"); enc.allocate_tensors()
enc.set_tensor(enc.get_input_details()[0]["index"], x); enc.invoke()
eo = {tuple(d["shape"]): enc.get_tensor(d["index"]) for d in enc.get_output_details()}

# 2) host prompt-encode one positive tap (px, py) in 1024-space (constants: this repo)
px, py = 512, 384
posmat, pe1, nap = np.split(np.fromfile("prompt_encode_const.bin", np.float32), [256, 512])
coord = 2 * np.pi * ((2 * (np.array([px, py], np.float32) + 0.5) / 1024 - 1) @ posmat.reshape(2, 128))
tok0 = np.concatenate([np.sin(coord), np.cos(coord)]) + pe1
sparse = np.stack([tok0, nap])[None].astype(np.float32)               # [1,2,256]

# 3) decode masks
dec = Interpreter(model_path="sam2_tiny_mask_decoder_v2_fp16.tflite"); dec.allocate_tensors()
feed = {(1,2,256): sparse}; feed.update(eo)                           # match inputs by shape
for d in dec.get_input_details(): dec.set_tensor(d["index"], feed[tuple(d["shape"])])
dec.invoke()
o = {len(d["shape"]): dec.get_tensor(d["index"]) for d in dec.get_output_details()}
masks, iou = o[4], o[2]                                               # [1,3,256,256], [1,3]
best = masks[0, iou[0].argmax()] > 0                                  # [256,256] binary mask
Image.fromarray(best.astype(np.uint8) * 255).resize(Image.open("photo.jpg").size).save("mask.png")
```

## Training data & PII

SAM 2 was trained by Meta on **SA-1B** (licensed photos) and **SA-V** (licensed videos) with
model-in-the-loop mask annotation. No new training was performed for this conversion — it is a
weights-faithful format change of the public `facebook/sam2.1-hiera-tiny` checkpoint. Because the source
data is real-world imagery it may incidentally contain people, faces, vehicles, signage and other PII; no
PII was deliberately collected and this conversion adds none. Apply your own content/PII filtering as
appropriate. See the [SAM 2 release](https://github.com/facebookresearch/sam2) and
[paper](https://arxiv.org/abs/2408.00714) for full dataset details.

## License

Apache-2.0, inherited from the upstream [SAM 2.1](https://huggingface.co/facebook/sam2.1-hiera-tiny).
This is a format conversion; all credit to the original authors (Meta AI).