luping-liu / pndm

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

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

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:

• Pseudo Numerical Methods for Diffusion Models on Manifolds (PNDM)
• Denoising Diffusion Implicit Models (DDIM)
• Score-Based Generative Modeling through Stochastic Differential Equations (PF)
• Improved Denoising Diffusion Probabilistic Models (iDDPM)

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.

Object	Option	Role
method	DDIM, S-PNDM, F-PNDM, FON, PF	the numerical method used to generate samples
schedule	linear, quad, cosine	the schedule of adding noise to images
model	DDIM, iDDPM, PF, PF_deep	the 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

• runner (train|sample): choose the mode of runner
• method (DDIM|FON|S-PNDM|F-PNDM|PF): choose the numerical methods
• sample_speed: control the total generation step
• device (cpu|cuda:0): choose the device to use
• config: choose the config file
• image_path: choose the path to save images
• model_path: choose the path of model

Train our models through main.py.

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

• train_path: choose the path to save training status

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:

• Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems 33 (2020): 6840-6851.
• Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising Diffusion Implicit Models. International Conference on Learning Representations. 2021.
• Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations. International Conference on Learning Representations. 2021.

pndm's Issues

Training time

Hi, thanks for your interesting work.

What is your training time on each dataset?

A question of “Pseudo Numerical Method“

Hi, thank you for your interesting work.

I have a question. Is the "pseudo numerical method" one of the methods of classical numerical? Or is it just named specifically for your research? As I searched for the name in the numerical methods domain, but I didn't find it. So I guess it is a specific term you named specifically for your work?

Looking forward to your reply. Thanks!

The aliyundrive is empty.

I am writing a paper and need to use your code, but the checkpoint and the fid is empty in the aliyundirve.

How to train a model for new Data

Hi @luping-liu! Thank you for your job. It´s really interesting!

Could you help me to know how could I train a model with new data? I have Image in 256x256 size and I have divided it into validation and train because it is medical data and must not be mixed because the train and the validation data will be very correlated.

Regards.

Segmentation fault (11)

Hey I am back, so i moved the program over to a google colab. Ran this command:

!python main.py --runner sample --method F-PNDM --sample_speed 50 --device cuda --config ddim_cifar10.yml --image_path results --model_path models/ddim_cifar10.ckpt
Everything was going fine and then ->

any advice?

The link of onedrive for models is out of date

AttributeError: 'Namespace' object has no attribute 'exp'

Predicted x_0 during denoising

Hi, thanks for the great work!
I'm wondering is it possible to get a predicted x_0 during the denoising process like DDIM? If yes, could you maybe expalin a bit?
Thanks a lot!

Is the implementation for FON really 4-th order?

Thanks for your inspiring work.

I found the implementation of FON is confusing, which may not what claimed in the paper.

Paper claims

However,

def gen_fon(img, t, t_next, model, alphas_cump, ets):
    t_list = [t, (t + t_next) / 2.0, t_next]
    if len(ets) > 2:
        noise = model(img, t)
        img_next = transfer(img, t, t-1, noise, alphas_cump)
        delta = img_next - img
        ets.append(delta)
    else:
        noise = model(img, t_list[0])
        img_ = transfer(img, t, t - 1, noise, alphas_cump)
        delta_1 = img_ - img
        ets.append(delta_1)

        img_2 = img + delta_1 * (t - t_next).view(-1, 1, 1, 1) / 2.0
        noise = model(img_2, t_list[1])
        img_ = transfer(img, t, t - 1, noise, alphas_cump)
        delta_2 = img_ - img

        img_3 = img + delta_2 * (t - t_next).view(-1, 1, 1, 1) / 2.0
        noise = model(img_3, t_list[1])
        img_ = transfer(img, t, t - 1, noise, alphas_cump)
        delta_3 = img_ - img

        img_4 = img + delta_3 * (t - t_next).view(-1, 1, 1, 1)
        noise = model(img_4, t_list[2])
        img_ = transfer(img, t, t - 1, noise, alphas_cump)
        delta_4 = img_ - img
        delta = (1 / 6.0) * (delta_1 + 2*delta_2 + 2*delta_3 + delta_4)

    img_next = img + delta * (t - t_next).view(-1, 1, 1, 1)
    return img_next

After len(ets) > 2, delta seems only use first order info?

How do I generate a .ckpt file after training?

After training a model, how do I create a .ckpt file to use with the sample runner? The training loop writes out .pth files that seem almost correct, but they are slightly mismatched. After a bit of research, it appears that the model is saved with the nn.DataParallel module, but is loaded without it. Reference: https://discuss.pytorch.org/t/solved-keyerror-unexpected-key-module-encoder-embedding-weight-in-state-dict/1686/2

Conditional PNDM

Hi, I wanted to use PNDM for conditional image generation and would like to know how to integrate that. Thanks