Latent Diffusion & Stable Diffusion

Latent diffusion is diffusion run on the compressed latent code of a VAE instead of on raw pixels. The reason to do it is cost. Pixel-space diffusion at 512² means the U-Net runs on tensors of shape [B, 3, 512, 512]. Each sampling step is ~100 GFLOPS for a 500M-param U-Net, so fifty steps is 5 TFLOPS per image. Train on a billion images and the compute bill is absurd.

Most of those FLOPs go to pushing perceptually unimportant detail through the net, like the high-frequency texture a lossy VAE could compress away. Rombach's idea: train a VAE once (the first stage), freeze it, and run diffusion in the 4-channel 64×64 latent space (the second stage). Same U-Net, but 8× smaller along each axis (64×64 positions instead of 512×512, so 64× fewer spatial positions) and 4 latent channels instead of 3 pixel channels (~48× fewer values in total). Per-step compute falls by roughly the position-count factor, ~64×, for comparable quality.

That is the Stable Diffusion recipe. SD 1.x / 2.x used an 860M U-Net over 64×64×4 latents, SDXL used a 2.6B U-Net over 128×128×4, and SD3 swapped the U-Net for a Diffusion Transformer (DiT) with flow matching. Flux.1-dev (Black Forest Labs, 2024) ships a 12B-param DiT-MMDiT. All of them run on the same two-stage substrate.

Latent diffusion splits generation into two stages, trained separately.

1. Stage 1: VAE. Encoder E(x) → z, decoder D(z) → x. The target is 8× downsampling along each spatial axis, with a small channel count (4 for SD 1.x), so a 64×64×4 latent holds ~1/48 as many values as a 512×512×3 image. Loss = reconstruction (L1 + LPIPS perceptual) + KL. The KL weight is small so z isn't forced too close to Gaussian, since we don't need exact sampling from z. An adversarial loss is often added so decoded images come out sharp.

2. Stage 2: diffusion on z. Treat z = E(x_real) as the data and train a U-Net (or DiT) to denoise z_t. At inference, sample z_0 with diffusion, then decode x = D(z_0).

Text conditioning adds two components. A frozen text encoder produces the tokens (CLIP-L for SD 1.x, CLIP-L + OpenCLIP-G for SD 2/XL, T5-XXL for SD3 and Flux). Cross-attention injects them: each U-Net block takes [Q = image features, K = V = text tokens] and mixes them in. Text reaches the image through those tokens.

The loss function is identical to Lesson 06. Same DDPM / flow matching MSE on noise, with a different data domain underneath.

The trend across these models: U-Net gives way to DiT (a transformer over latent patches), the text encoder gets bigger (T5 beats CLIP for prompt adherence), and latent channels go up (4 → 16 gives more headroom for detail).

code/main.py stacks a toy 1-D first stage (a fixed rescale, encode(x) = 0.5x and decode(z) = 2z, standing in for a real conv-net VAE) under a small diffusion model and adds class conditioning with classifier-free guidance. The denoiser is a 2-hidden-layer tanh MLP with a hand-written backward pass, and the script carries its own linear beta schedule (T=40), so it needs no code from Lesson 06. It shows the key insight: the same diffusion loss works on raw 1-D values and on encoded values.

Step 1: encoder/decoder

def encode(x):    return x * 0.5          # toy "compression" to smaller scale
def decode(z):    return z * 2.0

A real VAE has trained weights. Here a linear map is enough to show that diffusion operates on z without caring about the original data space.

Step 2: diffusion in z-space

Same noise-prediction objective as Lesson 06, just with the script's own schedule and MLP in place of a U-Net. The data the net sees is z = E(x). After sampling z_0, decode with D(z_0).

Step 3: classifier-free guidance

Classifier-free guidance blends a conditional and an unconditional prediction. During training, drop the class label 10% of the time (replace with a null token). At inference, compute both ε_cond and ε_uncond, then:

eps_cfg = (1 + w) * eps_cond - w * eps_uncond

w = 0 is no guidance (the plain conditional prediction). This is the (1+w) form from Ho & Salimans, and it is what main.py computes. Many libraries write the same blend as eps_uncond + s * (eps_cond - eps_uncond) and call s the guidance scale. Expanding shows the two are identical when s = 1 + w: eps_uncond + (1 + w) * (eps_cond - eps_uncond) = (1 + w) * eps_cond - w * eps_uncond. So s = 1 is no guidance, and a diffusers guidance_scale of 7.5 is w = 6.5 here. The w values in this lesson (the sweep, the 3-7 range, the w > 10 warning) use the (1+w) convention. As a production rule of thumb, w around 3 is a common default and w of 7 or more starts to look saturated and over-sharp. The toy script does not measure either claim: it sweeps w in {0, 1, 3, 7} and prints the per-class mean over 200 samples at each setting.

