# Stable Diffusion — Administrator & Developer Guide

## 1. Architecture Overview

Stable Diffusion implements the Latent Diffusion Model (LDM) architecture, which performs the diffusion process in the latent space of a pretrained autoencoder rather than directly in pixel space. This design dramatically reduces computational cost while preserving high-resolution image quality.

The system consists of three core components working in sequence: a **first-stage autoencoder** compresses images into a compact latent representation, a **UNet-based diffusion model** learns to denoise in that latent space, and a **conditioning encoder** (e.g., CLIP for text) provides cross-attention guidance.

```text
                                ┌──────────────────┐
                                │  Text Prompt      │
                                └────────┬─────────┘
                                         │
                                         ▼
┌──────────┐    encode     ┌──────────────────────────┐    decode     ┌──────────┐
│  Input   │──────────────►│   Latent Space (64×64×4)  │──────────────►│  Output  │
│  Image   │               │                          │               │  Image   │
│ 512×512  │               │  UNet Diffusion Model    │               │ 512×512  │
└──────────┘               │  + Cross-Attention from  │               └──────────┘
      ▲                    │    CLIP Text Embeddings   │                     │
      │                    └──────────────────────────┘                     │
      │                              ▲                                      │
      │                              │                                      │
      └──── AutoencoderKL ───────────┴──── AutoencoderKL.decode ───────────┘
             (Encoder)                       (Decoder)
```

### Component Inventory

| Component | Path | Role |
|---|---|---|
| Training entry point | `main.py` | PyTorch Lightning training loop with argparse CLI |
| Text-to-image inference | `scripts/txt2img.py` | Generate images from text prompts |
| Image-to-image inference | `scripts/img2img.py` | Transform images with text guidance |
| Inpainting inference | `scripts/inpaint.py` | Fill masked image regions |
| Unconditional sampling | `scripts/sample_diffusion.py` | Sample from unconditional models |
| Retrieval-augmented | `scripts/knn2img.py` | KNN-retrieval conditioned generation |
| Retrieval training | `scripts/train_searcher.py` | Train retrieval index |
| Notebook helpers | `notebook_helpers.py` | Colab/Jupyter interactive sampling |
| LDM library | `ldm/` | Core model implementations |
| Model configs | `configs/` | YAML architecture definitions |
| Pre-stage model configs | `models/` | First-stage and LDM model configs |
| Dataset utilities | `data/` | ImageNet synset maps, example images |

### Technology Stack

| Layer | Technology |
|---|---|
| Language | Python 3.8.5 |
| Deep learning | PyTorch 1.11.0, torchvision 0.12.0 |
| Training framework | PyTorch Lightning 1.4.2 |
| Configuration | OmegaConf 2.1.1 (YAML) |
| Text encoder | OpenAI CLIP (frozen), Transformers 4.19.2 |
| Image autoencoder | taming-transformers (VQ-VAE codebook) |
| Image processing | albumentations 0.4.3, OpenCV 4.1.2.30, kornia 0.6 |
| Tensor ops | einops 0.3.0 |
| Metrics | torch-fidelity 0.3.0, torchmetrics 0.6.0 |
| Logging | Test Tube, Weights & Biases (optional) |
| Debugging | pudb 2019.2 |

> [!NOTE]
> The codebase uses a dynamic object instantiation pattern via `instantiate_from_config()` in `ldm/util.py`. All model, dataset, callback, logger, and loss objects are constructed from YAML configs using a `target` key (full Python import path) and optional `params` dict. This is the central extension mechanism — adding a new model or dataset requires only a new class and a config entry.

### Data Flow: Text-to-Image Generation

1. User provides a text prompt
2. `FrozenCLIPEmbedder` encodes text into a sequence of 77×768 context vectors
3. A random noise tensor is sampled in latent space (shape: `[B, 4, 64, 64]`)
4. The UNet iteratively denoises using DDIM/PLMS, conditioned via cross-attention on CLIP embeddings
5. The denoised latent is scaled by `1/scale_factor` (0.18215)
6. `AutoencoderKL.decode()` maps the latent back to pixel space (512×512×3)
7. Output is saved as PNG

### Data Flow: Training

1. Config is loaded via `OmegaConf.load()` and merged with CLI overrides
2. Model instantiated via `instantiate_from_config(config.model)`
3. Data module instantiated via `DataModuleFromConfig` — wraps dataset classes from config
4. Learning rate computed: `lr = accumulate_grad_batches × ngpu × batch_size × base_lr` (if `scale_lr=True`)
5. PyTorch Lightning `Trainer.fit()` runs the training loop
6. Callbacks handle: directory setup (`SetupCallback`), image logging (`ImageLogger`), LR monitoring (`LearningRateMonitor`), CUDA stats (`CUDACallback`), checkpointing (`ModelCheckpoint`)
7. Signal handlers: `SIGUSR1` → save checkpoint, `SIGUSR2` → enter debugger

## 2. System Requirements

