Simple Denoising Diffusion

Summary

Simple Denoising Diffusion 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 Overview 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 Examples from the dataset: Examples generated by the diffusion model (rotations are because of the data augmentation, it's hilarious though): Example diffusion process: Outro 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...