Home

Awesome

Paint by Example: Exemplar-based Image Editing with Diffusion Models

Teaser

Paper | Huggingface Demo

<!-- <br> -->

Binxin Yang, Shuyang Gu, Bo Zhang, Ting Zhang, Xuejin Chen, Xiaoyan Sun, Dong Chen and Fang Wen.

<!-- <br> -->

Abstract

Language-guided image editing has achieved great success recently. In this paper, for the first time, we investigate exemplar-guided image editing for more precise control. We achieve this goal by leveraging self-supervised training to disentangle and re-organize the source image and the exemplar. However, the naive approach will cause obvious fusing artifacts. We carefully analyze it and propose an information bottleneck and strong augmentations to avoid the trivial solution of directly copying and pasting the exemplar image. Meanwhile, to ensure the controllability of the editing process, we design an arbitrary shape mask for the exemplar image and leverage the classifier-free guidance to increase the similarity to the exemplar image. The whole framework involves a single forward of the diffusion model without any iterative optimization. We demonstrate that our method achieves an impressive performance and enables controllable editing on in-the-wild images with high fidelity.

News

Requirements

A suitable conda environment named Paint-by-Example can be created and activated with:

conda env create -f environment.yaml
conda activate Paint-by-Example

Pretrained Model

We provide the checkpoint (Google Drive | Hugging Face) that is trained on Open-Images for 40 epochs. By default, we assume that the pretrained model is downloaded and saved to the directory checkpoints.

Testing

To sample from our model, you can use scripts/inference.py. For example,

python scripts/inference.py \
--plms --outdir results \
--config configs/v1.yaml \
--ckpt checkpoints/model.ckpt \
--image_path examples/image/example_1.png \
--mask_path examples/mask/example_1.png \
--reference_path examples/reference/example_1.jpg \
--seed 321 \
--scale 5

or simply run:

sh test.sh

Visualization of inputs and output:

Training

Data preparing

The data structure is like this:

dataset
├── open-images
│  ├── annotations
│  │  ├── class-descriptions-boxable.csv
│  │  ├── oidv6-train-annotations-bbox.csv
│  │  ├── test-annotations-bbox.csv
│  │  ├── validation-annotations-bbox.csv
│  ├── images
│  │  ├── train_0
│  │  │  ├── xxx.jpg
│  │  │  ├── ...
│  │  ├── train_1
│  │  ├── ...
│  │  ├── validation
│  │  ├── test
│  ├── bbox
│  │  ├── train_0
│  │  │  ├── xxx.txt
│  │  │  ├── ...
│  │  ├── train_1
│  │  ├── ...
│  │  ├── validation
│  │  ├── test

Download the pretrained model of Stable Diffusion

We utilize the pretrained Stable Diffusion v1-4 as initialization, please download the pretrained models from Hugging Face and save the model to directory pretrained_models. Then run the following script to add zero-initialized weights for 5 additional input channels of the UNet (4 for the encoded masked-image and 1 for the mask itself).

python scripts/modify_checkpoints.py

Training Paint by Example

To train a new model on Open-Images, you can use main.py. For example,

python -u main.py \
--logdir models/Paint-by-Example \
--pretrained_model pretrained_models/sd-v1-4-modified-9channel.ckpt \
--base configs/v1.yaml \
--scale_lr False

or simply run:

sh train.sh

Test Benchmark

We build a test benchmark for quantitative analysis. Specifically, we manually select 3500 source images from MSCOCO validation set, each image contains only one bounding box. Then we manually retrieve a reference image patch from MSCOCO training set. The reference image usually shares a similar semantic with mask region to ensure the combination is reasonable. We named it as COCO Exemplar-based image Editing benchmark, abbreviated as COCOEE. This test benchmark can be downloaded from Google Drive.

Quantitative Results

By default, we assume that the COCOEE is downloaded and saved to the directory test_bench. To generate the results of test bench, you can use scripts/inference_test_bench.py. For example,

python scripts/inference_test_bench.py \
--plms \
--outdir results/test_bench \
--config configs/v1.yaml \
--ckpt checkpoints/model.ckpt \
--scale 5

or simply run:

bash inference_test_bench.sh

FID Score

By default, we assume that the test set of COCO2017 is downloaded and saved to the directory dataset. The data structure is like this:

dataset
├── coco
│  ├── test2017
│  │  ├── xxx.jpg
│  │  ├── xxx.jpg
│  │  ├── ...
│  │  ├── xxx.jpg

Then convert the images into square images with 512 solution.

python scripts/create_square_gt_for_fid.py

To calculate FID score, simply run:

python eval_tool/fid/fid_score.py --device cuda \
test_bench/test_set_GT \
results/test_bench/results

QS Score

Please download the model weights for QS score from Google Drive and save the model to directory eval_tool/gmm. To calculate QS score, simply run:

python eval_tool/gmm/gmm_score_coco.py results/test_bench/results \
--gmm_path eval_tool/gmm/coco2017_gmm_k20 \
--gpu 1

CLIP Score

To calculate CLIP score, simply run:

python eval_tool/clip_score/region_clip_score.py \
--result_dir results/test_bench/results

Citing Paint by Example

@article{yang2022paint,
  title={Paint by Example: Exemplar-based Image Editing with Diffusion Models},
  author={Binxin Yang and Shuyang Gu and Bo Zhang and Ting Zhang and Xuejin Chen and Xiaoyan Sun and Dong Chen and Fang Wen},
  journal={arXiv preprint arXiv:2211.13227},
  year={2022}
}

Acknowledgements

This code borrows heavily from Stable Diffusion. We also thank the contributors of OpenAI's ADM codebase and https://github.com/lucidrains/denoising-diffusion-pytorch.

Maintenance

Please open a GitHub issue for any help. If you have any questions regarding the technical details, feel free to contact us.

License

The codes and the pretrained model in this repository are under the CreativeML OpenRAIL M license as specified by the LICENSE file.

The test benchmark, COCOEE, belongs to the COCO Consortium and are licensed under a Creative Commons Attribution 4.0 License.