| Resource | Inference | Training |
|---|---|---|
| GPU | NVIDIA, 8+ GB VRAM | NVIDIA, 24+ GB VRAM |
| CUDA | 11.3+ | 11.3+ |
| Python | 3.8.5 | 3.8.5 |
| RAM | 16 GB | 32+ GB |
| Disk (code) | ~500 MB | ~500 MB |
| Disk (models) | 4-7 GB per checkpoint | 4-7 GB per checkpoint |
| Disk (datasets) | — | 150+ GB (ImageNet) |
| Network | GitHub access for deps | GitHub + dataset downloads |

### Environment Specification

The complete environment is defined in `environment.yaml`:

```yaml
name: ldm
channels:
  - pytorch
  - defaults
dependencies:
  - python=3.8.5
  - pip=20.3
  - cudatoolkit=11.3
  - pytorch=1.11.0
  - torchvision=0.12.0
  - numpy=1.19.2
  - pip:
    - albumentations==0.4.3
    - diffusers
    - opencv-python==4.1.2.30
    - pudb==2019.2
    - invisible-watermark
    - imageio==2.9.0
    - imageio-ffmpeg==0.4.2
    - pytorch-lightning==1.4.2
    - omegaconf==2.1.1
    - test-tube>=0.7.5
    - streamlit>=0.73.1
    - einops==0.3.0
    - torch-fidelity==0.3.0
    - transformers==4.19.2
    - torchmetrics==0.6.0
    - kornia==0.6
    - -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
    - -e git+https://github.com/openai/CLIP.git@main#egg=clip
    - -e .
```

> [!IMPORTANT]
> The two `-e git+` entries install `taming-transformers` and `clip` in editable mode from their respective GitHub repositories. The final `-e .` installs the `latent-diffusion` package itself. All three require Git to be installed and network access to GitHub during environment creation.

## 3. Configuration Reference

### Config File Inventory

| File | Format | Purpose |
|---|---|---|
| `environment.yaml` | Conda YAML | Python environment and dependencies |
| `configs/stable-diffusion/v1-inference.yaml` | OmegaConf YAML | Stable Diffusion v1 inference config |
| `configs/latent-diffusion/*.yaml` (7 files) | OmegaConf YAML | Latent diffusion model configs |
| `configs/autoencoder/*.yaml` (4 files) | OmegaConf YAML | First-stage autoencoder configs |
| `configs/retrieval-augmented-diffusion/768x768.yaml` | OmegaConf YAML | RAD model config |
| `models/first_stage_models/*/config.yaml` (9 files) | OmegaConf YAML | Pre-trained first-stage model configs |
| `models/ldm/*/config.yaml` (12 files) | OmegaConf YAML | Pre-trained LDM model configs |

### Environment Variables

| Variable | Type | Default | Description | Required? |
|---|---|---|---|---|
| `XDG_CACHE_HOME` | `str` | `~/.cache` | Base cache directory for auto-downloaded datasets (ImageNet) | No |

### CLI Flags (main.py)

| Flag | Type | Default | Description |
|---|---|---|---|
| `-n` / `--name` | `str` | `""` | Postfix for log directory name |
| `-r` / `--resume` | `str` | `""` | Resume from log directory or checkpoint in logdir |
| `-b` / `--base` | `str[]` | `[]` | Paths to base configs (merged left-to-right) |
| `-t` / `--train` | `bool` | `False` | Enable training mode |
| `--no-test` | `bool` | `False` | Disable test evaluation after training |
| `-p` / `--project` | `str` | — | Name or path to project |
| `-d` / `--debug` | `bool` | `False` | Enable post-mortem debugging with pudb/pdb |
| `-s` / `--seed` | `int` | `23` | Seed for `pytorch_lightning.seed_everything()` |
| `-f` / `--postfix` | `str` | `""` | Additional postfix for log directory name |
| `-l` / `--logdir` | `str` | `logs` | Root directory for training logs |
| `--scale_lr` | `bool` | `True` | Scale LR by `ngpu × batch_size × accumulate_grad_batches` |

All PyTorch Lightning `Trainer` arguments are also accepted (e.g., `--gpus`, `--max_epochs`, `--precision`, `--accumulate_grad_batches`).

### Stable Diffusion v1 Inference Config

Complete annotated reference from `configs/stable-diffusion/v1-inference.yaml`:

