litert-community/CLIPSeg-rd64-LiteRT

CLIPSeg rd64 — LiteRT on-device text-prompted segmentation

CLIPSeg (CVPR 2022, Apache-2.0) re-authored for LiteRT: type what you want to segment ("a cat", "the sky") and get a mask — no fixed class list. Three graphs — CLIP text and vision encoders on the CompiledModel GPU, the tiny 3-layer decoder on CPU (its 4-head/head_dim-16 attention fp16-miscomputes on the Mali delegate; the 12-head/head_dim-64 vision encoder survives at 0.998).

Input | prompt "a dog" | prompt "the grass" — the same image, two prompts, masks from the on-device model. Photo: "Lily the Golden Retriever in the grass" (Wikimedia Commons, Public Domain).

Verified on a Pixel 8a: text 761/761 GPU (8.7 ms) + vision 613/613 GPU (8.2 ms) + decoder CPU (exact); end-to-end device-vs-PyTorch logits corr 0.99998, mask IoU 0.9986.

Files

file	graph	delegate
clipseg_text_fp16.tflite	token-emb [1,77,512] → hidden [1,77,512]	GPU
clipseg_vision_fp16.tflite	image [1,3,352,352] → t3,t6,t9 [1,485,768]	GPU
clipseg_decoder.tflite (fp32)	t3,t6,t9,cond[512] → logits [1,352,352]	CPU
token_embedding_f16.bin, text_projection_f16.bin, vocab.json, merges.txt	host assets	—

Minimal usage (Python)

import numpy as np, torch
from PIL import Image
from transformers import CLIPSegProcessor
from ai_edge_litert.interpreter import Interpreter

proc = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
img = proc(images=Image.open("photo.jpg"), return_tensors="pt")["pixel_values"].numpy()  # [1,3,352,352]

vis = Interpreter("clipseg_vision_fp16.tflite"); vis.allocate_tensors()
vis.set_tensor(vis.get_input_details()[0]["index"], img); vis.invoke()
t = [vis.get_tensor(o["index"]) for o in sorted(vis.get_output_details(), key=lambda o: o["index"])]  # t3,t6,t9

# cond[512] from the text encoder (token-emb lookup -> text graph -> EOT row @ text_projection)
dec = Interpreter("clipseg_decoder.tflite"); dec.allocate_tensors()   # CPU (exact)
ins = dec.get_input_details()
for d, arr in zip(ins, [t[0], t[1], t[2], cond]):   # cond: [1,512] float32
    dec.set_tensor(d["index"], arr.astype(np.float32))
dec.invoke()
mask = 1 / (1 + np.exp(-dec.get_tensor(dec.get_output_details()[0]["index"])[0]))   # sigmoid, [352,352]

Kotlin (Android)

// vision + text: Accelerator.GPU; decoder: Accelerator.CPU
val vis = CompiledModel.create(File(dir,"clipseg_vision_fp16.tflite").path, CompiledModel.Options(Accelerator.GPU), null)
val dec = CompiledModel.create(File(dir,"clipseg_decoder.tflite").path, CompiledModel.Options(Accelerator.CPU), null)
// vision: image[1,3,352,352] -> t3,t6,t9   decoder: (t3,t6,t9,cond[512]) -> logits[1,352,352]
// text graph + host BPE/emb-lookup/text_projection produce cond; see ClipSeg.kt in the LiteRT sample.
val logits = decOut[0].readFloat()   // sigmoid -> mask

Conversion

Re-authored with litert-torch: qkv-3D-BMM attention, quick-GELU, baked interpolated pos-embed (14²→22² @352), host-side token-embedding lookup, safe_ln_up (up-scaled LayerNorm keeping the eps fp16-normal), convT4x4 (exact non-overlapping ConvTranspose as 1×1-conv + 4-D interleave). The decoder ships on CPU because its small-head-dim attention fp16-miscomputes on the Mali GPU delegate (re-authoring is exact — desktop fp16 corr 0.999996).

Upstream

CIDAS/clipseg-rd64-refined (Apache-2.0). Please cite Lüddecke & Ecker, Image Segmentation Using Text and Image Prompts (CVPR 2022).