How to Implement Image Inpainting with Diffusion Models in Ailia-Models

You can implement image inpainting with diffusion models by using the latent diffusion inpainting pipeline in the ailia-models repository, which combines a conditional encoder, DDIM sampling in latent space, and an autoencoder decoder to fill masked regions while preserving context.

The ailia-models repository by axinc-ai provides a production-ready implementation of latent diffusion inpainting that runs entirely on ONNX models. This guide explains how to implement image inpainting with diffusion models using the provided Python scripts and how to integrate the pipeline into your own applications.

Understanding the Latent Diffusion Inpainting Architecture

The implementation in diffusion/latent-diffusion-inpainting/latent-diffusion-inpainting.py follows a three-stage pipeline that operates in the latent space of an autoencoder, significantly reducing computational requirements compared to pixel-space diffusion.

The Three-Stage Pipeline

1. Conditional Encoding: The masked image is processed by the cond_stage_model (a conditional encoder) to produce a latent conditioning vector. This encoder expects pixel values in the range [-1, 1] and outputs the initial conditioning state used by the diffusion process.

2. Latent Diffusion: The core generation happens in the diffusion_model, a UNet-style denoising network that predicts noise εₜ at each timestep. The script uses DDIM (Denoising Diffusion Implicit Models) sampling via make_ddim_timesteps() and make_ddim_sampling_parameters() to traverse the diffusion trajectory deterministically or stochastically based on the ddim_eta parameter.

3. Decoding: The final latent sample is passed through the autoencoder (specifically decode_first_stage()) to reconstruct the RGB image. The output is then blended with the original unmasked regions to preserve the context outside the inpainting mask.

DDIM Sampling Strategy

The DDIM implementation relies on pre-computed alpha values stored in diffusion/latent-diffusion-inpainting/constants.py. The alphas_cumprod array drives the variance schedule, while make_ddim_sampling_parameters() computes the scaling factors (ddim_alphas, ddim_alphas_prev, ddim_sigmas) required for the sampling equations from the original DDIM paper.

Setting Up the Environment and Models

Before running inference, you must obtain the three ONNX model files and initialize the runtime environment.

Downloading Required ONNX Files

The repository uses check_and_download_models() from util/model_utils.py to manage model artifacts. The required files are:

• cond_stage_model.onnx – Conditional encoder weights
• diffusion_model.onnx – UNet denoising network
• autoencoder.onnx – VAE decoder

These are automatically fetched from https://storage.googleapis.com/ailia-models/latent-diffusion-inpainting/ if not present locally.

Loading Models with Ailia

The main() function in latent-diffusion-inpainting.py demonstrates proper initialization:

import ailia

# Memory optimization flags

memory_mode = ailia.get_memory_mode(
    reduce_constant=True,
    ignore_input_with_initializer=True,
    reduce_interstage=False,
    reuse_interstage=True
)

# Initialize networks

cond_stage = ailia.Net(MODEL_COND_STAGE_PATH, WEIGHT_COND_STAGE_PATH, 
                      env_id=0, memory_mode=memory_mode)
diffusion = ailia.Net(MODEL_DFSN_PATH, WEIGHT_DFSN_PATH, 
                     env_id=0, memory_mode=memory_mode)
autoenc = ailia.Net(MODEL_AUTO_ENC_PATH, WEIGHT_AUTO_ENC_PATH, 
                   env_id=0, memory_mode=memory_mode)

models = {
    "cond_stage_model": cond_stage,
    "diffusion_model": diffusion,
    "autoencoder": autoenc,
}

For CPU-only or alternative runtime environments, you can substitute Ailia with ONNX Runtime by passing the --onnx flag.

Implementing the Inpainting Pipeline

The core logic resides in the predict() function, which orchestrates preprocessing, sampling, and post-processing.

Preprocessing Images and Masks

The preprocess() function (lines 16-35) handles normalization:

import cv2
import numpy as np

def preprocess(img, mask):
    # Scale to [0, 255] then normalize to [-1, 1]

    img = img.astype(np.float32) / 127.5 - 1.0
    mask = mask.astype(np.float32) / 255.0
    mask[mask < 0.5] = 0
    mask[mask >= 0.5] = 1
    mask = mask * 2 - 1  # Scale to [-1, 1]

    return img, mask

The masked image is computed as (1 - mask) * img and fed to the conditional encoder.

Running DDIM Sampling

The ddim_sampling() method implements the iterative denoising loop. At each step, p_sample_ddim() calls apply_model() (lines 207-218) to predict the noise εₜ:

def apply_model(x_noisy, t, cond):
    # x_noisy: latent with noise

    # t: timestep tensor

    # cond: conditioning from encoder

    eps = diffusion_model.run(None, {'x': x_noisy, 'timesteps': t, 'context': cond})[0]
    return eps

The predicted noise is used to compute pred_x0 (the estimated clean latent) and update the sample for the previous timestep using the DDIM formula with ddim_alphas and ddim_sigmas.

Post-processing and Blending

After sampling completes, decode_first_stage() (lines 221-232) converts the latent to RGB:

def decode_first_stage(z):
    # z: latent tensor [B, C, H, W]

    dec = autoencoder.run(None, {'input': z})[0]
    # Scale from [-1, 1] to [0, 255]

    dec = (dec + 1.0) / 2.0 * 255.0
    return dec.astype(np.uint8)

The final output blends the original unmasked regions with the generated content:

result = (1 - mask) * original_image + mask * generated_image

Command-Line and Python API Usage

You can run the pipeline via command line or integrate it programmatically.

