SAM2.1 Hiera-Tiny LiteRT (CompiledModel GPU), corr 1.0

README.md

@@ -0,0 +1,91 @@
+---
+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).

sam2_decoder.tflite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f00e15545c887e20a2886cb88f6afab92ceb8e35a81e9fbf67629cf1de56902d
+size 16866608

sam2_encoder.tflite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:433d488bbf13116bf45f51e98efdfd2c047c48aa2693366079a1fdf73b101b1d
+size 79847760

sam2_prompt.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c1ac798f0cc0bd5e4b0dc94efe26f2f9dfe4d1d4cec141ab0350580ce7a48588
+size 3072