README.md

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

# SAM 2.1 Hiera-Tiny — LiteRT (CompiledModel GPU)

[SAM 2.1](https://ai.meta.com/sam2/) (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

| File | Size (fp16) | Input | Output | Runtime |
|---|---|---|---|---|
| `sam2_encoder.tflite` | 80 MB | `[1, 3, 1024, 1024]` NCHW | flat `[1, 4194304]` (`image_embed \| fpn0 \| fpn1`) | CompiledModel GPU |
| `sam2_decoder.tflite` | 17 MB | flat `[1, 4194816]` (`image_embed \| sparse \| fpn0 \| fpn1`) | masks `[1, 3, 256, 256]` | CompiledModel GPU |
| `sam2_prompt.bin` | 3 KB | — | prompt-encoder constants for the Kotlin point encoder | — |

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)

```kotlin
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)

```python
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
```

## 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](https://github.com/john-rocky/LiteRT-Models) → `sam2/`.

## License & credits

Apache-2.0, following the original [SAM 2](https://github.com/facebookresearch/sam2) (Meta,
Apache-2.0). Conversion by [@john-rocky](https://github.com/john-rocky).