Diffusion Inpainting for Object Editing

A Quick Recap: From SDEdit to Inpainting

Remember the masked editing algorithm from SDEdit? The core idea was to apply the reverse diffusion process only to the masked regions while preserving the unmasked areas. Inpainting builds directly on this principle but refines it for more seamless and high-fidelity results. We'll explore two main families of inpainting techniques:

1. Training-Free Inpainting: An iterative method that works with any standard diffusion model.
2. Fine-Tuned Inpainting: A more advanced approach that uses models specifically trained for the inpainting task.

Part 1: Training-Free Inpainting (The "RePaint" Method)

The simplest way to perform inpainting is to extend the logic from SDEdit. At each step of the reverse diffusion process, we can force the known, unmasked regions of the image to stay consistent with the original image, while letting the model generate new content for the masked region.

Here's the process:

1. Take the original image and a mask.
2. Start the reverse diffusion process from pure noise.
3. At each denoising step $t\to t-1$:
   a. Predict the denoised image ${\hat{x}}_{t-1}$ using the U-Net.
   b. Noise the original image up to step $t-1$ to get a "ground truth" noisy version, $x_{t-1}^{\text{original}}$.
   c. Combine them: use the generated pixels from ${\hat{x}}_{t-1}$ inside the mask, and the ground truth pixels from $x_{t-1}^{\text{original}}$ outside the mask.
4. Repeat until you reach step 0.

This iterative "resampling" ensures the generated content remains consistent with the unmasked surroundings.

This technique is clever because it requires no changes to the model architecture. However, as the model was not explicitly trained to harmonize the boundary between the generated and original pixels, this method can sometimes produce visible seams or artifacts at the mask's edge.

Part 2: Fine-Tuned Inpainting Models

To overcome the limitations of the training-free approach, specialized inpainting models are trained. The key innovation is in the model's architecture. Instead of just receiving the noisy latent image $z_t$ and the timestep $t$, the U-Net is given additional information.

Specifically, the input to the U-Net is modified to be a concatenation of:

1. The noisy latent image $z_t$.
2. The VAE-encoded masked original image.
3. The mask itself (downsampled to the latent resolution).

By providing the model with the context of the surrounding pixels (via the masked image) and the exact location of the hole (via the mask), it learns to generate content that seamlessly blends with the existing image.

This fine-tuning approach is what powers most high-quality inpainting tools and libraries, including the diffusers pipelines we'll use next.

Part 3: Practical Application with HuggingFace Diffusers

Now, let's translate this theory into practice. We'll use the diffusers library to perform both object replacement and removal.

Step 1: Setting up the Pipeline

First, you need an initial image, a mask specifying the area to edit, and an inpainting pipeline. The AutoPipelineForInpainting class in diffusers makes it easy to load a pre-trained, fine-tuned inpainting model.

import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid

# Use a GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
generator = torch.Generator(device).manual_seed(42) # for reproducibility

# Load initial image and mask
image_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"

init_image = load_image(image_url).resize((512, 512))
mask_image = load_image(mask_url).resize((512, 512))

# Load a fine-tuned inpainting model
pipeline = AutoPipelineForInpainting.from_pretrained(
    "runwayml/stable-diffusion-inpainting", 
    torch_dtype=torch.float16, 
    variant="fp16"
).to(device)

Step 2: Object Replacement

To replace an object, you provide a text prompt describing the new content you want to appear in the masked area. The strength parameter is crucial here. It controls how much noise is added to the masked region of the original image, similar to the $t_0$ hyperparameter from SDEdit.

• strength = 1.0: The masked area is completely replaced. This is ideal for object replacement.
• strength < 1.0: The model will try to blend the new content with the original content in the masked area.

# Prompt for the new object
prompt = "a cat sitting on a park bench"

# Generate the image
image = pipeline(
    prompt=prompt, 
    image=init_image, 
    mask_image=mask_image,
    strength=1.0, # Replace the content entirely
    guidance_scale=7.5,
    num_inference_steps=50,
    generator=generator
).images[0]

# Display the original, mask, and result
make_image_grid([init_image, mask_image, image], rows=1, cols=3)

Step 3: Object Removal

Object removal is a special case of inpainting where you want to fill the masked region with plausible background content. The most effective way to do this is to simply describe the background in your prompt.

Let's use a practical example from a UI-based tool to illustrate the different modes.

In diffusers, to remove an object, you would keep strength=1.0 and change the prompt to describe the background. For example, to remove the man from the bench in our previous example:

prompt = "a park bench" # Describe the background
image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, strength=1.0).images[0]

Step 4: Refining the Output

Achieving perfect results often requires some fine-tuning. Here are two common techniques.

1. Mask Blurring:
A sharp, binary mask can sometimes create a hard edge in the final output. Blurring the mask slightly can help the model create a smoother transition.

2. Preserving Unmasked Areas:
During the diffusion process, even unmasked areas can undergo slight changes. To ensure they remain identical to the original, you can do a final post-processing step where you paste the unmasked part of the original image back onto the generated output.

Here's how you might apply the overlay in code:

# Assuming 'inpainted_image' is the output from the pipeline
final_image = pipeline.image_processor.apply_overlay(
    mask_image=mask_image,
    original_image=init_image,
    inpainted_image=image
)

This ensures pixel-perfect preservation of the areas you didn't want to change.

Conclusion

In this lesson, you've learned how to apply diffusion models for inpainting, a powerful technique for object removal and replacement.

Key Takeaways:

• Inpainting can be done with training-free methods (like RePaint) or with more robust fine-tuned models.
• Fine-tuned models achieve better results by modifying the U-Net architecture to accept the masked image and the mask itself as conditioning.
• In practice, you can use the diffusers AutoPipelineForInpainting to easily work with these models.
• The strength parameter is key to controlling whether you replace (strength=1.0) or modify (strength<1.0) the content within the mask.
• Object removal is achieved by masking the object and prompting for the background.
• Techniques like mask blurring and overlaying the original image can help refine the final output and eliminate artifacts.

Next Lesson Preview:

Now that you can fill in holes inside an image, what about extending an image outwards? In the next lesson, we will explore outpainting, a technique that uses the same underlying principles as inpainting to expand an image's canvas and generate new content beyond its original borders.