HoumanAHuman/conditional-diffusion-mnist

Conditional Diffusion on Modified MNIST

This project trains a class-conditional diffusion model on a modified MNIST dataset where class 1 includes samples from both MNIST digit 1 and FashionMNIST class 'trouser'.

📦 Setup

git clone https://github.com/yourusername/conditional-diffusion-mnist.git
cd conditional-diffusion-mnist
pip install -r requirements.txt

📄 Dataset

Prepare the custom dataset:

python prepare_dataset.py

This generates Dataset/shuffled_mnist_with_trousers.pt, a balanced dataset of MNIST 1s and FashionMNIST trousers labeled as 1.

🧨 Train the Diffusion Model

python train_diffusion.py

This trains a conditional DDPM model and saves checkpoints under the DDPM/ directory, including unet_final.pt.

🎨 Sample from the Model

python sample_images.py

This loads the trained model from DDPM/unet_final.pt and generates class-conditioned samples (by default, class 1).

📁 Folder Structure

conditional-diffusion-mnist/
├── DDPM/                                                                 # Model checkpoints
│   └── unet_final_ema.pt
│   └── class_embedder.pt
├── Samples/                                                              # Sample generation
│   └── class {condition_class}/                                          # Samples from the original model with respective conditioned class
│   └── class {condition_class}-Unlearned with constant lambda/           # Samples from the unlearned (constant lambda) model with respective conditioned class
│   └── class {condition_class}-Unlearned with dynamic lambda/            # Samples from the unlearned (dynamic lambda) model with respective conditioned class
├── DDPM_Unlearned/                                                       # Model checkpoints after Unlearning with constant lambda
│   └── unet_unlearned_ema.pt
│   └── class_embedder_unlearned_ema.pt
├── DDPM_Unlearned_dynamic/                                               # Model checkpoints after Unlearning with dynamic lambda
│   └── unet_unlearned_dynamic_ema.pt
│   └── class_embedder_unlearned_dynamic_ema.pt
├── Dataset/                                                              # Custom dataset generation and files
│   └── shuffled_augmented_mnist.pt
│   └── original_mnist.pt
│   └── trousers_subset.pt
├── train_diffusion.py                                                    # Training script
├── preprocess_dataset.py                                                 # Dataset creation script
├── requirements.txt
└── README.md

🧪 Requirements

Install all dependencies using:

pip install -r requirements.txt

📝 License

MIT License