DDPM Model for Image Generation and Training

This repository is a simple demo implementation of the DDPM (Denoising Diffusion Probabilistic Model) model for image generation and training. Based on the 2020 paper Denoising Diffusion Probabilistic Models

Features

• Diffusion model implementation using a custom UNet model.
• Support conditional (classes specified) and unconditional generation.
• Support for image generation from a pre-trained model.
• Support for training from scratch or continuing from a checkpoint.
• Customizable dataset, image size, batch size, and other training parameters via command-line arguments.

Table of Contents

1. Results
2. Installation
3. Usage
4. Command-line Arguments
5. Examples
6. Project Structure

Results

Celeb dataset: 64x64 images (Unconditional Generation).

Cifar10 dataset: 32x32 images (Conditional Generation).

Fashion mnist: 32x32 images (Conditional Generation).

Fashion mnist: 32x32 images (Unconditional Generation).

Installation

Clone the Repository

git clone https://github.com/EzzatEsam/TestingDiffusionModel.git

Install Dependencies

Make sure you have PyTorch installed. You can install the required Python packages using:

pip install -r requirements.txt

Modify dataset locations

You will need to manually modify the dataset locations in the loader.py file.

Usage

Generating Images

To generate images from a pre-trained model, simply run:

python main.py --img_size 64 --version {version} --saving_path {saving_path} --start_epoch {start_epoch}

This will load a pre-trained model (starting from the specified epoch) and generate a set of images.

Training the Model

If you want to train the model, use the --training flag:

python script.py --training --dataset celeb --img_size 64 --batch_size 64 --start_epoch 0 --epochs 1001 --learning_rate 1e-4 --save_every_n 2

The model will be trained for the specified number of epochs, and checkpoints will be saved at regular intervals.

Command-line Arguments

Argument	Type	Default	Description
--device	str	"cuda" (if available)	Device to run the model on (cuda or cpu).
--training	flag	False	If set, the model will be trained. If not, images will be generated.
--saving_path	str	"/teamspace/studios/this_studio/runs"	Directory where the model checkpoints and images will be saved.
--version	int	0	Version number for saving models and images.
--dataset	str	"celeb"	Dataset to use for training or image generation ("celeb" or "bridge" or "cifar10" or "fashion" or "cifar10").
--img_size	int	64	Size of the images (both height and width).
--batch_size	int	64	Batch size for training and validation.
--start_epoch	int	0	Starting epoch for training, useful for resuming from a checkpoint.
--epochs	int	1001	Maximum number of epochs to train the model.
--save_every_n	int	2	Save model checkpoint every n epochs during training.
--learning_rate	float	1e-4	Learning rate for the optimizer.
--conditional	flag	False	If set, the model will be trained with conditional generation.

Examples

Image Generation (Default Settings)

python script.py  --version {version} --saving_path {saving_path} --start_epoch {start_epoch}

Training the Model

python script.py --training --dataset celeb --img_size 64 --batch_size 64 --start_epoch 33 --epochs 1001 --learning_rate 1e-4 --save_every_n 2

Custom Device (CPU)

python script.py --device cpu --training --dataset celeb --img_size 64 --batch_size 32 --start_epoch 0 --epochs 500

Project Structure

.
├── DiffusionTest/
│   ├── model.py              # Contains the SmallUNetWithEmb model definition
│   ├── diffusion.py          # Contains the diffusion process and generation logic
│   ├── loader.py             # Contains dataset loading logic
│   ├── utils.py              # Utility functions like saving images
│   ├── __init__.py           # Package initialization
├── main.py                 # Main script to run training or image generation
├── requirements.txt          # Python dependencies
├── README.md                 # This README file

Contributing

Contributions are welcome! Please open an issue or submit a pull request for any improvements or bug fixes.

License

This project is licensed under the MIT License. See the LICENSE file for details.