This is work in progress -- for now development has been paused.

See also: nicecv#2

ComfyUI_DiffHarmony: ComfyUI implementation of DiffHarmony

Installation

Note: We assume that you need to create a Docker container, so we do not download the model weights at runtime.

1. Install this repo as a custom node in ComfyUI.
2. Manually download weights from Google Drive and save them in the checkpoints subdirectory -- your file tree should look like this:

.../ComfyUI/custom_nodes/ComfyUI_DiffHarmony/checkpoints
├── base
│   ├── feature_extractor
│   │   └── preprocessor_config.json
│   ├── model_index.json
│   ├── safety_checker
│   │   ├── config.json
│   │   └── pytorch_model.bin
│   ├── scheduler
│   │   └── scheduler_config.json
│   ├── text_encoder
│   │   ├── config.json
│   │   └── pytorch_model.bin
│   ├── tokenizer
│   │   ├── merges.txt
│   │   ├── special_tokens_map.json
│   │   ├── tokenizer_config.json
│   │   └── vocab.json
│   ├── unet
│   │   ├── config.json
│   │   └── diffusion_pytorch_model.bin
│   └── vae
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── condition_vae
│   ├── config.json
│   └── diffusion_pytorch_model.safetensors
├── inverse
│   ├── feature_extractor
│   │   └── preprocessor_config.json
│   ├── model_index.json
│   ├── safety_checker
│   │   ├── config.json
│   │   └── pytorch_model.bin
│   ├── scheduler
│   │   └── scheduler_config.json
│   ├── text_encoder
│   │   ├── config.json
│   │   └── pytorch_model.bin
│   ├── tokenizer
│   │   ├── merges.txt
│   │   ├── special_tokens_map.json
│   │   ├── tokenizer_config.json
│   │   └── vocab.json
│   ├── vae
│   │   ├── config.json
│   │   └── diffusion_pytorch_model.bin
│   └── weights-102800
│       ├── config.json
│       └── diffusion_pytorch_model.safetensors
└── refinement
    ├── config.json
    └── diffusion_pytorch_model.safetensors

3. ...
4. Profit

Original README below

---

DiffHarmony: Latent Diffusion Model Meets Image Harmonization

The official pytorch implementation of DiffHarmony and DiffHarmony++ (paper release soon).

Full Conference Poster is here.

Preparation

enviroment

First, prepare a virtual env. You can use conda or anything you like.

python 3.10
pytorch 2.2.0
cuda 12.1
xformers 0.0.24

Then, install requirements.

pip install -r requirements.txt

dataset

Download iHarmony4 dataset from here.

Make sure the structure is just like that:

data/iHarmony4
|- HCOCO
    |- composite_images
    |- masks
    |- real_images
    |- ...
|- HAdobe5k
|- HFlickr
|- Hday2night
|- train.jsonl
|- test.jsonl

The content in train.jsonl fit the following format

{"file_name": "HAdobe5k/composite_images/a0001_1_1.jpg", "text": ""}
{"file_name": "HAdobe5k/composite_images/a0001_1_2.jpg", "text": ""}
{"file_name": "HAdobe5k/composite_images/a0001_1_3.jpg", "text": ""}
{"file_name": "HAdobe5k/composite_images/a0001_1_4.jpg", "text": ""}
...

All file_name are from the original IHD_train.txt. Same way with test.jsonl and IHD_test.txt.

Training

Train diffharmony model

sh scripts/train_diffharmony.sh

Train refinement model

sh scripts/train_refinement_stage.sh

Train condition vae (cvae)

sh scripts/train_cvae.sh

Train diffharmony-gen and cvae-gen

Just add this in your training args:

$script
    ...
    --mode "inverse"

Basically it will use ground truth images as condition instead of composite images.

(optional) online training of condition vae

refer to scripts/train/cvae_online.py

(optional) train cvae with generated data

refer to scripts/train/cvae_with_gen_data.py

Purpose here is trying to improve cvae performance further on specific domain, i.e. our generated dataset.

Inference

Inference iHarmony4 dataset

sh scripts/inference.sh

use diffharmony-gen and cvae-gen to augment HFlickr and Hday2night

sh scripts/inference_generate_data.sh

The all_mask_metadata.jsonl file as its name fits following format:

{"file_name": "masks/f800_1.png", "text": ""}
{"file_name": "masks/f801_1.png", "text": ""}
{"file_name": "masks/f803_1.png", "text": ""}
{"file_name": "masks/f804_1.png", "text": ""}
...

Make HumanHarmony dataset

First, generate some candidate composite images.

Then, use harmony classifier to select the most unharmonious images.

python scripts/misc/classify_cand_gen_data.py

Evaluation

sh scripts/evaluate.sh

Pretrained Models

Baidu, code: aqqd

Citation

If you find this work useful, please consider citing:

@inproceedings{zhou2024diffharmony,
  title={DiffHarmony: Latent Diffusion Model Meets Image Harmonization},
  author={Zhou, Pengfei and Feng, Fangxiang and Wang, Xiaojie},
  booktitle={Proceedings of the 2024 International Conference on Multimedia Retrieval},
  pages={1130--1134},
  year={2024}
}

Contact

If you have any questions, please feel free to contact me via zhoupengfei@bupt.edu.cn .