Repository files navigation

Paper |Project

This is an unofficial implementation ofPalette: Image-to-Image Diffusion Models byPytorch, and it is mainly inherited from its super-resolution versionImage-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 inGuided-Diffusion, which give a substantial boost to sample quality.
• We used the attention mechanism in low-resolution features (16×16) like vanillaDDPM.
• We encode the$\gamma$ rather than$t$ inPalette and embed it with affine transformation.
• We fix the variance$Σ_\theta(x_t, t)$ to a constant during the inference as described inPalette.

• 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)

I try to finish following tasks in order:

• Inpainting onCelebaHQ🚀 (Google Colab)
• Inpainting onPlaces2 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

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

Results with 200 epochs and 930K iterations, and the first 100 samples incentering mask andirregular mask.

Results with 16 epochs and 660K iterations, and the severalpicked samples incentering mask.

Results with 8 epochs and 330K iterations, and the severalpicked samples inuncropping.

pipinstall-rrequirements.txt

Bs indicates sample size per gpu.

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 inCeleba-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 aboutdataloader andvalidation split also can be found indatasets part of configure file.

1. Download the checkpoints from given links.
2. Setresume_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 inload_everything function ofmodel.py, default isNetwork. 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:

pythonrun.py-ptrain-cconfig/inpainting_celebahq.json

We test the U-Net backbone used inSR3 andGuided Diffusion, andGuided Diffusion one have a more robust performance in our current experiments. More choices aboutbackbone,loss andmetric can be found inwhich_networks part of configure file.

1. Modify the configure file to point to your data following the steps inData Prepare part.
2. Set your model path following the steps inResume Training part.
3. Run the script:

pythonrun.py-ptest-cconfig/inpainting_celebahq.json

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:

pythoneval.py-s [groundimagepath]-d [sampleimagepath]

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