Amortized Guidance for Image Inpainting with Pretrained Diffusion Models

TL;DR

This paper introduces Amortized Inpainting with Diffusion (AID), a method that efficiently guides pretrained diffusion models for image inpainting without per-instance optimization, improving quality-speed trade-offs.

Contribution

AID is a novel approach that trains a small guidance module offline, enabling fast, high-quality inpainting across images with minimal additional training.

Findings

AID improves quality-speed trade-off over baselines on AFHQv2, FFHQ, and ImageNet.

AID adds less than 1% trainable overhead.

The method is effective across multiple mask types and diffusion pipelines.

Abstract

We study image inpainting with generative diffusion models. Existing methods typically either train dedicated task-specific models, or adapt a pretrained diffusion model separately for each masked image at deployment. We introduce a middle-ground model, termed Amortized Inpainting with Diffusion (AID), which keeps a pretrained diffusion backbone fixed, trains a small reusable guidance module offline, and then reuses it across masked images without per-instance optimization. We formulate it as a deterministic guidance problem with a supervised terminal objective. To make this problem learnable in high dimensions, we derive an auxiliary Gaussian formulation and prove that solving this randomized problem recovers the optimal deterministic guidance field. This bridge yields a principled continuous-time actor--critic algorithm for learning the guidance module in a fully data-driven manner. Empirically, on AFHQv2 and FFHQ under the pixel EDM pipeline and on ImageNet under the latent EDM2 pipeline, AID consistently improves the quality--speed trade-off over strong fixed-backbone and amortized inpainting baselines across multiple mask types, while adding less than one percent trainable overhead.