Imagen Editor and EditBench: Advancing and Evaluating Text-Guided Image Inpainting

In plain words

Text-guided inpainting means you paint over a part of a picture, type what should go there, and the model fills in that part so it matches both your words and the rest of the image. This paper takes a large text-to-image model, Imagen, and retrains it for this job. The model is called Imagen Editor.

A simple way to train such a model is to hide random patches of photos and ask the model to fill them back in. The trouble is that many random patches can be filled in from the surrounding picture alone, so the model learns to ignore the text. The authors' fix is to hide whole objects, found with an off-the-shelf object detector, instead of random patches. Training then depends more on the words.

The authors also built EditBench, a test set of images with masks and prompts. Prompts describe a single object, several objects, or the whole picture. People then judged whether each object, colour, material, size, shape and count was actually present. In those human tests, Imagen Editor beat DALL-E 2 and Stable Diffusion 1.5 on matching the text. Counting and shape were the weakest areas for every model.

TL;DR

• Fine-tunes Imagen, a cascaded diffusion model, for text-guided inpainting, conditioned on the full-resolution image and mask §3 fig. 2.
• Changes training masks: a random mask is combined with boxes from an object detector, so whole objects get hidden §3 app. A.2.
• Introduces EditBench, 240 image-mask pairs with three prompt types, scored on each attribute and object by humans §4 §5.1.
• Author-reported: Imagen Editor is preferred for text-image alignment over Stable Diffusion 1.5 (78%) and DALL-E 2 (77%) in side-by-side tests §5.2 fig. 8.

Problem

Inpainting with a text prompt needs two things at once: the filled region must fit the rest of the image, and it must show what the prompt says §1. Training with random masks can teach a model to paint from the image context and ignore the prompt, most often when masks are small or cut through an object §3 fig. 3.

Evaluation was also weak. Earlier work often showed cherry-picked samples or ran side-by-side comparisons that asked one overall question. No standard benchmark separated the attributes and objects a prompt asks for §2 §4.

Method

Base model. Imagen Editor is Imagen, a cascaded diffusion model, fine-tuned for inpainting §1 §3. The paper does not describe the Imagen cascade in its active text. The fine-tuning data, learning rate, step count and batch size are not stated.

Conditioning on image and mask (§3). The 1024×1024 source image and the binary mask are both passed to every diffusion stage, the base model and the super-resolution models fig. 2.

1. Each stage receives the image and mask concatenated along the channel axis, with the diffusion latents §3.
2. Before concatenation, the conditioning is downsampled with a learned strided convolution rather than a fixed filter. The authors report that bicubic downsampling left artifacts along mask edges §3.
3. The new input channels start at zero. At initialisation the model therefore ignores the image and mask and behaves like Imagen §3.
4. The intro says three convolutional downsampling image encoders add this context to each stage §1 fig. 2. The paper does not give their kernel sizes or channel counts.

# Image and mask conditioning for one diffusion stage  [§3] [fig. 2]
# Input:  noisy_latent (B, C, H, W) at this stage's resolution
#         source_image (B, 3, 1024, 1024); mask (B, 1, 1024, 1024)
#         text_embedding (B, L, D); t is the diffusion timestep
# Output: predicted noise (B, C, H, W)

def conditioned_denoise(noisy_latent, source_image, mask, text_embedding, t):
    # 1. Stack image and mask at full 1024 px resolution  [§3]
    full_res = concat([source_image, mask], axis=channel)   # (B, 4, 1024, 1024)
    # 2. Learned strided conv to the stage resolution, not bicubic  [§3]
    #    kernel size, stride and channels: not stated
    cond = strided_conv(full_res)                           # (B, C_cond, H, W)
    # 3. Concatenate with the diffusion latent along channels  [§3]
    model_input = concat([noisy_latent, cond], axis=channel)
    # 4. Added input channels start at zero, so the model begins as Imagen  [§3]
    return unet(model_input, t, text_embedding)

Object masking (§3, app. A.2). Training masks are built on the fly.

1. Draw a random mask. The paper names random box and/or random stroke masks, but does not give their distribution §3.
2. Run an off-the-shelf detector, SSD MobileNet v2 from TensorFlow Hub, to get object boxes §3 (footnote).
3. Combine the random mask with the detected box(es) by union, so both regions are hidden app. A.2 fig. 14.
4. The number of boxes used, the score threshold and the object classes are not stated in the active text. A commented-out listing in the source takes only the top box [src: obj_detect_masking.tex, commented]. That listing is not in the compiled paper.

# Object-masking training mask for one image  [§3] [app. A.2] [fig. 14]
# Input:  image (1024, 1024, 3)
# Output: mask (1024, 1024); 0 marks the hidden region (see Pitfalls)

def make_training_mask(image):
    # 1. Random box and/or stroke mask, the usual inpainting policy  [§3]
    random_mask = sample_random_mask(image)
    # 2. Object detector proposes boxes  [§3]
    #    number of boxes and score threshold: not stated
    boxes = detect_objects(image)
    # 3. Start from a fully visible mask  [app. A.2]
    object_mask = ones_like(random_mask)
    for box in boxes:
        object_mask[box] = 0
    # 4. Union of the two hidden regions  [app. A.2]
    return random_mask * object_mask

Classifier-free guidance (§3). Guidance is used heavily. The authors follow Ho et al. (2022, Imagen Video) and use high weights with oscillation §3. In the base model the guidance weight oscillates between 1 and 30 §3. The shape of the oscillation and the weights of the other stages are not stated.

