MNIST Diffusion Model

A diffusion-based generative model for MNIST digits implemented in PyTorch.

Installation

pip install -r requirements.txt

Usage

Train Diffusion Model

python main.py train --epochs 10 --lr 2e-4 --batch-size 64

Train with Self-Attention

python main.py train --epochs 10 --attention

Train Autoencoder (for latent diffusion)

python main.py train-ae --epochs 10

Generate Samples

python main.py sample --checkpoint checkpoints/diffusion_latest.pt

Visualize Diffusion Process

python main.py visualize --all

Configuration

All hyperparameters can be found in src/config.py: - DIFFU_STEPS: Number of diffusion steps (default: 1000) - EPOCHS: Training epochs (default: 10) - BATCH_SIZE: Batch size (default: 64) - LEARNING_RATE: Learning rate (default: 2e-4)

Model Architecture

The diffusion model uses a UNet architecture with: - Timestep embedding - Label conditioning (for class-conditional generation) - Optional self-attention layers

License

See LICENSE for details.