SmartBrush: Text and Shape Guided Object Inpainting with Diffusion Model

In plain words

Image inpainting fills a hole in a picture so that it blends in. Earlier text-guided tools can fill a hole from a sentence, but they often paint background into the hole instead of the object you named, and they ignore the shape you drew. This paper's model, SmartBrush, takes a short label such as "cat" plus a mask and fills the mask with that object.

The mask can be exact, following the outline of the object, or rough, like a box. The model is trained by blurring real object masks to different degrees, so it learns how closely to follow a given outline. It also learns to guess where the object sits inside a rough mask, so that the pixels around the object inside the box are left alone. Training uses real object masks from segmentation datasets, which tell the model exactly which region holds which object.

TL;DR

• Text- and mask-guided object inpainting with a diffusion model; the mask precision is an explicit input §1 §4.2.
• Training uses instance masks from OpenImages v6 with label or BLIP captions, and an extra mask-prediction output that keeps background inside coarse masks §4.1 §4.3 §4.4.
• Author-reported: best local FID and FID in Tables 1 and 2 on OpenImages and MSCOCO; CLIP scores are close to the baselines tab. 1 tab. 2.
• No official code was found (see Paper vs code).

Problem

Text-guided inpainting baselines show three failures, which the authors name text misalignment, mask misalignment and background preservation §1. Text misalignment: the model fills the hole with background instead of the named object §1. Mask misalignment: a global CLIP-guided method does not place the object to the shape of the mask §1 §2. Background preservation: inside a coarse mask, the model changes pixels that are not part of the object §1 §5.4.

The paper assumes a mask is available and that the object class is known. The prompt is a local label or caption for the masked instance, not a description of the whole image §1 §4.4.

Method

Overview. The model is a conditional diffusion model. Its inputs are a noisy image, a timestep, a blurred mask with a precision level, and a text prompt. Its outputs are the predicted noise and an extra channel that predicts the accurate object mask §4.1 §4.3. The paper fine-tunes a Stable Diffusion text-to-image model v1.2 with these conditions §4.4.

Step 1: text and shape guided diffusion §4.1. Training samples come from segmentation datasets. For each instance mask m, the class label is its text prompt §4.1. Noise is added only inside the mask. The background stays clean, so the model can use it while it generates the object §4.1 eq. 5. The loss is the usual noise-prediction loss, computed with the mask and the text as extra inputs §4.1 eq. 6.

Step 2: shape precision control §4.2. For each training object, the accurate mask m is blurred with a Gaussian kernel of size k_s and standard deviation σ_s, for a precision level s in [0, S] eq. 7. At s = 0 the mask is unchanged. At s = S the mask is the bounding box of the instance §4.2. The model is conditioned on the blurred mask, its level s, the text, and the noisy image §4.2 eq. 8. At test time the user chooses s, which sets how closely the object follows the mask §4.2 fig. 5.

Step 3: background preservation §4.3. An extra output channel of the same network predicts the accurate instance mask from the blurred input §4.3. It is trained with the DICE loss on the predicted mask §4.3 eq. 9. During sampling, the coarse mask is first used, then the model's predicted mask replaces it, so that only the predicted object region is denoised §4.3 fig. 7.

Step 4: training strategy §4.4. The final objective is the noise loss plus λ times the mask loss eq. 10. The text is a segmentation label or a BLIP caption, paired at random with the mask §4.4. The model is also trained for text-to-image generation, with the mask covering the whole image, on LAION-Aesthetics v2 5+ image-text pairs §4.4.

Training and evaluation split §5.1. Inpainting and text-to-image tasks are sampled with probabilities 80% and 20% §5.1.

# Training step  [§4.1] [§4.2] [§4.3] [§4.4]
# Input:  image x0 (3, H, W); accurate instance mask m (H, W);
#         label or caption c for that instance
# Output: scalar loss

def training_step(x0, m, c):
    # 1. Sample a precision level and blur the instance mask  [§4.2] [eq. 7]
    # s in [0, S]; sampler not stated
    s = sample_precision_level()
    # kernel, sigma not stated
    m_s = gaussian_blur(m, kernel[s], sigma[s])
    # 2. Add noise only inside the mask; keep background clean  [§4.1] [eq. 5]
    t = sample_timestep()
    eps = sample_gaussian_noise(x0)
    x_t_tilde = sqrt(alpha_bar[t]) * x0 + sqrt(1 - alpha_bar[t]) * eps
    x_t = x_t_tilde * m + x0 * (1 - m)
    # 3. Predict the noise and an extra mask channel  [§4.2] [§4.3] [eq. 8]
    eps_pred, mask_pred = model(x_t, t, m_s, c, s)
    loss_seg = squared_l2(eps - eps_pred)
    # 4. Supervise the extra channel with the accurate mask  [§4.3] [eq. 9]
    loss_pred = dice_loss(mask_pred, m)
    # 5. Weighted sum with lambda = 0.01  [§4.4] [eq. 10] [§5.1]
    return loss_seg + 0.01 * loss_pred

# Inference  [§4.1] [§4.3] [fig. 7]
# Input:  image x0; coarse input mask m_s; prompt c; precision level s
# Output: inpainted image

def inpaint(x0, m_s, c, s):
    # 1. Noise inside the coarse mask, clean background outside  [§4.1]
    noise = sample_gaussian_noise(x0)
    x_T = noise * m_s + x0 * (1 - m_s)
    # 2. Denoise; use the coarse mask at first  [§4.3]
    mask = m_s
    for t in sampling_schedule:
        eps_pred, mask_pred = model(x_T, t, mask, c, s)
        # 3. Switch to the predicted mask; the switch step is not stated  [§4.3]
        if switch_to_predicted(t):
            mask = mask_pred
        x_T = sampler_step(x_T, eps_pred, t)
    return x_T