Command-Line Interface

Execute the script with specific parameters to control quality and speed:

python diffusion/latent-diffusion-inpainting/latent-diffusion-inpainting.py \
    --input path/to/photo.png \
    --mask path/to/photo_mask.png \
    --ddim_steps 50 \
    --ddim_eta 0.0 \
    --savepath outputs/

Key Parameters:

• --input – Source image path (BGR format).
• --mask – Binary mask where white pixels (255) indicate regions to inpaint.
• --ddim_steps – Number of sampling steps (default 50; increase for quality).
• --ddim_eta – Stochasticity parameter (0.0 = deterministic, higher = more random).
• --savepath – Output directory for *_inpainted.png files.

The script automatically downloads ONNX weights from https://storage.googleapis.com/ailia-models/latent-diffusion-inpainting/ if missing.

Python API Integration

For programmatic use, import the functions directly:

import cv2
from diffusion.latent_diffusion_inpainting.latent_diffusion_inpainting import (
    preprocess, predict, check_and_download_models,
    WEIGHT_COND_STAGE_PATH, MODEL_COND_STAGE_PATH,
    WEIGHT_DFSN_PATH, MODEL_DFSN_PATH,
    WEIGHT_AUTO_ENC_PATH, MODEL_AUTO_ENC_PATH,
    REMOTE_PATH
)

# Download models if needed

for w, m in [(WEIGHT_COND_STAGE_PATH, MODEL_COND_STAGE_PATH),
             (WEIGHT_DFSN_PATH, MODEL_DFSN_PATH),
             (WEIGHT_AUTO_ENC_PATH, MODEL_AUTO_ENC_PATH)]:
    check_and_download_models(w, m, REMOTE_PATH)

# Load image and mask

img = cv2.imread("damaged_photo.png")
mask = cv2.imread("mask.png", cv2.IMREAD_GRAYSCALE)

# Run inference (models dict would be initialized as shown in previous sections)

result = predict(models, img, mask)
cv2.imwrite("restored.png", result)

This approach allows you to embed the diffusion inpainting pipeline into web services, batch processing scripts, or GUI applications.

Customising the Sampling Schedule

If you need a different number of DDIM steps, adjust the global variables before calling predict():

from diffusion.latent_diffusion_inpainting.latent_diffusion_inpainting import (
    args, make_ddim_timesteps, make_ddim_sampling_parameters, alphas_cumprod
)

args.ddim_steps = 100          # higher quality

args.ddim_eta = 0.1            # add stochasticity

ddim_num_steps = args.ddim_steps
ddpm_num_timesteps = 1000
ddim_timesteps = make_ddim_timesteps(ddim_num_steps, ddpm_num_timesteps)

ddim_eta = args.ddim_eta
ddim_sigmas, ddim_alphas, ddim_alphas_prev = \
    make_ddim_sampling_parameters(alphas_cumprod, ddim_timesteps, ddim_eta)

These variables are used internally by ddim_sampling(); redefining them gives you fine-grained control over speed vs. fidelity.

Summary

• The ailia-models repository provides a complete latent diffusion inpainting implementation using three ONNX networks: a conditional encoder, a UNet diffusion model, and an autoencoder decoder.
• The pipeline runs DDIM sampling in latent space via make_ddim_timesteps() and make_ddim_sampling_parameters(), offering deterministic or stochastic generation controlled by the ddim_eta parameter.
• Preprocessing in preprocess() normalizes images and masks to the [-1, 1] range required by the diffusion model, while post-processing blends generated content with original unmasked regions.
• You can execute the pipeline via command line (latent-diffusion-inpainting.py) or integrate it programmatically using the predict() function and Ailia/ONNX Runtime APIs.

Frequently Asked Questions

What is the difference between DDIM and DDPM sampling in this implementation?

DDIM (Denoising Diffusion Implicit Models) sampling is used instead of DDPM because it allows for deterministic generation when ddim_eta is set to 0.0, and requires fewer steps (typically 50) to achieve high-quality results compared to the original DDPM formulation. The implementation uses make_ddim_timesteps() to create a subsequence of timesteps and make_ddim_sampling_parameters() to compute the variance schedule from the alphas_cumprod values defined in constants.py.

How do I prepare the mask image for the inpainting pipeline?

The mask should be a binary image (single channel grayscale) where white pixels (255) indicate regions to be inpainted and black pixels (0) indicate regions to preserve. The preprocess() function in latent-diffusion-inpainting.py automatically binarizes the mask (threshold 0.5) and normalizes values to the range [-1, 1]. Ensure your mask dimensions match the input image exactly to avoid alignment errors during the blending stage.

Can I run this inpainting model on CPU or do I need a GPU?

You can run the model on CPU using ONNX Runtime by passing the --onnx flag when executing the script. By default, the code uses the Ailia SDK, which automatically selects the best available backend (CUDA, Vulkan, or CPU). The memory_mode settings in the model initialization optimize memory usage for both GPU and CPU environments, making the pipeline feasible on devices without dedicated graphics hardware, though inference will be significantly slower on CPU.

Where are the model weights stored and how does the automatic download work?

The ONNX weights are stored remotely at https://storage.googleapis.com/ailia-models/latent-diffusion-inpainting/ and include cond_stage_model.onnx, diffusion_model.onnx, and autoencoder.onnx. The check_and_download_models() function in util/model_utils.py verifies local file existence and checksums, automatically fetching missing files from the REMOTE_PATH bucket. This ensures the pipeline works immediately without manual weight management, while caching models locally for subsequent runs.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}