This project implements a Stable Diffusion model from scratch using PyTorch. The model is capable of generating images from textual descriptions (text-to-image) and transforming existing images based on textual prompts (image-to-image). The implementation is based on the foundational principles of diffusion models and includes features like classifier-free guidance and latent space representations.
- Text-to-Image Generation: Create images from textual descriptions.
- Image-to-Image Transformation: Modify existing images based on textual prompts.
- Inpainting: Fill in missing parts of images using textual descriptions.
- Classifier-Free Guidance: Improved generation quality by balancing conditioned and unconditioned outputs.
- Latent Space Representation: Efficient computation using a variational autoencoder.
- Basic understanding of probability and statistics (e.g., multivariate Gaussian, conditional probability, Bayes' rule).
- Basic knowledge of PyTorch and neural networks.
- Familiarity with attention mechanisms and convolution layers.
- Clone the repository:
git clone https://github.com/AK2k30/Stable_Diffusion_model_from_scratch.git cd Stable_Diffusion_model_from_scratch - Install the required dependencies:
pip install -r requirements.txt
Training
To train the model from scratch, run:
python train.py --config configs/config.yamlCustomize the configuration file (configs/config.yaml) to suit your dataset and training preferences.
Text-to-Image
To generate an image from a text prompt:
python generate.py --mode text-to-image --prompt "A dog with glasses"
Image-to-Image
To transform an existing image based on a text prompt:
python generate.py --mode image-to-image --input_image path/to/image.jpg --prompt "A dog with glasses"
Inpainting
To inpaint a missing part of an image using a text prompt:
python inpaint.py --input_image path/to/image.jpg --prompt "A dog running"
The VAE is responsible for encoding the input images into a latent space.
It consists of two main parts:
-
Encoder: Compresses the input image into a latent vector.
-
Decoder: Reconstructs the image from the latent vector.
The VAE helps in reducing the dimensionality of the data, making it computationally efficient to process through the diffusion model.
The LDM operates on the latent representations obtained from the VAE. It is designed to learn the data distribution in the latent space. The key idea is to progressively denoise a sample from a simple distribution (e.g., Gaussian noise) to match the target data distribution.
The UNet architecture is used as the core neural network within the diffusion model.
It consists of:
-
Encoder Path: A series of convolutional layers that downsample the input.
-
Bottleneck: The middle part of the network that captures the most compressed representation.
-
Decoder Path: A series of convolutional layers that upsample the representation back to the original size.
The UNet allows for efficient handling of high-resolution images by capturing multi-scale features through its symmetric design.
Attention mechanisms are integrated within the UNet to enhance the model's ability to focus on relevant parts of the image during processing.
This includes:
-
Self-Attention: Helps the model to consider dependencies between different parts of the image.
-
Cross-Attention: Facilitates the interaction between the text embeddings and the image features, crucial for text-to-image generation.
Classifier-free guidance improves the generation quality by balancing conditioned (text-prompt) and unconditioned (random noise) outputs. This technique allows the model to generate more coherent and high-quality images by guiding the diffusion process with the text prompt while still considering the general data distribution.
Denoising Diffusion Probabilistic Models by Ho et al.
U-Net: Convolutional Networks for Biomedical Image Segmentation by Ronneberger et al.