Step 4: text conditioning (concept, not code)

Replace the class label with a frozen text encoder output. Feed the text embedding to the U-Net via cross-attention:

h = h + CrossAttention(Q=h, K=text_embed, V=text_embed)

Swapping the class label for a text embedding and adding this attention step is what separates a class-conditional diffusion model from Stable Diffusion.

• VAE-scale mismatch. SD 1.x VAEs have a scaling constant (scaling_factor ≈ 0.18215) applied after encoding (you multiply the latents by it, which is the same as dividing by ~5.5). Forget it and the U-Net trains on latents with wildly wrong variance. Each VAE family has its own value (SDXL's differs from SD 1.x's), so read it from the checkpoint's config.
• Text encoder silently wrong. SD3 needs T5-XXL with >=128 tokens, and the fallback to CLIP-only is lossy. Check use_t5=True, or prompt fidelity craters.
• Mixing latent spaces. SDXL, SD3, and Flux use different VAEs, so a LoRA trained on SDXL latents will not work on SD3. Hugging Face diffusers 0.30+ refuses to load mismatched checkpoints.
• CFG too high. w > 10 produces saturated, oily images and over-fits the prompt at the cost of diversity. The sweet spot is w = 3-7.
• Negative prompts leaking. An empty negative prompt becomes the null token, while a filled one becomes ε_uncond. These are different things, and some pipelines silently default to the null.

These are the production stacks in 2026, by target:

Save outputs/skill-sd-prompter.md. The skill takes a text prompt and a target style, and outputs a model + checkpoint, CFG scale, sampler, negative prompt, resolution, an optional ControlNet/IP-Adapter combo, and a per-step QA checklist.

1. Easy. Run code/main.py with the guidance sweep w ∈ {0, 1, 3, 7} (the script's own list). Compare the printed per-class means against the real data means of -2 and +2. At what w does a class mean first move past its real mean? Then add 15 to the list and see whether the trend continues.
2. Medium. Swap the toy linear encoder for a tanh-MLP encoder/decoder pair with a reconstruction loss. Retrain diffusion on the new latents. Does sample quality change?
3. Hard. Set up a real Stable Diffusion inference with diffusers: load sdxl-base, run 30 Euler steps with CFG=7, time it. Now switch to sdxl-turbo with 4 steps and CFG=0. Same subject, different quality. Describe what changed and why.

The reference Flux integration is the canonical "I have a consumer GPU, can I ship this?" recipe. It applies the same three-knob recipe the production inference literature lists to a diffusion DiT:

1. Staggered loading. Flux has four networks that do not need to sit in VRAM at the same time: the T5-XXL text encoder (~10 GB in fp32), CLIP-L (small), the 12B MMDiT, and the VAE. Encode the prompt first, delete the encoders, load the DiT, denoise, delete the DiT, load the VAE, decode. A consumer 8GB GPU fits one stage at a time.
2. 4-bit quantization via bitsandbytes. Use BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16) on both the T5 encoder and the DiT. It cuts memory 8×, and the quality drop is imperceptible for text-to-image per Aritra's benchmarks (linked in the notebook).
3. CPU offload. pipe.enable_model_cpu_offload() swaps modules between CPU and GPU as each forward pass advances. It adds 10-20% latency but lets the pipeline run at all.

The memory accounting: 10 GB T5 / 8 = 1.25 GB quantized, 12 B params × 0.5 bytes = ~6 GB for the quantized DiT, plus activations. In stas00's terms this is the extreme end of TP=1 inference, with no model parallelism and maximum quantization. For production you'd run TP=2 or TP=4 on H100s. For a single dev laptop, this is the recipe.

• Rombach et al. (2022). High-Resolution Image Synthesis with Latent Diffusion Models: Stable Diffusion.
• Podell et al. (2023). SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis: SDXL.
• Peebles & Xie (2023). Scalable Diffusion Models with Transformers (DiT): DiT.
• Esser et al. (2024). Scaling Rectified Flow Transformers for High-Resolution Image Synthesis: SD3, MMDiT.
• Ho & Salimans (2022). Classifier-Free Diffusion Guidance: CFG.
• Labs (2024). Flux.1 — Black Forest Labs announcement: Flux.1 family.
• Hugging Face Diffusers docs: reference implementation for every checkpoint above.