Overpainting

Localized Context-aware Diffusion Image Editing

arXivComing soonCodeComing soonModelComing soonTest setAll results
A source image followed by five localized overpainting edits, each with a trimap inset.
A sequence of localized overpainting edits. Gray represents areas where change may occur and white represents areas where change must occur.

Abstract

We present “overpainting”, an image editing operation which offers both control over the location of the edit and awareness of the previous content in that location. The overpainted area is given by a trimap, where white-annotated pixels must be edited, gray-annotated pixels may be edited, and black-annotated pixels must not be edited. This enables both precise and loose control, depending on user intent.

We implement overpainting by adapting a pretrained image editing diffusion model using a combination of joint attention and low-rank adaption across input images with attention-dropout to balance the information flow between noise, source and mask images. We present a novel, automated, training data generation pipeline that (1) generates a set of candidate image pairs leveraging existing language-based editing models, (2) carefully curates those pairs, and (3) extracts a trimap from each usable pair. We demonstrate the versatility of our overpainting model on a wide range of editing tasks.

Trimap control

Specifying an exact editing mask around existing image features is cumbersome and error-prone (e.g., for thin features or precise object boundaries). It is also often ambiguous for edits with ill-defined boundaries (e.g., frost forming on a window or smoke dissipating) and for to-be-created features whose location depends on the target edit itself (e.g., whiskers protruding from an object on which the user asked the model to grow hair). We tackle this challenge by introducing a trimap that marks not only pixels which must be edited and must not be edited , but also pixels that may or may not be edited .

The same source edited with several precise and approximate trimaps.
Trimap variations on an image of an apple with the prompt “wrapped in black twine”. The user can place individual twines by drawing thin white lines, control some twine placement but also leave the model free to add/remove extra, or mark the entire apple white and force the model to cover it entirely.

Data Generation

To train our overpainting model we require paired data in the form of an original image, an edited version, a mask, and a prompt describing the edit operation. We follow the well-established process of leveraging VLMs and image diffusion models to generate training data. Specifically, we:

  1. randomly sample freely licensed photographs from pexels.com as source photographs
  2. use Qwen2.5-VL to suggest edit prompts
  3. feed the source images and generated propts into FLUX.1 Kontext to produce edited photographs
  4. filter out over-edited, under-edited, and misaligned results
  5. estimate which pixels were actually edited using an adapted BiRefNet finetune
  6. degrade the precise binary masks into human-like trimaps
Overview of the training data generation pipeline.
Overview of the training data generation pipeline.

Finetuning

We demonstrate overpainting by finetuning a Teamwork LoRA for the FLUX.1 Kontext model. Base FLUX Kontext accepts a prompt, latents for any number of condition images, and noisy latents for the image to generate. We provide both the source image and trimap image as VAE-encoded conditions.

Empirically we found that a naive LoRA-adapted model adhears poorly to the user-provided trimap. By using a Teamwork LoRA, we impose a strong positional bias (i.e., edited image latents at a particular position exchange information directly with the correspoinding source latents and mask latents). On the other hand, Teamwork alone allows limited global reasoning, and so provides poor edit quality. We therefore use both the attention of the base FLUX model and the aggredation of the Teamwork LoRAs, and balance between them with attention dropout during training.

Architecture overview.
Architecture overview. We pass the source image and trimap as additional condition images, and communicate across the conditions and the noise using both teamwork and joint attention.

Examples

A range of localized edits from the paper. Browse all 180 comparisons →

Selected overpainting results showing sources, prompts, trimaps, and edited outputs.
Selected overpainting results showing the source, prompt, trimap, and edited output.

Resources

Model outputs on our test set are avalible now. Source code and checkpoints will be made available once published.

Raw test setRaw result imagesResult viewer