Two variants compared (§5). Imagen Editor (IM) uses object masking. Imagen Editor RM uses random masking. Everything else is described as the same §5. Stable Diffusion v1.5 and DALL-E 2 are evaluated as baselines §5.

Hyperparameters

Results

Author-reported, on EditBench, using human raters and CLIP-based metrics. Models: Imagen Editor (IM), Imagen Editor with random masking (IM RM), Stable Diffusion v1.5 (SD), DALL-E 2 (DL) §5.

Human single-image evaluation (§5.2, fig. 7). Imagen Editor has the highest text-image alignment rate on every prompt type, 10–13% above the second-best model §5.2. The rest are ordered IM RM > DL > SD, with 3–6% gaps, except on Mask-Simple, where IM RM is 4–8% behind §5.2. Strict scoring: on Mask-Simple and Mask-Rich, a rating counts as correct only when every attribute and object is present and bound correctly fig. 7 caption.

Human side-by-side (§5.2, fig. 8). On Mask-Rich prompts, Imagen Editor is preferred for text-image alignment over SD in 78%, over DL in 77%, and over IM RM in 68% of comparisons §5.2. Image-quality preference differs by 0–6% between models §5.2.

By object type, Mask-Simple prompts (§5.2, fig. 9). Imagen Editor leads on common, rare and text-rendering objects, by 10%, 11% and 11% over the second-best model §5.2. SD scores 59% on common, 44% on rare and 26% on text-rendering objects §5.2.

By attribute type, Mask-Simple prompts (§5.2, fig. 10). Imagen Editor leads by 13–16% on the attribute types, except count, where DL is 1% behind §5.2. Imagen Editor gains least over IM RM on count (14%; the next smallest gain is 22%) §5.2.

Automatic metrics, aggregated (§5.3, tab. 3). Imagen Editor ranks first on every automatic metric among the four models. The reference images score highest on CLIP-R-Prec and NIMA.

Metric agreement with humans, best image of two (§5.3, tab. 1). Agreement with human picks, in percent. Chance is 50.0.

Metric agreement with humans, best of four models (§5.3, tab. 2). Chance is 25.0.

Paper vs code

No official code.

The arXiv metadata and abstract carry no code link. The paper's footnote points to a project page, imagen.research.google/editor. That page links an EditBench dataset archive and no code for Imagen Editor or EditBench. GitHub searches for "imagen editor editbench", "editbench" and "imagen" inside the google-research organisation found no repository from the authors. Matches for "editbench" belong to other owners and were not used.

Imagen Editor has no released code or weights, so the architecture and training recipe above come only from the paper.

Pitfalls

• Training cannot be reproduced from this paper. The learning rate, optimizer, step count, batch size, fine-tuning data and mask-size distribution are not stated. Treat any reimplementation as your own choice, and say so.
• The base model is not described in the active text. The Imagen cascade (stages, sizes, training) is only in the cited Imagen work. A commented-out passage in the source describes it, but that passage is not part of the compiled paper.
• Mask polarity is not stated in the active text. The commented-out listing treats 0 as hidden. Under that convention, multiplying the random mask by the box mask gives the same union the text describes. Check the convention in your own code and be explicit about it.
• Detector details are unclear. The active text says "boxes" (plural) and "object detector". The number of boxes, classes and score threshold are not stated. The commented-out listing uses the top prediction only.
• Object masking helps less for large masks. The authors report that the gain is largest for small and medium masks app. A.1 fig. 13. The abstract's "across-the-board" claim refers to text-image alignment in the main results §5.2.
• Strict prompt scoring. Overall rates for Mask-Simple and Mask-Rich count a rating as correct only if every attribute and object is right and bound correctly fig. 7 caption. Rates on different prompt types are not directly comparable, as the caption also says.
• Count and shape remain weak for all models. Count was the least improved attribute by object masking §5.2. Appendix failure cases describe missing objects, often more than half app. A.3 fig. 23.
• CLIP metrics are a weak proxy. The best image-level agreement with humans is 68–76% tab. 1, so automatic scores do not replace the human ratings. On the aggregate table, the reference images score higher than every model on CLIP-R-Prec and NIMA, but not on CLIPScore T2I, where IM is 31.5 and the reference is 31.0 tab. 3.
• DALL-E 2 is a moving target. It was accessed through a web UI in October 2022. The underlying model is unpublished and may differ from the cited paper §5 footnote. SD was version 1.5 §5.
• Small rater pool. The human studies used 18 US-based annotators, with about 30 seconds per prompt app. A.1. Results depend on this pool and on single-image binary judgements.
• EditBench data licence. The project page links an archive of about 379 MB. The paper does not state a dataset licence, and this card does not use any EditBench image or prompt text. Check the dataset's terms before reusing it.
• Societal impact. The authors say they add watermarks to generated images and avoid generating photorealistic human faces §6. The paper gives no detail on how either is done, and does not evaluate them.
• Figure and table numbers follow the order of captions in the source. The appendix's own references (Figs. 7, 9 and 10) agree with that order.

Lineage

• Builds on: Imagen, "Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding" (Saharia et al., NeurIPS 2022; cascaded text-to-image diffusion, base of this fine-tune). The source gives no arXiv id, and no card exists yet §1 §3. Channel-concatenation conditioning follows SR3, Palette and GLIDE §3. Random-mask training policy follows earlier inpainting work app. A.2.
• Compared against: DALL-E 2 §5; Stable Diffusion v1.5 (Rombach et al., CVPR 2022) §5. Related masked and mask-free editing work, including Blended Diffusion, DiffEdit, Paint by Word, Text2LIVE, Prompt-to-Prompt and Imagic, is discussed in §2.
• Superseded by: not known from the content of this wiki.