```yaml
model:
  base_learning_rate: 1.0e-04          # Base LR (scaled if --scale_lr)
  target: ldm.models.diffusion.ddpm.LatentDiffusion
  params:
    linear_start: 0.00085              # β_start for linear noise schedule
    linear_end: 0.0120                 # β_end for linear noise schedule
    num_timesteps_cond: 1              # Number of conditioning timesteps
    log_every_t: 200                   # Log images every N timesteps during sampling
    timesteps: 1000                    # Total diffusion steps T
    first_stage_key: "jpg"             # Dataset key for images
    cond_stage_key: "txt"              # Dataset key for conditioning (text)
    image_size: 64                     # Latent spatial resolution (NOT pixel res)
    channels: 4                        # Latent channels
    cond_stage_trainable: false        # Freeze conditioning encoder
    conditioning_key: crossattn        # Conditioning via cross-attention
    monitor: val/loss_simple_ema       # Metric to monitor for checkpointing
    scale_factor: 0.18215              # Latent scaling (unit variance normalization)
    use_ema: False                     # Disable EMA for inference

    scheduler_config:                  # LR scheduler
      target: ldm.lr_scheduler.LambdaLinearScheduler
      params:
        warm_up_steps: [10000]
        cycle_lengths: [10000000000000]
        f_start: [1.e-6]
        f_max: [1.]
        f_min: [1.]

    unet_config:                       # Denoising UNet
      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
      params:
        image_size: 32                 # Unused (legacy)
        in_channels: 4                 # Must match latent channels
        out_channels: 4                # Must match latent channels
        model_channels: 320            # Base channel count
        attention_resolutions: [4, 2, 1]  # Downsample factors for attention
        num_res_blocks: 2              # ResBlocks per level
        channel_mult: [1, 2, 4, 4]     # Channel multipliers per level
        num_heads: 8                   # Attention heads
        use_spatial_transformer: True   # Enable spatial transformer blocks
        transformer_depth: 1           # Transformer blocks per resolution
        context_dim: 768               # CLIP embedding dimension
        use_checkpoint: True           # Gradient checkpointing (saves VRAM)
        legacy: False

    first_stage_config:                # Autoencoder (VAE)
      target: ldm.models.autoencoder.AutoencoderKL
      params:
        embed_dim: 4
        monitor: val/rec_loss
        ddconfig:
          double_z: true               # Encode mean + logvar
          z_channels: 4                # Latent channels
          resolution: 256              # Input resolution
          in_channels: 3               # RGB input
          out_ch: 3                    # RGB output
          ch: 128                      # Base channels
          ch_mult: [1, 2, 4, 4]        # Channel multipliers
          num_res_blocks: 2
          attn_resolutions: []         # No attention in autoencoder
          dropout: 0.0
        lossconfig:
          target: torch.nn.Identity    # No loss for inference

    cond_stage_config:                 # Text encoder
      target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
```

> [!NOTE]
> The `image_size: 32` in `unet_config.params` is a legacy field that is not actively used — the actual spatial resolution is determined by the latent dimensions passed during forward computation.

### First-Stage Model Configurations

Nine pre-built first-stage model configurations in `models/first_stage_models/`:

| Config | Type | Embed Dim | Z Channels | Resolution | Codebook (VQ only) |
|---|---|---|---|---|---|
| `kl-f4/config.yaml` | AutoencoderKL | 3 | 3 | 256 | — |
| `kl-f8/config.yaml` | AutoencoderKL | 4 | 4 | 256 | — |
| `kl-f16/config.yaml` | AutoencoderKL | 16 | 16 | 256 | — |
| `kl-f32/config.yaml` | AutoencoderKL | 64 | 64 | 256 | — |
| `vq-f4/config.yaml` | VQModelInterface | 3 | 3 | 256 | 8192 |
| `vq-f4-noattn/config.yaml` | VQModelInterface | 3 | 3 | 256 | 8192 |
| `vq-f8/config.yaml` | VQModelInterface | 4 | 4 | 256 | 16384 |
| `vq-f8-n256/config.yaml` | VQModelInterface | 4 | 4 | 256 | 256 |
| `vq-f16/config.yaml` | VQModelInterface | 8 | 8 | 256 | 16384 |

### LDM Model Configurations

Twelve pre-built LDM configs in `models/ldm/`:

| Config | Target Dataset | Conditioning |
|---|---|---|
| `bsr_sr/config.yaml` | Super-resolution (BSR) | Image (concat) |
| `celeba256/config.yaml` | CelebA faces 256 | Unconditional |
| `cin256/config.yaml` | ImageNet 256 | Class label |
| `ffhq256/config.yaml` | FFHQ faces 256 | Unconditional |
| `inpainting_big/config.yaml` | Inpainting | Image + mask (concat) |
| `layout2img-openimages256/config.yaml` | Layout-to-image | Layout |
| `lsun_beds256/config.yaml` | LSUN bedrooms | Unconditional |
| `lsun_churches256/config.yaml` | LSUN churches | Unconditional |
| `semantic_synthesis256/config.yaml` | Semantic synthesis 256 | Segmentation map |
| `semantic_synthesis512/config.yaml` | Semantic synthesis 512 | Segmentation map |
| `text2img256/config.yaml` | Text-to-image 256 | BERT text |

### Autoencoder Training Configurations

Four autoencoder training configs in `configs/autoencoder/`:

| Config | Latent Shape | Loss | Batch Size |
|---|---|---|---|
| `autoencoder_kl_8x8x64.yaml` | 8×8×64 | LPIPS + Disc | 12 |
| `autoencoder_kl_16x16x16.yaml` | 16×16×16 | LPIPS + Disc | 12 |
| `autoencoder_kl_32x32x4.yaml` | 32×32×4 | LPIPS + Disc | 12 |
| `autoencoder_kl_64x64x3.yaml` | 64×64×3 | LPIPS + Disc | 12 |

