🎨 Stable Diffusion — Image Generation Model

A PyTorch implementation of Stable Diffusion from scratch, including the full pipeline: CLIP text encoder, VAE encoder/decoder, UNet diffusion model, and DDPM scheduler.

---

🖼️ Sample Output

Image generated using the pipeline with a text prompt via the CLIP encoder and DDPM sampler.

---

Download Pretrained Weights

Download the v1-5-pruned-emaonly.ckpt file from Hugging Face:

👉 Click here to download

On that page, click v1-5-pruned-emaonly.ckpt → then click the download icon on the right.

Once downloaded, place it inside the data/ folder:

🧠 Model Architecture

Pipeline overview:

Text Prompt
    │
    ▼
┌─────────────┐
│  CLIP       │  ← Text Encoder (tokenizer + transformer)
│  Encoder    │
└──────┬──────┘
       │  text embeddings
       ▼
┌─────────────┐     ┌─────────────┐
│   VAE       │────▶│   UNet /    │
│  Encoder    │     │  Diffusion  │  ← Denoising over T timesteps
└─────────────┘     └──────┬──────┘
                           │
                    ┌──────▼──────┐
                    │    VAE      │
                    │   Decoder   │
                    └──────┬──────┘
                           │
                    ┌──────▼──────┐
                    │  output.png │
                    └─────────────┘

---

📁 Project Structure

Diffusion/
├── data/
│   ├── vocab.json          # CLIP tokenizer vocabulary
│   └── merges.txt          # BPE merge rules
├── diffusion_model/
│   ├── pipeline.py         # Main inference pipeline
│   ├── clip.py             # CLIP text encoder
│   ├── encoder.py          # VAE encoder
│   ├── decoder.py          # VAE decoder
│   ├── diffusion.py        # UNet diffusion model
│   ├── ddpm.py             # DDPM noise scheduler
│   ├── attention.py        # Attention mechanisms
│   ├── model_loader.py     # Load pretrained weights
│   ├── model_converter.py  # Convert model formats
│   ├── add_noise.py        # Add noise to images (DDPM forward process)
│   ├── add_noise.ipynb     # Interactive noise visualization notebook
│   ├── check.py            # Model sanity checks
│   └── output.png          # Sample generated image
└── README.md

---

⚙️ Requirements

• Python 3.10+
• PyTorch
• transformers
• numpy
• Pillow
• tqdm
• jupyter (for the notebook)

Install all dependencies:

pip install torch torchvision transformers numpy Pillow tqdm jupyter

---

🚀 Getting Started

Step 1 — Download Pretrained Weights

You need the Stable Diffusion v1.5 weights (.ckpt file). Download from Hugging Face:

# Option A: using huggingface_hub
pip install huggingface_hub

python -c "
from huggingface_hub import hf_hub_download
hf_hub_download(
    repo_id='runwayml/stable-diffusion-v1-5',
    filename='v1-5-pruned-emaonly.ckpt',
    local_dir='./data'
)
"

Or manually download from: https://huggingface.co/runwayml/stable-diffusion-v1-5

Place the downloaded .ckpt file inside the data/ folder:

data/
├── v1-5-pruned-emaonly.ckpt   ← place here
├── vocab.json
└── merges.txt

---

Step 2 — Run the Add Noise Script

To visualize the DDPM forward process (adding noise to an image step by step):

Using the Python script:

cd diffusion_model
python add_noise.py

Using the Jupyter Notebook (recommended for visualization):

cd diffusion_model
jupyter notebook add_noise.ipynb

This will show how a clean image gradually becomes pure noise across T timesteps — the forward diffusion process.

---

Step 3 — Generate an Image

Run the full text-to-image pipeline:

from diffusion_model.pipeline import generate
from diffusion_model.model_loader import load_models
from PIL import Image

# Load models
models = load_models("data/v1-5-pruned-emaonly.ckpt")

# Generate
output = generate(
    prompt="a photograph of an astronaut riding a horse",
    uncond_prompt="",
    models=models,
    device="cuda",   # or "cpu"
    num_inference_steps=20,
    cfg_scale=7.5,
    seed=42
)

Image.fromarray(output).save("diffusion_model/output.png")

---

📊 How It Works

Step	Process	File
1	Tokenize text prompt	clip.py
2	Encode tokens to embeddings	clip.py
3	Sample random latent noise	pipeline.py
4	Denoise over T steps (reverse diffusion)	diffusion.py, ddpm.py
5	Decode latents to pixel space	decoder.py
6	Save output image	pipeline.py

---

🔧 Configuration

You can tweak the following in pipeline.py:

num_inference_steps = 20      # More steps = better quality, slower
cfg_scale = 7.5               # Classifier-free guidance scale
sampler = "ddpm"              # Noise scheduler
seed = 42                     # Reproducibility

---

📜 License

This project is for educational purposes. Pretrained weights are subject to the CreativeML Open RAIL-M license.