MI-GAN β€” LiteRT (on-device image inpainting / object removal, fully-GPU)

MI-GAN (Picsart AI Research, ICCV 2023) β€” a mobile "magic eraser": paint over an object and it is removed and inpainted. Converted to LiteRT and running fully on the CompiledModel GPU (ML Drift) on Android (512Γ—512, Places2).

MI-GAN β€” original / mask / inpainted (on-device LiteRT GPU)

On-device (Pixel 8a, Tensor G3 β€” verified)

nodes on GPU 509 / 509 LITERT_CL (full residency)
inference ~6 ms (512Γ—512)
size 16.3 MB (fp16)
accuracy device-vs-PyTorch corr 0.99998, no NaN
in[1,4,512,512] = concat(mask-0.5, rgbΒ·mask)  β†’[GPU: MI-GAN]β†’  out[1,3,512,512] (inpainted, [-1,1])

How it converts (litert-torch) β€” clean in one shot, no re-authoring

The MI-GAN inference generator (the re-parametrized mobile model) is already GPU-friendly: depthwise-separable Conv2d, nn.Upsample(nearest) + a fixed FIR-filter grouped conv (no transposed conv), leaky-ReLU with gain/clamp (β†’ MAXIMUM/MINIMUM), and no normalization layers (StyleGAN-style). Banned ops NONE, all tensors ≀4D, tflite-vs-torch corr 1.0, device-vs-torch corr 0.99998.

I/O

  • Input (4 ch): concat(mask βˆ’ 0.5, rgb Β· mask) β€” rgb ∈ [βˆ’1,1] (pixel/127.5 βˆ’ 1); mask = 1 keep, 0 erase.
  • Output (3 ch): generated image in [βˆ’1,1]; composite as rgbΒ·mask + outΒ·(1βˆ’mask).

Preprocessing: center-crop, resize 512Γ—512.

Minimal usage

Android (Kotlin, CompiledModel GPU)

val model = CompiledModel.create(context.assets, "migan_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()
inputs[0].writeFloat(x)            // [1,4,512,512] = concat(mask-0.5, rgb*mask)
model.run(inputs, outputs)
val out = outputs[0].readFloat()    // [1,3,512,512] in [-1,1]; composite rgb*mask + out*(1-mask)

Python (desktop verification)

import numpy as np
from PIL import Image
from ai_edge_litert.interpreter import Interpreter

rgb = (np.asarray(Image.open("photo.jpg").convert("RGB").resize((512, 512)), np.float32)
       / 127.5 - 1).transpose(2, 0, 1)                            # [3,512,512], [-1,1]
m = np.asarray(Image.open("mask.png").convert("L").resize((512, 512)), np.float32)
mask = (m < 128).astype(np.float32)[None]                          # 1 = keep, 0 = erase (painted)
x = np.concatenate([mask - 0.5, rgb * mask])[None]                 # [1,4,512,512]

it = Interpreter(model_path="migan_fp16.tflite"); it.allocate_tensors()
it.set_tensor(it.get_input_details()[0]["index"], x); it.invoke()
out = it.get_tensor(it.get_output_details()[0]["index"])[0]        # [3,512,512], [-1,1]
comp = rgb * mask + out * (1 - mask)
Image.fromarray(((comp.transpose(1, 2, 0) + 1) * 127.5).clip(0, 255).astype(np.uint8)).save("inpainted.png")

Performance

Measured on a Pixel 8a (Tensor G3, Android 16) with the standard TFLite benchmark_model tool β€” 10 warm-up runs then 50 timed runs, reported as the tool's mean.

Runtime Backend Graph on GPU Latency
LiteRT CompiledModel (LITERT_CL) GPU 509 / 509 ~6 ms
TFLite benchmark_model (TfLiteGpuDelegateV2) GPU (OpenCL) 509 / 509 68.1 ms
TFLite benchmark_model CPU (XNNPACK, 4 threads) β€” XNNPACK declined the graph

The two GPU rows are different runtimes, not a contradiction. The LITERT_CL figure is the one recorded when this model shipped, taken through LiteRT's own CompiledModel accelerator β€” the path the Kotlin sample app and the LiteRT API use. The TfLiteGpuDelegateV2 figure is the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. They agree on how much of the graph the GPU takes; they disagree on speed, and the classic delegate is the slower of the two here. Read the TfLiteGpuDelegateV2 row as a reproducible floor, not as this model's speed on LiteRT.

XNNPACK declines these fp16 graphs β€” it reports failed to delegate DEPTHWISE_CONV_2D and then fails to allocate tensors β€” so there is no usable CPU number. Disabling XNNPACK falls back to reference kernels, which measured about 20Γ— slower than the GPU on models of this size and would not represent CPU inference anyone would ship.

Snapdragon NPU (Hexagon)

The GPU is faster: 26.64 ms against 30.85 ms on the NPU, a factor of 1.16. The NPU still loads 7.34x faster (168 ms against 1231 ms).

backend compiled inference (median / min) load
NPU (Hexagon v81) on-device JIT 30.85 ms / 30.17 ms 168 ms
GPU (Adreno) β€” 26.64 ms / 24.54 ms 1231 ms

Measured on a Samsung Galaxy S26 (Snapdragon 8 Elite Gen 5 / SM8850, Hexagon v81, Android 16) with LiteRT CompiledModel 2.2.0, one accelerator per process, 5 warm-up runs then N=50 timed runs, median reported. Every run held thermal status NONE throughout. Headroom 0.80, where 1.0 is the throttling threshold.

The NPU rows ran the published file unchanged. LiteRT compiled it for the Hexagon on the device at first load. That first compile took 35 s here. The load column above is the cached load every later run pays. Recipe and the runtime libraries it needs: NPU guide.

GPU wiring: GPU guide.

License

MIT. Upstream: Picsart-AI-Research/MI-GAN.

Downloads last month
98
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Collection including litert-community/MI-GAN-512-Places2-LiteRT