All autoencoder configs use:
- `disc_start: 50001` (discriminator kicks in after 50K steps)
- `kl_weight: 0.000001` (very small KL penalty)
- `disc_weight: 0.5`
- `accumulate_grad_batches: 2`
- ImageNet super-resolution dataset with `pil_nearest` degradation

> [!IMPORTANT]
> The autoencoder uses a dual-optimizer setup: one for the encoder/decoder/quantizer (`Adam, lr=lr_g, betas=(0.5, 0.9)`) and one for the discriminator (`Adam, lr=lr_d, betas=(0.5, 0.9)`). The `lr_g_factor` parameter scales the generator learning rate relative to the base LR.

## 4. Deployment & Scaling

### Multi-GPU Training

The training script defaults to DDP (Distributed Data Parallel) when GPUs are configured:

```bash
python main.py \
    --base configs/latent-diffusion/lsun_churches-ldm-kl-8.yaml \
    -t --gpus 0,1,2,3
```

When GPUs are specified, the trainer sets `accelerator: "ddp"` automatically. Without GPUs, the accelerator is removed and training falls back to CPU.

> [!WARNING]
> The `CUDACallback` references `trainer.root_gpu`, which may raise `AttributeError` on newer PyTorch Lightning versions. If upgrading Lightning beyond 1.4.2, verify GPU callback compatibility.

### Gradient Accumulation

Increase effective batch size without additional VRAM:

```bash
python main.py --base config.yaml -t --gpus 0, \
    --accumulate_grad_batches 4
```

The effective batch size becomes `batch_size × accumulate_grad_batches × num_gpus`. Learning rate is scaled accordingly when `--scale_lr` is enabled.

### Gradient Checkpointing

The UNet config includes `use_checkpoint: True` by default, enabling PyTorch's gradient checkpointing. This trades compute for memory — roughly halving VRAM usage during training at the cost of ~30% slower backward passes.

### Log Directory Structure

Training creates the following directory layout:

```text
logs/
└── {timestamp}_{name}{postfix}/
    ├── checkpoints/
    │   ├── last.ckpt
    │   ├── {epoch:06}.ckpt
    │   └── trainstep_checkpoints/   (if metrics_over_trainsteps_checkpoint enabled)
    ├── configs/
    │   ├── {timestamp}-project.yaml
    │   └── {timestamp}-lightning.yaml
    └── images/
        ├── train/
        └── val/
```

> [!TIP]
> To enable checkpointing every N training steps (in addition to per-epoch), add `metrics_over_trainsteps_checkpoint` to the `lightning.callbacks` section of your config. Note: this saves weight-only checkpoints every 10,000 steps by default, which can consume significant disk space.

## 5. Operations

### Checkpoint Management

**Automatic checkpointing:** `ModelCheckpoint` saves the best model (by `monitor` metric, default `val/loss_simple_ema`) and `last.ckpt` automatically.

**Manual checkpoint via signal:**
```bash
# Save immediately without stopping training
kill -USR1 $(pgrep -f "python main.py")
```

**Resume from checkpoint:**
```bash
# From a directory (loads last.ckpt + saved configs)
python main.py -r logs/2024-01-15T12-30-00_experiment -t --gpus 0,

# From a specific checkpoint file
python main.py --resume_from_checkpoint logs/.../checkpoints/epoch=5.ckpt \
    --base configs/stable-diffusion/v1-inference.yaml -t -n resumed_run --gpus 0,
```

> [!NOTE]
> When resuming via `-r`, saved config files from the checkpoint's `configs/` directory are automatically prepended to the config chain. The `-n`/`--name` flag cannot be used with `-r`; use `--resume_from_checkpoint` for a fresh run starting from weights.

### Image Logging

The `ImageLogger` callback saves sample images during training:

- **Default frequency:** Every 750 batches (`batch_frequency: 750`)
- **Max images per log:** 4 (`max_images: 4`)
- **Warmup schedule:** Logarithmic steps (1, 2, 4, 8, ..., 750) for early logging, then every 750
- **Output location:** `{logdir}/images/{train|val}/`
- **Filename format:** `{key}_gs-{global_step:06d}_e-{epoch:06d}_b-{batch:06d}.png`

> [!TIP]
> Disable image logging for faster training by setting `disabled: True` in the image_logger callback config, or by reducing `max_images` to 0.

### Monitoring

**Built-in logging:** TestTubeLogger is the default (TensorBoard-compatible). Switch to Weights & Biases by overriding the logger config:

```yaml
lightning:
  logger:
    target: pytorch_lightning.loggers.WandbLogger
    params:
      name: my_experiment
      save_dir: logs
      offline: false
```

**CUDA memory tracking:** The `CUDACallback` reports peak GPU memory and epoch duration:

