Home

Awesome

Palette: Image-to-Image Diffusion Models

Paper | Project

Brief

This is an unofficial implementation of Palette: Image-to-Image Diffusion Models by Pytorch, and it is mainly inherited from its super-resolution version Image-Super-Resolution-via-Iterative-Refinement. The code template is from my another seed project: distributed-pytorch-template.

There are some implementation details with paper descriptions:

Status

Code

Task

I try to finish following tasks in order:

The follow-up experiment is uncertain, due to lack of time and GPU resources:

Results

The DDPM model requires significant computational resources, and we have only built a few example models to validate the ideas in this paper.

Visuals

Celeba-HQ

Results with 200 epochs and 930K iterations, and the first 100 samples in centering mask and irregular mask.

Process_02323Process_02323

Places2 with 128×128 centering mask

Results with 16 epochs and 660K iterations, and the several picked samples in centering mask.

Mask_Places365_test_00209019.jpgMask_Places365_test_00143399.jpgMask_Places365_test_00263905.jpgMask_Places365_test_00144085.jpg
Out_Places365_test_00209019Out_Places365_test_00143399.jpgOut_Places365_test_00263905.jpgOut_Places365_test_00144085.jpg

Uncropping on Places2

Results with 8 epochs and 330K iterations, and the several picked samples in uncropping.

Process_Places365_test_00309553Process_Places365_test_00042384

Metrics

TasksDatasetEMAFID(-)IS(+)
Inpainting with centering maskCeleba-HQFalse5.78733.0705
Inpainting with irregular maskCeleba-HQFalse5.40263.1221

Usage

Environment

pip install -r requirements.txt

Pre-trained Model

DatasetTaskIterationsGPUs×Days×BsURL
Celeba-HQInpainting930K2×5×3Google Drive
Places2Inpainting660K4×8×10Google Drive

Bs indicates sample size per gpu.

Data Prepare

We get most of them from Kaggle, which may be slightly different from official version, and you also can download them from official website.

We use the default division of these datasets for training and evaluation. The file lists we use can be found in Celeba-HQ, Places2.

After you prepared own data, you need to modify the corresponding configure file to point to your data. Take the following as an example:

"which_dataset": {  // import designated dataset using arguments 
    "name": ["data.dataset", "InpaintDataset"], // import Dataset() class
    "args":{ // arguments to initialize dataset
    	"data_root": "your data path",
    	"data_len": -1,
    	"mask_mode": "hybrid"
    } 
},

More choices about dataloader and validation split also can be found in datasets part of configure file.

Training/Resume Training

  1. Download the checkpoints from given links.
  2. Set resume_state of configure file to the directory of previous checkpoint. Take the following as an example, this directory contains training states and saved model:
"path": { //set every part file path
	"resume_state": "experiments/inpainting_celebahq_220426_150122/checkpoint/100" 
},
  1. Set your network label in load_everything function of model.py, default is Network. Follow the tutorial settings, the optimizers and models will be loaded from 100.state and 100_Network.pth respectively.
netG_label = self.netG.__class__.__name__
self.load_network(network=self.netG, network_label=netG_label, strict=False)
  1. Run the script:
python run.py -p train -c config/inpainting_celebahq.json

We test the U-Net backbone used in SR3 and Guided Diffusion, and Guided Diffusion one have a more robust performance in our current experiments. More choices about backbone, loss and metric can be found in which_networks part of configure file.

Test

  1. Modify the configure file to point to your data following the steps in Data Prepare part.
  2. Set your model path following the steps in Resume Training part.
  3. Run the script:
python run.py -p test -c config/inpainting_celebahq.json

Evaluation

  1. Create two folders saving ground truth images and sample images, and their file names need to correspond to each other.

  2. Run the script:

python eval.py -s [ground image path] -d [sample image path]

Acknowledge

Our work is based on the following theoretical works:

and we are benefiting a lot from the following projects: