This repository contains a bare-bone implementation of denoising diffusion [1,2] in PyTorch, with majority of its code taken from The Annotated Diffusion and Phil Wang's diffusion repository. Both resources are great to get started with diffusion models but they were still a bit convoluted for me when I first started learning about diffusion models so I refactored majority of The Annotated Diffusion's implementation and made a bare-bone implementation with functions and classes logically separated into different files as a learning exercise. My goal was to understand the building blocks of diffusion models in order to use them in some upcoming projects. I'm sharing this repo in hopes that my exercise will be useful for you in understanding more complex implementations.
Code is organized under src folder as follows:
-
funct_diffusion.py - Contains all necessary functions for forward and backward diffusion process, including the scheduler.
-
cls_dataset.py - Contains data-related functions and classes. I used a single class (n01443537 - Carassius auratus - Goldfish) with some augmentations (e.g., rotations and flips), that's why generated images have several upside down fishes.
-
cls_model.py - Contains the model. The model in this repo is basically a copy paste of The Annotated Diffusion's implementation, except for dim_mults=(1, 2, 4, 8) and channe=3 (RGB).
-
main_train_diffusion.py - I wanted to separate training and generation into two different files to digest what parameter is needed for what. This file is used to train the diffusion model.
-
main_generate_images.py - Generates images using the trained model.
Examples from the dataset:
Examples generated by the diffusion model (rotations are because of the data augmentation, it's hilarious though):
Example diffusion process:
As you can see, generated images are not as crisp as those from the dataset. There are many improvements that can be incorporated to improve the image quality, with each of those adding more to the complexity. Phil Wang's diffusion repository is a great place to discover some of those methods.
torch
torchvision
datasets
PIL
numpy
[1] Song and Ermon, Generative Modeling by Estimating Gradients of the Data Distribution
[2] Ho et al., Denoising Diffusion Probabilistic Models