```text
Average Epoch time: 1234.56 seconds
Average Peak memory 22150.00MiB
```

### Debug Mode

Enable with `-d` / `--debug`:
1. W&B logger runs in offline mode
2. On exception: `pudb` (or `pdb` fallback) post-mortem debugging
3. After completion: log directory moved to `debug_runs/`

## 6. Performance

### Memory Optimization

| Technique | Impact | Config |
|---|---|---|
| Gradient checkpointing | ~50% VRAM reduction | `unet_config.params.use_checkpoint: True` |
| Reduce batch size | Linear VRAM reduction | `data.params.batch_size` |
| Gradient accumulation | Same effective batch, less VRAM | `--accumulate_grad_batches N` |
| Mixed precision | ~40% VRAM reduction | `--precision 16` |
| EMA disable | Slight VRAM saving | `model.params.use_ema: False` |

### Sampling Performance

| Sampler | Typical Steps | Relative Speed | Quality |
|---|---|---|---|
| DDIM | 50-250 | Baseline | Good |
| PLMS | 20-50 | 2-5x faster | Good-Excellent |
| DPM-Solver | 10-25 | 5-10x faster | Excellent |

## 7. Security

### Model Supply Chain

The codebase loads model weights via `torch.load()` with `map_location="cpu"`, which uses Python's `pickle` module for deserialization.

> [!CAUTION]
> `torch.load()` uses `pickle` under the hood, which can execute arbitrary code during deserialization. Only load checkpoints from trusted sources. Malicious `.ckpt` files can compromise the host system. This applies to both model weights and any state_dict loaded via `init_from_ckpt()`.

### External Downloads

The following components download files at runtime:

| Component | URL Pattern | When |
|---|---|---|
| `ImageNetBase` | `heibox.uni-heidelberg.de/f/...` | First dataset use (synset labels) |
| `notebook_helpers.download_models()` | `heibox.uni-heidelberg.de/f/...` | Colab notebook initialization |
| `taming-transformers` | Various | Model weights for VQ-VAE |
| `CLIP` | Hugging Face / OpenAI | Pre-trained CLIP weights |

### Dependency Provenance

Two dependencies are installed from Git repositories in editable mode:

| Package | Source | Risk |
|---|---|---|
| `taming-transformers` | `github.com/CompVis/taming-transformers@master` | Mutable branch reference |
| `clip` | `github.com/openai/CLIP@main` | Mutable branch reference |

> [!WARNING]
> Both Git dependencies reference mutable branches (`master`, `main`). A compromised or updated repository could inject unexpected code. For production deployments, pin to specific commit hashes instead of branch names.

### Input Validation

The codebase performs minimal input validation:
- Boolean arguments use a custom `str2bool` parser that raises `ArgumentTypeError` on invalid values
- Checkpoint paths are validated for existence before loading
- No sanitization of file paths constructed from user inputs in dataset classes

## 8. Developer Internals

### Repository Structure