Key equations (in the paper's notation) eq. 5 eq. 7 eq. 8 eq. 9:

x_t = \tilde{x}_t \odot m + x_0 \odot (1-m) \quad [\text{eq. 5}]

m_s = \text{GaussianBlur}(m, k_s, \sigma_s) \quad [\text{eq. 7}]

\mathcal{L}_{\text{seg-DM}} = \mathbb{E}\left[\|\epsilon - \epsilon_\theta(x_t, t, m_s, c, s)\|_2^2\right] \quad [\text{eq. 8}]

\mathcal{L}_{\text{prediction}} = H(\epsilon_\theta(m_s), m), \quad H(X,Y) = 1 - \frac{2|X \cap Y|}{|X|+|Y|} \quad [\text{eq. 9}]

\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{seg-DM}} + \lambda \mathcal{L}_{\text{prediction}} \quad [\text{eq. 10}]

Hyperparameters

Results

Author-reported. Setup: OpenImages and MSCOCO, text prompt is the class label, metrics Local FID (lower is better), CLIP score (higher is better) and FID (lower is better) §5.1. Baselines: Blended Diffusion, GLIDE, Stable Diffusion, Stable Inpainting §5.1. DALL-E 2 is compared only on a few images, since no public model exists §5.1 fig. 6.

Bounding box mask tab. 1:

Object layout mask tab. 2:

Bold marks the best value in the source tables tab. 1 tab. 2. On OpenImages CLIP score, Stable Diffusion ties SmartBrush at 0.263 in tab. 2.

User study §5.2 fig. 4. Over 300 Amazon Mechanical Turk workers judged which output follows the mask, which follows the text, and which looks most realistic. The paper states that more than 50% of votes went to SmartBrush on each question. Per-method vote counts are not in the text.

Qualitative results. Comparison with five methods on object and scene prompts fig. 3. Mask precision levels 0 to 4 for the astronaut example fig. 5. Background preservation against DALL-E 2 and an ablation without the mask prediction fig. 6. Predicted masks at sampling steps 0 to 50 fig. 7.

Paper vs code

No official code.

Searched, in order: the paper's text, footnotes and bibliography have no code link (only the Stable Diffusion repository and the OpenImages page). The arXiv metadata has no comment or project page field. GitHub repository search for "smartbrush" returned 12 repositories, none by the authors. Search for "smartbrush inpainting" returned 0. Search within the adobe-research organisation returned 0. A GitHub user search for "shaoan xie" returned one account that could not be linked to the author, so it was not used. The MIT-licensed repository IjustAi/smartbrush and the other search hits were not cloned, since none is shown to belong to the authors.

Pitfalls

• Unstated core values. The precision range S, the blur kernels k_s and σ_s, how s is fed to the network, and the mask-switch step are all absent §4.2 §4.3. An implementer has to pick them and will not reproduce the paper's mask-control behaviour exactly.
• Sampler and optimizer. Neither the sampler, number of steps, guidance scale, nor the optimizer and learning rate are stated §5.1. Section 3 names DDIM and PNDM only as generic options §3.
• Which mask is used in the loss. The noise loss uses the blurred mask m_s as an input, while the mask loss targets the accurate mask m eq. 8 eq. 9. Eq. 9 writes the mask head as H(ε_θ(m_s), m); the paper does not say at which timesteps the mask loss is computed or how the extra channel is supervised per step.
• Mask switch at inference. The paper says to switch from the coarse mask to the predicted mask, but gives no step and no binarisation threshold §4.3. Fig. 7 says the prediction sharpens after about 10 steps fig. 7.
• Noise only in the mask. Eq. 5 keeps the background clean during training §4.1. The inference start point uses the same mask §4.1. Whether the loss is averaged over the whole image or over the mask is not stated.
• Text labels versus captions. Labels are random-paired with BLIP captions, with no ratio given §4.4. The paper says short labels degrade quality compared with long sentences §4.4; no numbers are given for that.
• Task mix. The 80/20 split is stated §5.1, but not whether it is per sample or per batch.
• Background claim is qualitative. The background-preservation benefit is shown in figures only fig. 6 §5.4. No background metric is reported. The DALL-E 2 comparison uses a few images §5.1.
• Unequal resolution for baselines. Blended Diffusion and GLIDE are run at 256 px and resized, while SmartBrush runs at 512 px §5.1. The gap may be partly resolution-related.
• Known failure. Long shadows that extend beyond the object mask are not covered, because the coarsest mask is the bounding box §6.
• No variance reported. Tables 1 and 2 give single numbers with no seeds or confidence intervals tab. 1 tab. 2.

Lineage

• Builds on: diffusion models (Sohl-Dickstein 2015; Ho et al. 2020, cited in §3) and the Stable Diffusion text-to-image model v1.2 used as the base §4.4. Those works are cited in the paper and have no card here.
• Compared against: Blended Diffusion, GLIDE, Stable Diffusion and Stable Inpainting tab. 1 tab. 2; DALL-E 2 on a few examples §5.1 fig. 6. DALL-E 2 has a card: /papers/2204.06125.
• Uses as data sources: BLIP captions §4.4 and LAION-Aesthetics v2 §4.4.
• Superseded by: not known from content in this wiki.