YashNagraj75/Diffusion-Transformer

Diffusion Transformer (DiT) — CelebA-HQ Face Generation

A Diffusion Transformer (DiT) trained on CelebA-HQ for unconditional face image generation at 128x128 resolution. The model uses a Vision Transformer backbone in place of the conventional U-Net for the denoising network, operating in the latent space of a VQ-VAE.

Model Description

This is a two-stage pipeline:

1. 1.Stage 1 — VQ-VAE: Compresses 128x128 RGB images into a 4-channel discrete latent space with codebook size 8192.
2. 2.Stage 2 — DiT: A transformer-based denoising model that operates on flattened image patches in the VQ-VAE latent space.

DiT Architecture

The DiT (Peebles & Xie, 2023) replaces the U-Net backbone with a standard Vision Transformer (ViT) encoder. Each image latent is divided into non-overlapping patches, linearly embedded, and processed by a stack of transformer blocks with time-step conditioning via adaptive layer norm.

VQ-VAE Architecture

Diffusion Process

Training Details

• —Dataset: CelebA-HQ, center-cropped and resized to 128x128, normalized to [-1, 1]
• —Data loaded from parquet files via a custom ParquetImageDataset

Generated Samples

The repository includes generated face samples in celebhq/samples/ (x0_*.jpg), produced by running the trained DiT in reverse diffusion from Gaussian noise.

Repository Contents

References

• —Peebles & Xie (2023). Scalable Diffusion Models with Transformers
• —Ho et al. (2020). Denoising Diffusion Probabilistic Models
• —van den Oord et al. (2017). Neural Discrete Representation Learning (VQ-VAE)

License

MIT