```text
stable-diffusion/
├── main.py                          # Training entry point (argparse CLI)
├── setup.py                         # Package setup (latent-diffusion 0.0.1)
├── environment.yaml                 # Conda environment specification
├── notebook_helpers.py              # Colab/Jupyter interactive helpers
├── Stable_Diffusion_v1_Model_Card.md
├── LICENSE
├── configs/
│   ├── stable-diffusion/
│   │   └── v1-inference.yaml        # SD v1 inference config
│   ├── latent-diffusion/
│   │   ├── celebahq-ldm-vq-4.yaml
│   │   ├── cin-ldm-vq-f8.yaml
│   │   ├── cin256-v2.yaml
│   │   ├── ffhq-ldm-vq-4.yaml
│   │   ├── lsun_bedrooms-ldm-vq-4.yaml
│   │   ├── lsun_churches-ldm-kl-8.yaml
│   │   └── txt2img-1p4B-eval.yaml
│   ├── autoencoder/
│   │   ├── autoencoder_kl_8x8x64.yaml
│   │   ├── autoencoder_kl_16x16x16.yaml
│   │   ├── autoencoder_kl_32x32x4.yaml
│   │   └── autoencoder_kl_64x64x3.yaml
│   └── retrieval-augmented-diffusion/
│       └── 768x768.yaml
├── ldm/
│   ├── lr_scheduler.py              # Custom LR schedulers
│   ├── util.py                      # Core utilities (instantiate_from_config)
│   ├── data/
│   │   ├── base.py                  # Txt2ImgIterableBaseDataset
│   │   ├── imagenet.py              # ImageNet dataset classes
│   │   └── lsun.py                  # LSUN dataset classes
│   ├── models/
│   │   ├── autoencoder.py           # VQModel, AutoencoderKL, IdentityFirstStage
│   │   └── diffusion/
│   │       ├── ddpm.py              # LatentDiffusion (DDPM-based)
│   │       ├── ddim.py              # DDIM sampler
│   │       ├── plms.py              # PLMS sampler
│   │       ├── classifier.py        # Classifier guidance
│   │       └── dpm_solver/          # DPM-Solver sampler
│   └── modules/
│       ├── attention.py             # Cross/self-attention, SpatialTransformer
│       ├── ema.py                   # Exponential Moving Average (LitEma)
│       ├── x_transformer.py         # Full transformer implementation
│       ├── diffusionmodules/
│       │   ├── model.py             # Encoder/Decoder for autoencoders
│       │   ├── openaimodel.py       # UNetModel (OpenAI-style)
│       │   └── util.py              # Timestep embedding, checkpointing
│       ├── encoders/
│       │   └── modules.py           # FrozenCLIPEmbedder, BERTEmbedder, ClassEmbedder
│       ├── distributions/
│       │   └── distributions.py     # DiagonalGaussianDistribution
│       ├── losses/
│       │   ├── contperceptual.py    # LPIPS perceptual loss
│       │   └── vqperceptual.py      # VQ perceptual loss with discriminator
│       └── image_degradation/
│           ├── bsrgan.py            # BSRGAN degradation
│           ├── bsrgan_light.py      # BSRGAN-light degradation
│           └── utils_image.py       # Image processing utilities
├── models/                          # Pre-trained model configs
│   ├── first_stage_models/          # 9 autoencoder configs (kl-f4..f32, vq-f4..f16)
│   └── ldm/                         # 12 LDM model configs
├── scripts/
│   ├── txt2img.py                   # Text-to-image inference
│   ├── img2img.py                   # Image-to-image translation
│   ├── inpaint.py                   # Inpainting
│   ├── knn2img.py                   # Retrieval-augmented generation
│   ├── sample_diffusion.py          # Unconditional sampling
│   ├── train_searcher.py            # Retrieval index training
│   ├── download_first_stages.sh     # Download first-stage checkpoints
│   ├── download_models.sh           # Download model checkpoints
│   └── tests/
│       └── test_watermark.py        # Watermark test
└── data/
    ├── DejaVuSans.ttf               # Font for text-as-image rendering
    ├── index_synset.yaml            # ImageNet synset-to-index mapping (1000 classes)
    ├── imagenet_clsidx_to_label.txt # ImageNet class labels
    ├── imagenet_train_hr_indices.p   # Pickle: HR training image indices
    ├── imagenet_val_hr_indices.p     # Pickle: HR validation image indices
    ├── example_conditioning/         # Example conditioning inputs
    │   ├── superresolution/
    │   └── text_conditional/
    └── inpainting_examples/          # 8 image+mask pairs for inpainting demo
```

### Module Reference

#### `ldm/util.py` — Core Utilities

The most critical function in the codebase:

**`instantiate_from_config(config)`** — Universal factory function. Takes a dict with `target` (Python import path) and optional `params`, imports the class, and instantiates it. Special sentinel values: `__is_first_stage__` and `__is_unconditional__` return `None` (used in configs to signal that no first-stage/cond-stage is needed).

**`get_obj_from_str(string, reload=False)`** — Resolves a dotted import path (e.g., `ldm.models.diffusion.ddpm.LatentDiffusion`) to the actual class object.

Other utilities:
- `log_txt_as_img(wh, xc, size)` — Renders text captions as images using `DejaVuSans.ttf`
- `ismap(x)` — Checks if tensor is a spatial map (4D, >3 channels)
- `isimage(x)` — Checks if tensor is an image (4D, 1 or 3 channels)
- `count_params(model, verbose)` — Counts total model parameters
- `parallel_data_prefetch(func, data, n_proc, ...)` — Parallel data loading via multiprocessing or threading

#### `ldm/lr_scheduler.py` — Learning Rate Schedulers

Three scheduler classes, all implementing `__call__(n)` → lr multiplier:

| Class | Schedule | Supports Cycles |
|---|---|---|
| `LambdaWarmUpCosineScheduler` | Linear warmup → cosine decay | No |
| `LambdaWarmUpCosineScheduler2` | Linear warmup → cosine decay | Yes (list-based) |
| `LambdaLinearScheduler` | Linear warmup → linear decay | Yes (inherits Scheduler2) |

All schedulers use a base LR of 1.0 and return a multiplier. The SD v1 config uses `LambdaLinearScheduler` with 10,000 warmup steps and an effectively infinite cycle length.

#### `ldm/modules/ema.py` — Exponential Moving Average

`LitEma(model, decay=0.9999)` — Maintains exponential moving average of model parameters as PyTorch buffers. Key methods:

- `forward(model)` — Update shadow params: `shadow = decay * shadow + (1-decay) * param`
- `copy_to(model)` — Copy EMA weights into the model
- `store(parameters)` — Save current params for later restoration
- `restore(parameters)` — Restore previously stored params

The decay rate adapts during early training: `decay = min(decay, (1 + num_updates) / (10 + num_updates))`.

#### `ldm/modules/attention.py` — Attention Mechanisms

| Class | Purpose |
|---|---|
| `CrossAttention` | Multi-head cross-attention (self-attention when `context=None`) |
| `LinearAttention` | O(n) linear attention approximation |
| `SpatialSelfAttention` | 2D spatial self-attention with Conv2d projections |
| `BasicTransformerBlock` | Self-attn → Cross-attn → FFN with residual connections |
| `SpatialTransformer` | Projects image to sequence, applies transformer, projects back |
| `FeedForward` | GELU or GEGLU feed-forward network |
| `GEGLU` | Gated linear unit with GELU activation |

The `SpatialTransformer` is the key integration point: it sits inside the UNet and enables text conditioning via cross-attention. When no context is provided, cross-attention defaults to self-attention.

#### `ldm/modules/x_transformer.py` — Full Transformer

A complete transformer implementation adapted from `lucidrains/x-transformers`:

| Class | Purpose |
|---|---|
| `Attention` | Configurable attention (causal, talking heads, sparse topk, memory KV) |
| `AttentionLayers` | Stacked attention/FF layers with configurable ordering |
| `Encoder` | Non-causal `AttentionLayers` wrapper |
| `TransformerWrapper` | Token embedding + positional embedding + `AttentionLayers` + output projection |

Features: ScaleNorm, RMSNorm, ReZero, GRU gating, sandwich architectures, PAR (parallel attention/FF), memory tokens, and fixed/absolute positional embeddings.

#### `ldm/models/autoencoder.py` — First-Stage Models

| Class | Parent | Purpose |
|---|---|---|
| `VQModel` | `pl.LightningModule` | Vector-quantized autoencoder with codebook |
| `VQModelInterface` | `VQModel` | Interface variant (returns pre-quantized latents from `encode()`) |
| `AutoencoderKL` | `pl.LightningModule` | KL-regularized continuous autoencoder |
| `IdentityFirstStage` | `torch.nn.Module` | Pass-through (no encoding/decoding) |

Both `VQModel` and `AutoencoderKL` include:
- Dual optimizer setup (autoencoder + discriminator)
- EMA weight management via `LitEma`
- `log_images()` for training visualization
- `init_from_ckpt()` for weight loading with selective key ignoring

The `VQModelInterface` differs from `VQModel` in that `encode()` returns the pre-quantization tensor (before the codebook lookup), enabling external code to control quantization.

#### `ldm/data/` — Dataset Classes

**Base classes:**
- `Txt2ImgIterableBaseDataset(IterableDataset)` — Interface for text-image iterable datasets

**ImageNet:**
- `ImageNetBase(Dataset)` — Abstract base with auto-downloading, synset mapping, and class label resolution
- `ImageNetTrain(ImageNetBase)` — ILSVRC2012 training set (1,281,167 images, auto-downloads via Academic Torrents)
- `ImageNetValidation(ImageNetBase)` — ILSVRC2012 validation set (50,000 images)
- `ImageNetSR(Dataset)` — Super-resolution wrapper (crops → resizes → degrades)
- `ImageNetSRTrain/ImageNetSRValidation` — SR train/val splits using pre-selected HR indices from pickle files

**LSUN:**
- `LSUNBase(Dataset)` — Loads images from file lists with configurable interpolation and horizontal flip
- `LSUNChurchesTrain/Validation` — Church outdoor scenes
- `LSUNBedroomsTrain/Validation` — Bedroom scenes
- `LSUNCatsTrain/Validation` — Cat images

#### `main.py` — Training Callbacks

| Callback | Purpose |
|---|---|
| `SetupCallback` | Creates log/checkpoint/config directories, saves configs at training start |
| `ImageLogger` | Periodically saves sample images (train and validation) to disk and logger |
| `CUDACallback` | Tracks peak GPU memory and epoch time on rank 0 |
| `DataModuleFromConfig` | Instantiates datasets from YAML config, creates DataLoaders |

`DataModuleFromConfig` dynamically creates train/val/test/predict dataloaders based on config entries. It uses `num_workers = num_workers or batch_size * 2` and applies `worker_init_fn` for iterable datasets to ensure proper data sharding across workers.

### Code Patterns and Conventions

**Dynamic instantiation:** The `target` + `params` pattern is used everywhere. Any Python class can be plugged in via config:

```yaml
target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
params:
  version: "openai/clip-vit-large-patch14"
```

resolves to:

```python
FrozenCLIPEmbedder(version="openai/clip-vit-large-patch14")
```

**Config merging:** Configs are loaded left-to-right via `OmegaConf.merge(*configs, cli)`. CLI overrides (dot notation) take highest precedence. Lightning config is popped from the main config separately.

**Conditional compilation:** The code adapts to PyTorch Lightning version at runtime via `version.parse(pl.__version__)`. Key breakpoint: PL 1.4.0 changed `checkpoint_callback` handling.

**Signal-based control:** POSIX signals for training control (SIGUSR1 → checkpoint, SIGUSR2 → debugger). These are registered after the model and data are initialized.

**EMA scope:** The `ema_scope()` context manager (on VQModel/AutoencoderKL) temporarily swaps in EMA weights for inference, then restores training weights:

```python
with model.ema_scope("Plotting"):
    samples = model.sample(...)  # Uses EMA weights
# Back to training weights
```

### Testing

The repository contains minimal test infrastructure:

| Test File | Purpose |
|---|---|
| `scripts/tests/test_watermark.py` | Tests invisible watermark functionality |

> [!NOTE]
> Comprehensive unit tests are not included in this research codebase. The primary quality signal is the published paper results and pre-trained model outputs. Contributors should add regression tests for any modifications to the sampling pipeline.

### Error Messages Reference

| Error | Source | Cause |
|---|---|---|
| `Boolean value expected.` | `main.py:33` | Invalid value for `str2bool` argument |
| `-n/--name and -r/--resume cannot be specified both.` | `main.py:472` | Conflicting CLI flags |
| `Cannot find {path}` | `main.py:479` | Resume path does not exist |
| `Expected key 'target' to instantiate.` | `ldm/util.py:84` | Config dict missing `target` key |
| `list expected but function got ndarray.` | `ldm/util.py:116` | Type mismatch in `parallel_data_prefetch` |
| `Decay must be between 0 and 1` | `ldm/modules/ema.py:9` | Invalid EMA decay rate |
| `Check out entmax activation` | `ldm/modules/x_transformer.py:232` | `use_entmax15=True` not implemented |
| `invalid layer type {type}` | `ldm/modules/x_transformer.py:465` | Unknown layer type in `AttentionLayers` |
| `cond not implemented for mode{mode}` | `notebook_helpers.py:82` | Unsupported notebook mode |

### Potential Improvements

Based on the codebase analysis:

1. **Dependency pinning:** The two editable Git installs reference mutable branches. Pin to commit hashes for reproducibility.
2. **Test coverage:** Only one test file exists (`test_watermark.py`). Critical sampling and training code paths lack automated tests.
3. **PyTorch Lightning compatibility:** Several patterns (e.g., `trainer.root_gpu`, `TestTubeLogger`) are deprecated in newer PL versions. The code is pinned to PL 1.4.2.
4. **Checkpoint loading security:** `torch.load()` uses `pickle` which allows arbitrary code execution. Consider `weights_only=True` for untrusted checkpoints (requires PyTorch 1.13+).
5. **Error handling:** The main training loop uses a bare `except Exception` that can mask important errors in debug mode.
6. **Data pipeline:** Dataset classes download files without checksum verification (only file size checks). Academic Torrents hashes are checked but HTTP downloads are not.

### Contributing

The repository follows these conventions:
- Python 3.8 compatibility required
- Dependencies managed via Conda (`environment.yaml`)
- Configuration via OmegaConf YAML files
- All model components must implement the `target`/`params` instantiation pattern
- Training uses PyTorch Lightning `LightningModule` and `LightningDataModule`
- Image values normalized to `[-1, 1]` range throughout the pipeline
- Checkpoints saved as PyTorch state dicts via PL's `ModelCheckpoint`

## 9. Integration Points

### Extending with Custom Models

Add a new diffusion model by creating a class that inherits from `LatentDiffusion` (in `ldm/models/diffusion/ddpm.py`) and referencing it in a YAML config:

```yaml
model:
  target: my_module.MyCustomDiffusion
  params:
    # ... custom params
```

### Adding Custom Datasets

Implement a PyTorch `Dataset` or `IterableDataset`:

```python
class MyDataset(torch.utils.data.Dataset):
    def __init__(self, data_root, size=256):
        self.images = glob.glob(f"{data_root}/*.png")
        self.size = size

    def __len__(self):
        return len(self.images)

    def __getitem__(self, i):
        image = Image.open(self.images[i]).resize((self.size, self.size))
        image = np.array(image).astype(np.float32) / 127.5 - 1.0
        return {"image": image}
```

Reference in config:

```yaml
data:
  target: main.DataModuleFromConfig
  params:
    batch_size: 4
    train:
      target: my_module.MyDataset
      params:
        data_root: /path/to/images
        size: 256
```

### Adding Custom Samplers

Implement a sampler class with a `sample()` method compatible with `LatentDiffusion`. The existing samplers (`DDIMSampler`, `PLMSSampler`, `DPMSolverSampler`) provide the interface pattern.

### Downstream Ecosystems

This codebase is the foundation for:
- **Stability AI's Stable Diffusion** releases (v1.x-v2.x, different conditioning/resolution)
- **Hugging Face `diffusers`** library (reimplemented with a different architecture)
- **Automatic1111's webUI** (wraps these scripts with a Gradio frontend)
- **ComfyUI** (node-based workflow using these model definitions)
- **InvokeAI** (full-featured generator built on the LDM architecture)

---

_Generated May 2026 from CompVis/stable-diffusion at 21f890f. AI-generated documentation. The library is checked by sampling (three repositories so far); this guide was not individually reviewed. Verify against the project before relying on it. Unofficial; not affiliated with or reviewed by the project._