Home

Awesome

Pseudo Numerical Methods for Diffusion Models on Manifolds (PNDM, PLMS | ICLR2022)

PWC

This repo is the official PyTorch implementation for the paper Pseudo Numerical Methods for Diffusion Models on Manifolds (PNDM, PLMS | ICLR2022)

by Luping Liu, Yi Ren, Zhijie Lin, Zhou Zhao (Zhejiang University).

What does this code do?

This code is not only the official implementation for PNDM, but also a generic framework for DDIM-like models including:

Structure

This code contains three main objects including method, schedule and model. The following table shows the options supported by this code and the role of each object.

ObjectOptionRole
methodDDIM, S-PNDM, F-PNDM, FON, PFthe numerical method used to generate samples
schedulelinear, quad, cosinethe schedule of adding noise to images
modelDDIM, iDDPM, PF, PF_deepthe neural network used to fit noise

All of them can be combined at will, so this code provide at least 5x3x4=60 choices to generate samples.

Integration with 🤗 Diffusers library

PNDM is now also available in 🧨 Diffusers and accesible via the PNDMPipeline. Diffusers allows you to test PNDM in PyTorch in just a couple lines of code.

You can install diffusers as follows:

pip install diffusers torch accelerate

And then try out the sampler/scheduler with just a couple lines of code:

from diffusers import PNDMPipeline

model_id = "google/ddpm-cifar10-32"

# load model and scheduler
pndm = PNDMPipeline.from_pretrained(model_id)

# run pipeline in inference (sample random noise and denoise)
image = pndm(num_inference_steps=50).images[0]

# save image
image.save("pndm_generated_image.png")

The PNDM scheduler can also be used with more powerful diffusion models such as Stable Diffusion

You simply need to accept the license on the Hub, login with huggingface-cli login and install transformers:

pip install transformers

Then you can run:

from diffusers import StableDiffusionPipeline, PNDMScheduler

pndm = PNDMScheduler.from_config("runwayml/stable-diffusion-v1-5", subfolder="scheduler")
pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", scheduler=pndm)

image = pipeline("An astronaut riding a horse.").images[0]
image.save("astronaut_riding_a_horse.png")

How to run the code

Dependencies

Run the following to install a subset of necessary python packages for our code.

pip install -r requirements.txt

Tip: mpi4py can make the generation process faster using multi-gpus. It is not necessary and can be removed freely.

Usage

Evaluate our models through main.py.

python main.py --runner sample --method F-PNDM --sample_speed 50 --device cuda --config ddim_cifar10.yml --image_path temp/results --model_path temp/models/ddim/ema_cifar10.ckpt

Train our models through main.py.

python main.py --runner train --device cuda --config ddim_cifar10.yml --train_path temp/train

Checkpoints & statistics

All checkpoints of models and precalculated statistics for FID are provided in this Google Drive.

References

If you find the code useful for your research, please consider citing:

@inproceedings{liu2022pseudo,
    title={Pseudo Numerical Methods for Diffusion Models on Manifolds},
    author={Luping Liu and Yi Ren and Zhijie Lin and Zhou Zhao},
    booktitle={International Conference on Learning Representations},
    year={2022},
    url={https://openreview.net/forum?id=PlKWVd2yBkY}
}

This work is built upon some previous papers which might also interest you: