Skip to content

Janspiry/Palette-Image-to-Image-Diffusion-Models

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

39 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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:

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.

Process_02323 Process_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.jpg Mask_Places365_test_00143399.jpg Mask_Places365_test_00263905.jpg Mask_Places365_test_00144085.jpg
Out_Places365_test_00209019 Out_Places365_test_00143399.jpg Out_Places365_test_00263905.jpg Out_Places365_test_00144085.jpg

Uncropping on Places2

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

Process_Places365_test_00309553 Process_Places365_test_00042384

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.

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: