morbi25 / Palette-Image-to-Image-Diffusion-Models

Unofficial implementation of Palette: Image-to-Image Diffusion Models by Pytorch

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:

• We adapted the U-Net architecture used in Guided-Diffusion, which give a substantial boost to sample quality.
• We used the attention mechanism in low-resolution features (16×16) like vanilla DDPM.
• We encode the $\gamma$ rather than $t$ in Palette and embed it with affine transformation.
• We fix the variance $Σ_\theta(x_t, t)$ to a constant during the inference as described in Palette.

Status

Code

• Diffusion Model Pipeline
• Train/Test Process
• Save/Load Training State
• Logger/Tensorboard
• Multiple GPU Training (DDP)
• EMA
• Metrics (now for FID, IS)
• Dataset (now for inpainting, uncropping, colorization)
• Google colab script 🌟(now for inpainting)

Task

I try to finish following tasks in order:

• Inpainting on CelebaHQ🚀 (Google Colab)
• Inpainting on Places2 with 128×128 centering mask🚀

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

• Uncropping on Places2
• Colorization on ImageNet val set

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.

Places2 with 128×128 centering mask

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

Uncropping on Places2

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

Metrics

Tasks	Dataset	EMA	FID(-)	IS(+)
Inpainting with centering mask	Celeba-HQ	False	5.7873	3.0705
Inpainting with irregular mask	Celeba-HQ	False	5.4026	3.1221

Usage

Environment

pip install -r requirements.txt

Pre-trained Model

Dataset	Task	Iterations	GPUs×Days×Bs	URL
Celeba-HQ	Inpainting	930K	2×5×3	Google Drive
Places2	Inpainting	660K	4×8×10	Google 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.

• CelebA-HQ resized (256x256) Kaggle
• Places2 Official | Places2 Kaggle
• ImageNet Official

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" 
},

2. 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)

3. 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:

• Denoising Diffusion Probabilistic Models
• Palette: Image-to-Image Diffusion Models
• Diffusion Models Beat GANs on Image Synthesis

and we are benefiting a lot from the following projects:

• openai/guided-diffusion
• LouisRouss/Diffusion-Based-Model-for-Colorization