xdiffusion

A unified media (Image, Video, Audio, Text) diffusion repository, for education and learning.

If you are looking for just the lessons on image diffusion models, checkout mindiffusion for more image diffusion model paper implementations.

Requirements

This package built using PyTorch and written in Python 3. Due to package dependencies, this repo requires python 3.10 or greater. To setup an environment to run all of the lessons, we suggest using conda or venv:

> python3 -m venv xdiffusion_env
> source xdiffusion_env/bin/activate
> pip install --upgrade pip
> pip install -r requirements.txt

I find pyenv-virtualenv to be very helpful as well, in managing both the virtual environment as well as the python version dependencies.

> pyenv install 3.10.15
> pyenv virtualenv 3.10.15 xdiffusion_env
> pyvenv activate xdiffusion_env     
> pip install --upgrade pip
> pip install -r requirements.txt

All lessons are designed to be run from the root of the repository, and you should set your python path to include the repository root:

> export PYTHONPATH=$(pwd)

If you have issues with PyTorch and different CUDA versions on your instance, make sure to install the correct version of PyTorch for the CUDA version on your machine. For example, if you have CUDA 11.8 installed, you can install PyTorch using:

> pip install torch==2.1.0 torchvision --index-url https://download.pytorch.org/whl/cu118

Image Diffusion

Training Datasets

In this repository, we will be working with the MNIST dataset because it is simple and can be trained in real time with minimal GPU power and memory. The main difference between MNIST and other datasets is the single channel of the imagery, versus 3 channels in most other datasets. We will make sure that the models we build can easily accomodate 1- or 3-channel data, so that you can test the models we build on other datasets.

Image Models

The following is a list of the supported image models, their current results, and a link to their configuration files and documentation.

Video Diffusion

Training Datasets

Due to the resource constraints of most models, we have decided to use the Moving MNIST dataset to train on. Moving MNIST is a simple dataset similar to MNIST, of digits which move around the screen. It is an unlabeled dataset, so we do not have access to text labels to determine which digits are moving around the screen, but we will address that deficiency as well. We train at a reduced resolution of 32x32, due to the resource constraints that most models require. This allows us to train most diffusion models on a T4 instance, which is free to run on Google Colab. We limit training and sample generation to 16 frames, even though the source dataset contains 20 frames.

Sample from the original dataset:

Video Diffusion Models

The following is a list of the supported video models, their current results, and a link to their configuration files and documentation.

Audio Diffusion

Audio Diffusion Models

The following is a list of the supported audio diffusion models, their current results, and a link to their configuration files and documentation.

TODO

• Port the image diffusion repository
• Port the video diffusion respository
• Unify all of the different attention mechanisms under Transformer based (B, L, D) and pixel based (B, C, H, W).