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

> Master image inpainting with diffusion models using ailia-models latent diffusion pipeline. Learn to fill masked image areas effortlessly while preserving context. Get started today.

- Repository: [axinc-ai/ailia-models](https://github.com/axinc-ai/ailia-models)
- Tags: how-to-guide
- Published: 2026-02-26

---

**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`](https://github.com/axinc-ai/ailia-models/blob/main/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`](https://github.com/axinc-ai/ailia-models/blob/main/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`](https://github.com/axinc-ai/ailia-models/blob/main/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`](https://github.com/axinc-ai/ailia-models/blob/main/latent-diffusion-inpainting.py) demonstrates proper initialization:

```python
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:

```python
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 εₜ:

```python
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:

```python
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:

```python
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:

```bash
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:

```python
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()`:

```python
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`](https://github.com/axinc-ai/ailia-models/blob/main/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`](https://github.com/axinc-ai/ailia-models/blob/main/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`](https://github.com/axinc-ai/ailia-models/blob/main/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`](https://github.com/axinc-ai/ailia-models/blob/main/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.