mlboydaisuke / SAM2-hiera-tiny-LiteRT

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

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 .