litert-community/SAM2.1-Hiera-Tiny-Mask-Decoder

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 (Meta, Apache-2.0), running fully on the mobile GPU via the LiteRT CompiledModel API (ML Drift / LITERT_CL delegate).

This is the lightweight, per-click half of the SAM 2 image path. Pair it with the SAM 2.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)

// 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)

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 and paper for full dataset details.

License

Apache-2.0, inherited from the upstream SAM 2.1. This is a format conversion; all credit to the original authors (Meta AI).