This is the official impelmenation of the paper Towards Coherent Image Inpainting Using Denoising Diffusion Implicit Models.
Image inpainting refers to the task of generating a complete, natural image based on a partially re- vealed reference image. Recently, many research interests have been focused on addressing this problem using fixed diffusion models. These ap- proaches typically directly replace the revealed re- gion of the intermediate or final generated images with that of the reference image or its variants. However, since the unrevealed regions are not di- rectly modified to match the context, it results in incoherence between revealed and unrevealed regions. To address the incoherence problem, a small number of methods introduce a rigorous Bayesian framework, but they tend to introduce mismatches between the generated and the ref- erence images due to the approximation errors in computing the posterior distributions. In this paper, we propose COPAINT, which can coher- ently inpaint the whole image without introducing mismatches. COPAINT also uses the Bayesian framework to jointly modify both revealed and unrevealed regions, but approximates the poste- rior distribution in a way that allows the errors to gradually drop to zero throughout the denoising steps, thus strongly penalizing any mismatches with the reference image. Our experiments ver- ify that COPAINT can outperform the existing diffusion-based methods under both objective and subjective metrics.
Install conda environment by running
conda env create -f environment.yaml
conda activate copaint
bash scripts/download.sh
This script will download pretrained diffusion models from guided-diffusion and repaint.
Then it will download and preprocess CelebA-HQ
dataset according to the data split in Lama. The images we used for ImageNet
dataset are provided in datasets/imagenet100
.
To inpaint a specific image with our algorithm CoPaint
, you can run
python main.py:
--config_file: The configuration file, which specifies the model to use and some hyper-parameters for our method
--input_image: The path to input image
--mask: The path to mask file
--outdir: The path to output folder
--n_samples: The number of images to be generated
--algorithm: The algorithm to be used
Here is an example result on CelebA-HQ
dataset with CoPaint
:
python main.py --config_file configs/celebahq.yaml --input_image examples/celeb-sample.jpg --mask examples/celeb-mask.jpg --outdir images/example --n_samples 1 --algorithm o_ddim
Here is an example result on CelebA-HQ
dataset with CoPaint-TT
:
python main.py --config_file configs/celebahq.yaml --input_image examples/celeb-sample.jpg --mask examples/celeb-mask.jpg --outdir images/example --n_samples 1 --algorithm o_ddim --ddim.schedule_params.jump_length 10 --ddim.schedule_params.jump_n_sample 2 --ddim.schedule_params.use_timetravel
You can tune following hyper-parameters if you are not satisfied with the generated image:
-
--optimize_xt.num_iteration_optimize_xt
: The number of optimization steps$G$ . -
--optimize_xt.lr_xt
: The initial learning rate$\mu_T$ Below are hyper-paremters for the time-travel trick: -
--ddim.schedule_params.use_timetravel
: Whether to use time-travel trick. -
--ddim.schedule_parms.jump_length
: The time-travel interval$\tau$ . -
--ddim.schedule_parms.jum_n_sample
: The time-travel frequency$K$ .
Please refer to our paper for more details.
We provide shell scripts for all algorithms discussed in our paper, copaint, repaint, copaint-tt
, you can find them in scripts
folder. To run a specific algorithm on both datasets:
bash scripts/{algorithm}.sh
The generated images will be put in images/{algorithm}
.
Below is an example to run Copaint
on CelebA-HQ
dataset with half
mask:
python main.py --dataset_name celebahq --algorithm o_ddim --outdir celebresults --mask_type half --config_file configs/celebahq.yaml
If you find our work useful for your research, please consider citing our paper:
@misc{zhang2023coherent,
title={Towards Coherent Image Inpainting Using Denoising Diffusion Implicit Models},
author={Guanhua Zhang and Jiabao Ji and Yang Zhang and Mo Yu and Tommi Jaakkola and Shiyu Chang},
year={2023},
eprint={2304.03322},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
Our implementation is based on following repos: