A complete, from-scratch PyTorch implementation of Stable Diffusion v1.5, featuring both text-to-image and image-to-image generation capabilities. This project demonstrates the inner workings of diffusion models by implementing all components without relying on pre-built libraries.

🚀 Features

Text-to-Image Generation: Create high-quality images from text descriptions
Image-to-Image Generation: Transform existing images using text prompts
Complete Implementation: All components built from scratch in PyTorch
Flexible Sampling: Configurable inference steps and CFG scale
Model Compatibility: Support for various fine-tuned Stable Diffusion models
Clean Architecture: Modular design with separate components for each part of the pipeline

🏗️ Architecture

This implementation includes all the core components of Stable Diffusion:

CLIP Text Encoder: Processes text prompts into embeddings
VAE Encoder/Decoder: Handles image compression and reconstruction
U-Net Diffusion Model: Core denoising network with attention mechanisms
DDPM Sampler: Implements the denoising diffusion probabilistic model
Pipeline Orchestration: Coordinates all components for generation

📁 Project Structure

├── main/
│   ├── attention.py      # Multi-head attention implementation
│   ├── clip.py           # CLIP text encoder
│   ├── ddpm.py           # DDPM sampling algorithm
│   ├── decoder.py        # VAE decoder for image reconstruction
│   ├── diffusion.py      # U-Net diffusion model
│   ├── encoder.py        # VAE encoder for image compression
│   ├── model_converter.py # Converts checkpoint files to PyTorch format
│   ├── model_loader.py   # Loads and manages model weights
│   ├── pipeline.py       # Main generation pipeline
│   └── demo.py           # Example usage and demonstration
├── data/                 # Model weights and tokenizer files
└── images/               # Input/output images

🛠️ Installation

Prerequisites

Python 3.8+
PyTorch 1.12+
Transformers library
PIL (Pillow)
NumPy
tqdm

Setup

Clone the repository:

bash
1git clone https://github.com/https://github.com/ApoorvBrooklyn/Stable-Diffusion 2cd pytorch-stable-diffusion
Create virtual environment:

bash
1python -m venv venv 2source venv/bin/activate # On Windows: venv\Scripts\activate
Install dependencies:

bash
1pip install torch torchvision torchaudio 2pip install transformers pillow numpy tqdm
Download required model files:
- Download vocab.json and merges.txt from Stable Diffusion v1.5 tokenizer
- Download v1-5-pruned-emaonly.ckpt from Stable Diffusion v1.5
- Place all files in the data/ folder

🎯 Usage

Basic Text-to-Image Generation

python

import model_loader
import pipeline
from transformers import CLIPTokenizer

# Initialize tokenizer and load models
tokenizer = CLIPTokenizer("data/vocab.json", merges_file="data/merges.txt")
models = model_loader.preload_models_from_standard_weights("data/v1-5-pruned-emaonly.ckpt", "cpu")

# Generate image from text
output_image = pipeline.generate(
    prompt="A beautiful sunset over mountains, highly detailed, 8k resolution",
    uncond_prompt="",  # Negative prompt
    do_cfg=True,
    cfg_scale=8,
    sampler_name="ddpm",
    n_inference_steps=50,
    seed=42,
    models=models,
    device="cpu",
    tokenizer=tokenizer
)

Image-to-Image Generation

python

from PIL import Image

# Load input image
input_image = Image.open("images/input.jpg")

# Generate transformed image
output_image = pipeline.generate(
    prompt="Transform this into a watercolor painting",
    input_image=input_image,
    strength=0.8,  # Controls how much to change the input
    # ... other parameters
)

Advanced Configuration

CFG Scale: Controls how closely the image follows the prompt (1-14)
Inference Steps: More steps = higher quality but slower generation
Strength: For image-to-image, controls transformation intensity (0-1)
Seed: Set for reproducible results

🔧 Model Conversion

The model_converter.py script converts Stable Diffusion checkpoint files to PyTorch format:

python main/model_converter.py --checkpoint_path data/v1-5-pruned-emaonly.ckpt --output_dir converted_models/

🎨 Supported Models

This implementation is compatible with:

Stable Diffusion v1.5: Base model
Fine-tuned Models: Any SD v1.5 compatible checkpoint
Custom Models: Models trained on specific datasets or styles

Tested Fine-tuned Models:

InkPunk Diffusion: Artistic ink-style images
Illustration Diffusion: Hollie Mengert's illustration style

🚀 Performance Tips

Device Selection: Use CUDA for GPU acceleration, MPS for Apple Silicon
Batch Processing: Process multiple prompts simultaneously
Memory Management: Use idle_device="cpu" to free GPU memory
Optimization: Adjust inference steps based on quality vs. speed needs

🔬 Technical Details

Diffusion Process

Implements DDPM (Denoising Diffusion Probabilistic Models)
Uses U-Net architecture with cross-attention for text conditioning
VAE handles 512x512 image compression to 64x64 latents

Attention Mechanisms

Multi-head self-attention in U-Net
Cross-attention between text embeddings and image features
Efficient attention implementation for memory optimization

Sampling

Configurable number of denoising steps
Classifier-free guidance (CFG) for prompt adherence
Deterministic generation with seed control

🤝 Contributing

Contributions are welcome! Please feel free to submit pull requests or open issues for:

Bug fixes
Performance improvements
New sampling algorithms
Additional model support
Documentation improvements

📄 License

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

🙏 Acknowledgments

Stability AI for the original Stable Diffusion model
OpenAI for the CLIP architecture
CompVis for the VAE implementation
Hugging Face for the transformers library

📚 References

📞 Support

If you encounter any issues or have questions:

Open an issue on GitHub
Check the existing documentation
Review the demo code for examples

Note: This is a research and educational implementation. For production use, consider using the official Stable Diffusion implementations or cloud-based APIs.

ApoorvBrooklyn/stable-diffusion-implementation

PyTorch Stable Diffusion Implementation