Fast Diffusion Model

This is an official PyTorch implementation of Fast Diffusion Model. See the paper here. If you find our FDM helpful or heuristic to your projects, please cite this paper and also star this repository. Thanks!

@misc{wu2023fast,
      title={Fast Diffusion Model}, 
      author={Zike Wu and Pan Zhou and Kenji Kawaguchi and Hanwang Zhang},
      year={2023},
      eprint={2306.06991},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

Acknowledgement: This repo is based on the following amazing projects: EDM and DPM-Solver.

Results

Requirements

All experiments were conducted using PyTorch 1.13.0, CUDA 11.7.1, and CuDNN 8.5.0. We strongly recommend to use the provided Dockerfile to build an image to reproduce our experiments.

Pre-trained models

We provide pre-trained models for our FDMs along with the baseline models on Hugging Face. Download the checkpoints here.

Preparing datasets

CIFAR-10: Download the CIFAR-10 python version and convert to ZIP archive:

python dataset_tool.py --source=downloads/cifar10/cifar-10-python.tar.gz \
    --dest=datasets/cifar10-32x32.zip
python fid.py ref --data=datasets/cifar10-32x32.zip --dest=fid-refs/cifar10-32x32.npz

FFHQ: Download the Flickr-Faces-HQ dataset as 1024x1024 images and convert to ZIP archive at 64x64 resolution:

python dataset_tool.py --source=downloads/ffhq/images1024x1024 \
    --dest=datasets/ffhq-64x64.zip --resolution=64x64
python fid.py ref --data=datasets/ffhq-64x64.zip --dest=fid-refs/ffhq-64x64.npz

AFHQv2: Download the updated Animal Faces-HQ dataset (afhq-v2-dataset) and convert to ZIP archive at 64x64 resolution:

python dataset_tool.py --source=downloads/afhqv2 \
    --dest=datasets/afhqv2-64x64.zip --resolution=64x64
python fid.py ref --data=datasets/afhqv2-64x64.zip --dest=fid-refs/afhqv2-64x64.npz

Training from scratch

Train FDM for class-conditional CIFAR-10 using 8 GPUs:

# EDM-FDM
torchrun --standalone --nproc_per_node=8 train.py --outdir=training-output \
    --data=datasets/cifar10-32x32.zip --cond=1 --arch=ddpmpp \
    --precond=fdm_edm --warmup_ite=200 

# VP-FDM
torchrun --standalone --nproc_per_node=8 train.py --outdir=training-output \
    --data=datasets/cifar10-32x32.zip --cond=1 --arch=ddpmpp --cres=1,2,2,2 \
    --precond=fdm_vp --warmup_ite=400

# VE-FDM
torchrun --standalone --nproc_per_node=8 train.py --outdir=training-output \
    --data=datasets/cifar10-32x32.zip --cond=1 --arch=ncsnpp --cres=1,2,2,2 \
    --precond=fdm_ve --warmup_ite=400 

Train FDM for unconditional FFHQ using 8 GPUs:

# EDM-FDM
torchrun --standalone --nproc_per_node=8 train.py --outdir=training-output
    --data=datasets/ffhq-64x64.zip --cond=0 --arch=ddpmpp \
    --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 \
    --precond=fdm_edm --warmup_ite=800 --fdm_multipler=1

# VP-FDM
torchrun --standalone --nproc_per_node=8 train.py --outdir=training-output \
    --data=datasets/ffhq-64x64.zip --cond=0 --arch=ddpmpp \
    --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 \
    --precond=fdm_vp --warmup_ite=400 --fdm_multipler=1

# VE-FDM
torchrun --standalone --nproc_per_node=8 train.py --outdir=training-output \
    --data=datasets/ffhq-64x64.zip --cond=0 --arch=ncsnpp \
    --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 \
    --precond=fdm_ve --warmup_ite=400

Train FDM for unconditional AFHQv2 using 8 GPUs:

# EDM-FDM
torchrun --standalone --nproc_per_node=8 train.py --outdir=training-output
    --data=datasets/afhqv2-64x64.zip --cond=0 --arch=ddpmpp \
    --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.25 --augment=0.15 \
    --precond=fdm_edm --warmup_ite=400

# VP-FDM
torchrun --standalone --nproc_per_node=8 train.py --outdir=training-output \
    --data=datasets/afhqv2-64x64.zip --cond=0 --arch=ddpmpp \
    --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.25 --augment=0.15 \
    --precond=fdm_vp --warmup_ite=400

# VE-FDM
torchrun --standalone --nproc_per_node=8 train.py --outdir=training-output \
    --data=datasets/afhqv2-64x64.zip --cond=0 --arch=ncsnpp \
    --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.25 --augment=0.15 \
    --precond=fdm_ve --warmup_ite=400

Calculating FID

To compute Fréchet inception distance (FID) for a given model and sampler, first generate 50,000 random images and then compare them against the dataset reference statistics using fid.py, replace $PATH_TO_CHECKPOINT with the path to the checkpoint:

# Generate 50000 images 
torchrun --standalone --nproc_per_node=8 generate.py --outdir=fid \
    --seeds=0-49999 --subdirs --network=$PATH_TO_CHECKPOINT
# Calculate FID
torchrun --standalone --nproc_per_node=8 fid.py calc --images=fid \
    --ref=fid-refs/cifar10-32x32.npz

Note that the generated images should be evaluated against the same reference dataset that the model was originally trained on. Please ensure to replace the --ref option with the correct one (e.g., fid-refs/ffhq-64x64.npz or fid-refs/afhqv2-64x64.npz) to obtain the right FID score. Addtionally, you can use --solver=dpm option to generate images with DPM-Solver++.