SAM 2.1 Hiera-Tiny — LiteRT (CompiledModel GPU)

SAM 2.1 (Segment Anything 2, Meta) Hiera-Tiny converted to LiteRT and running fully on the GPU via the CompiledModel API (ML Drift). Tap a point on an image and it returns a segmentation mask — the image encoder runs once per image, the mask decoder runs per point.

Both graphs are fully GPU-accelerated on the Pixel 8a (Mali / ML Drift) and on Apple silicon (Metal), and the output is bit-exact (corr 1.0) vs the original PyTorch SAM 2.1.

Files

Preprocessing: resize to 1024×1024, ImageNet mean [0.485, 0.456, 0.406] / std [0.229, 0.224, 0.225], NCHW.

GPU compatibility

The Hiera image encoder is made GPU-clean with three numerically-identical rewrites (done at conversion time; the SAM 2 mask decoder converts unchanged):

1. Bake the windowed positional embedding (constant for a fixed 1024² input) — removes the bicubic interpolate (GATHER_ND) and the tiled window embed (BROADCAST_TO).
2. 4-D window partition / unpartition — the 6-D view+permute becomes split-H → transpose → split-W (ML Drift rejects > 4-D tensors).
3. 4-D multi-scale attention — the 5-D fused qkv reshape becomes a channel-wise q/k/v slice.

Usage (Kotlin, LiteRT CompiledModel)

import com.google.ai.edge.litert.Accelerator
import com.google.ai.edge.litert.CompiledModel

val encoder = CompiledModel.create(
    context.assets, "sam2_encoder.tflite", CompiledModel.Options(Accelerator.GPU), null)
val decoder = CompiledModel.create(
    context.assets, "sam2_decoder.tflite", CompiledModel.Options(Accelerator.GPU), null)

// Encode once per image (input = normalized NCHW floats).
val encIn = encoder.createInputBuffers()
encIn[0].writeFloat(inputFloats)                 // 3 * 1024 * 1024
val flat = encoder.run(encIn)[0].readFloat()     // [image_embed | fpn0 | fpn1]

// Build the flat decoder input [image_embed | sparse | fpn0 | fpn1] (sparse = point encoding
// from sam2_prompt.bin), then run the decoder per tap.
val decIn = decoder.createInputBuffers()
decIn[0].writeFloat(flatDecoderInput)
val masks = decoder.run(decIn)[0].readFloat()    // (3, 256, 256) logits; mask > 0 = foreground

Usage (Python, verify the graph)

from ai_edge_litert.interpreter import Interpreter
import numpy as np

enc = Interpreter(model_path="sam2_encoder.tflite"); enc.allocate_tensors()
enc.set_tensor(enc.get_input_details()[0]["index"], pixels_nchw.astype(np.float32))  # [1,3,1024,1024]
enc.invoke()
flat = enc.get_tensor(enc.get_output_details()[0]["index"]).flatten()  # image_embed | fpn0 | fpn1

Video tracking (SAM 2 video path)

The full SAM 2.1 tracking loop — memory attention, memory encoder, object pointers and the prompt-conditioned mask decoder — as four fixed-shape per-frame graphs on the CompiledModel GPU. The rolling memory bank and per-frame orchestration run on the host (Kotlin/Swift/Python); only tensor math touches the GPU. Tap once on the first frame and the mask follows the object.

Why it works on the GPU: SAM 2's memory attention runs its RoPE attention with the batch dim collapsed (rank 3), which the ML Drift delegate silently mis-computes — the graphs here are re-authored batch-first (rank 4), numerically identical on the host and correct on the GPU (exact under fp32 GPU compute). The residual fp16 accumulation over the memory keys does not reach the mask.

Fidelity: the assembled loop matches the PyTorch Sam2VideoModel reference at min mask-IoU 0.9999 over a 10-frame clip (7- and 2-slot banks). All four graphs are fully GPU-resident with no CPU fallback — Pixel 8a (Mali): encode 828/828, memcond 480/480, decode 462/462, memorize 145/145 nodes; iPhone 17 Pro (Metal): all fullyGPU.

Per tracked frame (encode + memcond + decode + memorize): iPhone 17 Pro ~471 ms (2-slot) / ~751 ms (7-slot); Pixel 8a ~1.0–1.5 s.

Usage (Python, CompiledModel)

import numpy as np
from ai_edge_litert.compiled_model import CompiledModel

enc = CompiledModel.from_file("sam2v_encode.tflite")
ins, outs = enc.create_input_buffers(0), enc.create_output_buffers(0)
ins[0].write(np.ascontiguousarray(frame_nchw.ravel().astype(np.float32)))
enc.run_by_index(0, ins, outs)
flat = outs[0].read(4_194_304, np.float32)  # pix_raw | hi0 | hi1
# memcond -> decode -> memorize per frame; the full host loop (bank assembly,
# best-IoU pick, no-object handling) is verify_video.py in the recipe below.

Usage (Kotlin)

The Android tracker (Sam2VideoTracker.kt: filesDir load, rolling bank, per-frame loop) and the demo app live in LiteRT-Models → sam2/. The conversion + verification recipe is in litert-samples: models/sam2/sam2_hiera_tiny_video/converted.

Conversion

Converted with litert-torch from the Hugging Face transformers SAM 2 model. The full conversion script (and Android sample app) is in LiteRT-Models → sam2/.

License & credits

Apache-2.0, following the original SAM 2 (Meta, Apache-2.0). Conversion by @john-rocky.

---

Want a different model on-device? Open a request — free, open weights only; the export and its measured numbers get published publicly.