MI-GAN brings mobile-sized image inpainting to a research baseline

The problem is a model that fits on a phone

MI-GAN is the official implementation of a paper presented at ICCV 2023, titled A Simple Baseline for Image Inpainting on Mobile Devices. Inpainting means filling in the masked region of an image with plausible content, and the reason the project exists is not that inpainting is new but that the state of the art models were too heavy to run on the hardware people actually carry. The teaser makes the trade-off explicit, describing results on both complex scene images and face images, and claiming a low FID while being one order of magnitude smaller and faster than recent state of the art approaches. The bubble chart in the same image encodes both axes that matter, with the bubble size showing the relative number of parameters and the number inside showing the absolute count. So the goal was never the highest possible benchmark score. The goal was a network whose size and latency let it live inside a mobile application, with quality good enough that the compromise is invisible to a user.

Distillation and re-parameterisation do the shrinking

The repository topics name the mechanism: knowledge distillation, re-parameterization, mobile inpainting, object removal and image inpainting. The two techniques do different jobs. Re-parameterisation is what makes the network cheap at inference: a training-time structure that folds into a simpler, faster form once training is over, so the deployed graph has fewer operations than the one that was optimised. Knowledge distillation is what keeps the small network close to the behaviour of a large one, by training the compact model against the outputs of a heavier teacher rather than against ground truth alone. Neither idea is new on its own, and the value of the paper is in the combination and in the engineering that made the combination run on a phone. The repository layout reflects the research tooling around that model, with configs/, a dnnlib/ directory that is the usual StyleGAN-style model library, lib/ for model definitions, torch_utils/ for training utilities, scripts/ for the demo, evaluation and mask generation entry points, examples/ holding the sample images and masks, and main.py at the top level.

Setting up the environment, with dependencies from 2021

The environment setup is a short sequence, and the documentation is upfront that Python 3.8 or newer is needed.

sudo apt install python3-venv
python3 -m venv venv
source ./venv/bin/activate
pip install pip --upgrade
pip install -r requirements.txt

The requirements file is where the real constraint lives, and it is a constraint on your environment rather than on the model. It pins torch at 1.8.0+cu111 and torchvision at 0.9.0+cu111, both pulled from the PyTorch stable wheel index, which means a CUDA 11.1 era build. Alongside that sit numpy at 1.23.5, opencv-python at 4.5.1.48, scikit-image at 0.17.2, protobuf at 3.20.*, tensorboard at 2.13.0, and the evaluation stack with lpips, pytorch_fid, fvcore, imgaug and omegaconf. onnxruntime is pinned at 1.13.1, which is the export and inference path you would use for a mobile target. Practically, this pins you to a fairly narrow corner of the ecosystem. If your environment already has a newer torch, the install will not simply resolve, and the sensible path is an isolated environment, which is why the instructions create one first. The pretrained models come from a linked Google Drive folder and go into a ./models directory, with a separate LICENSE-WEIGHTS file in the repository root that carries the licence for the weights rather than the code.

Running the demo on the bundled examples

The demo script takes a model, an image directory, a mask directory and an output directory, and the README gives a complete invocation for the 512 Places2 model.

python -m scripts.demo \
    --model-name migan-512 \
    --model-path ./models/migan_512_places2.pt \
    --images-dir ./examples/places2_512_object/images \
    --masks-dir ./examples/places2_512_object/masks \
    --output-dir ./examples/places2_512_object/results/migan \
    --device cuda \
    --invert-mask

Two details in that command are easy to get wrong. The --device flag selects cuda, so a machine without a GPU needs it changed. And the --invert-mask flag exists because of a convention, since by default white in a mask denotes the known region, and the example masks in the repository for the free-form sets treat white the other way around, so the note says the option should not be used with the provided free-form masks. The README also gives the 256 FFHQ variant, which uses migan-256, the migan_256_ffhq.pt weights and the ffhq_256_freeform examples directory, and that one runs without the invert flag. The examples directory holds four sample sets, ffhq_256_freeform, places2_256_freeform, places2_512_freeform and places2_512_object, and the same command shape applies to Co-Mod-GAN if you download those weights and change --model-name and --model-path.

The inference helper is not fully convolutional yet

There is a limitation the project states in its own words, and it is the kind that matters if you plan to ship the model. Even though the method is fully convolutional, the provided lib/model_zoo/migan_inference.py implementation currently performs some fixed-resolution operations. The documentation then names the fix: to get fully convolutional support you need to make the filter_const and noise_const computations dynamic so they depend on the input size, and it points at the self.register_buffer calls in the SeparableConv2d class and in the Upsample2d class. In other words, the technique generalises to arbitrary resolution but the reference inference code does not yet, and fixing it is a code change rather than a configuration flag. That is a small task for someone comfortable with PyTorch internals, and a real obstacle for anyone who wants to run at several resolutions from the same checkpoint. It also explains why the example commands are tied to specific model names and resolutions, since the pretrained weights are distributed per resolution, with 256 and 512 variants for both Places2 and FFHQ.

Evaluation, mask generation and the FID and LPIPS numbers

The evaluation path is a separate script, scripts.evaluate_fid_lpips, and it follows the same argument shape as the demo, with --real-dir pointing at a validation image set. For Places2 at 256x256 you pass migan-256 with migan_256_places2.pt, at 512x512 you pass migan-512 with migan_512_places2.pt, and for FFHQ at 256x256 you copy the first 10,000 samples into a separate validation directory and pass migan-256 with migan_256_ffhq.pt. The script generates free-form masks on the fly, and if you want the authors' exact masks you download the pre-generated sets for the three configurations and add a --mask-dir argument pointing at the right one. You can also produce your own masks with scripts.generate_masks, which takes an images directory, an output directory and a mask resolution. The metric of interest is FID and LPIPS, with pytorch_fid and lpips in the requirements, and the paper's own claim is a low FID relative to heavier models. The point worth holding on to is that the evaluation code is as reproducible as the demo code, so the numbers in the paper can be regenerated rather than taken on trust.

An ONNX path and third-party integrations

The repository documents a route for developers who want the model inside an application rather than in a Python process, pointing to a section for easy-to-use MI-GAN ONNX pipeline creation and for a pre-generated ONNX file that can be integrated into an app. The onnxruntime pin at 1.13.1 in the requirements confirms that path is a first-class part of the project. The README also lists third-party implementations, and the convention is that anyone with an app or demo can add it by opening an issue. Among them is an in-browser inpainting tool, a WeChat mini-program for on-device inpainting, and IOPaint, a desktop inpainting tool that can run MI-GAN directly with a model flag and a local web interface on port 8080. The IOPaint route is the shortest way to see the model in use without touching the research code, and the fact that three independent integrations exist says something useful about the interest in an inpainting model light enough to run on a phone.

Editorial conclusion

Adopt MI-GAN if you are working on on-device image inpainting and want the reference implementation, the pretrained weights and an ONNX export path, because the repository gives you the pieces a research baseline is supposed to give you. Skip it if you need a supported inference service today, since the requirements pin torch 1.8.0 with CUDA 11.1 and the inference helper still does fixed-resolution work. Before you run anything, read LICENSE-WEIGHTS separately from the code licence, because the weights have their own terms, and check whether the masks you plan to use treat white as the known region, since the default means the opposite of what many people assume.

Frequently asked questions

Official sources

1. Issues
2. License: MIT
3. Picsart-AI-Research/MI-GAN on GitHub
4. README