elin24 / paint-by-example

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

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

Paper | Huggingface Demo

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

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

• 2022-12-07 Release a Gradio demo on Hugging Face Spaces.
• 2022-11-29 Upload code.

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

• Download separate packed files of Open-Images dataset from CVDF's site and unzip them to the directory dataset/open-images/images.
• Download bbox annotations of Open-Images dataset from Open-Images official site and save them to the directory dataset/open-images/annotations.
• Generate bbox annotations of each image in txt format.

  python scripts/read_bbox.py
  

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

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.