Diffusion Transformers(DiT) Implementation in PyTorch

DiT Tutorial Video

Sample Output for DiT on CelebHQ

Trained for 200 epochs

---

This repository implements DiT in PyTorch for diffusion models. It provides code for the following:

• Training and inference of VAE on CelebHQ (128x128 to 32x32x4 latents)
• Training and Inference of DiT using trained VAE on CelebHQ
• Configurable code for training all models from DIT-S to DIT-XL

This is very similar to official DiT implementation except the following changes.

• Since training is on celebhq its unconditional generation as of now (but can be easily modified to class conditional or text conditional as well)
• Variance is fixed during training and not learned (like original DDPM)
• No EMA
• Ability to train VAE
• Ability to save latents of VAE for faster training

Setup

• Create a new conda environment with python 3.10 then run below commands
• conda activate <environment_name>
• git clone https://github.com/explainingai-code/DiT-PyTorch.git
• cd DiT-PyTorch
• pip install -r requirements.txt
• Download lpips weights by opening this link in browser(dont use cURL or wget) https://github.com/richzhang/PerceptualSimilarity/blob/master/lpips/weights/v0.1/vgg.pth and downloading the raw file. Place the downloaded weights file in models/weights/v0.1/vgg.pth

---

Data Preparation

CelebHQ

For setting up on CelebHQ, simply download the images from the official repo of CelebMASK HQ here. and add it to data directory. Ensure directory structure is the following

DiT-PyTorch
    -> data
        -> CelebAMask-HQ
            -> CelebA-HQ-img
                -> *.jpg

---

Configuration

Allows you to play with different components of DiT and autoencoder

• config/celebhq.yaml - Configuration used for celebhq dataset

Important configuration parameters

• autoencoder_acc_steps : For accumulating gradients if image size is too large and a large batch size cant be used.
• save_latents : Enable this to save the latents , during inference of autoencoder. That way DiT training will be faster

---

Training

The repo provides training and inference for CelebHQ (Unconditional DiT)

For working on your own dataset:

• Create your own config and have the path in config point to images (look at celebhq.yaml for guidance)
• Create your own dataset class which will just collect all the filenames and return the image or latent in its getitem method. Look at celeb_dataset.py for guidance

Once the config and dataset is setup:

• First train the auto encoder on your dataset using this section
• For training and inference of Unconditional DiT follow this section

Training AutoEncoder for DiT

• For training autoencoder on celebhq,ensure the right path is mentioned in celebhq.yaml
• For training autoencoder on your own dataset
  • Create your own config and have the path point to images (look at celebhq.yaml for guidance)
  • Create your own dataset class, similar to celeb_dataset.py
• Call the desired dataset class in training file here
• For training autoencoder run python -m tools.train_vae --config config/celebhq.yaml for training autoencoder with the desire config file
• For inference make sure save_latent is True in the config
• For inference run python -m tools.infer_vae --config config/celebhq.yaml for generating reconstructions and saving latents with right config file.

Training Unconditional DiT

Train the autoencoder first and setup dataset accordingly.

For training unconditional DiT ensure the right dataset is used in train_vae_dit.py

• python -m tools.train_vae_dit --config config/celebhq.yaml for training unconditional DiT using right config
• python -m tools.sample_vae_dit --config config/celebhq.yaml for generating images using trained DiT

Output

Outputs will be saved according to the configuration present in yaml files.

For every run a folder of task_name key in config will be created

During training of autoencoder the following output will be saved

• Latest Autoencoder and discriminator checkpoint in task_name directory
• Sample reconstructions in task_name/vae_autoencoder_samples

During inference of autoencoder the following output will be saved

• Reconstructions for random images in task_name
• Latents will be save in task_name/vae_latent_dir_name if mentioned in config

During training and inference of unconditional DiT following output will be saved:

• During training we will save the latest checkpoint in task_name directory
• During sampling, unconditional sampled image grid for all timesteps in task_name/samples/*.png . The final decoded generated image will be x0_0.png. Images from x0_999.png to x0_1.png will be latent image predictions of denoising process from T=999 to T=1. Generated Image is at T=0

Citations

@misc{peebles2023scalablediffusionmodelstransformers,
      title={Scalable Diffusion Models with Transformers}, 
      author={William Peebles and Saining Xie},
      year={2023},
      eprint={2212.09748},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2212.09748}, 
}