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

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).

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")

License

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