Latent Diffusion Models - The Architecture Behind Stable Diffusion

:::note Reading time: ~55 minutes | Interview relevance: Very High | Target roles: ML Engineer, Research Engineer, AI Engineer, Applied Scientist :::

---

The Real Interview Moment

The interviewer slides a diagram across the table. It shows three components connected by arrows: an encoder $\mathcal{E}$, a U-Net $\varepsilon_\theta$, and a decoder $\mathcal{D}$. "Walk me through the complete Stable Diffusion inference pipeline. What does each component do? What dimensionality does each operate at? Why does this three-part design produce better samples than running diffusion directly on pixels, and not just cheaper samples?"

This question has layers. Most engineers can describe the pipeline at a high level. Fewer can explain:

Why the VAE uses both a perceptual loss (LPIPS) and a PatchGAN discriminator - and why a plain L2 reconstruction loss would fail.

Why 8x spatial compression works but 16x does not - what information is preserved and what is lost.

How cross-attention transforms a sequence of 77 text tokens into spatial conditioning that routes different words to different image regions.

Why SDXL added a second text encoder, what OpenCLIP-bigG adds that CLIP-L cannot provide, and why aesthetic conditioning on image dimensions matters.

Why the VAE has a scaling factor of 0.18215 and what happens if you forget it.

Robin Rombach and colleagues at Ludwig Maximilian University published Latent Diffusion Models (LDMs) in 2022. The paper introduced a single organizing insight: separate perceptual compression (removing imperceptible pixel-level redundancy) from semantic compression (learning the image distribution). Do each with the best tool. The result was Stable Diffusion - open-source, runs on consumer GPUs, and launched the open generative AI ecosystem.

---

Why This Exists - The Pixel-Space Bottleneck

Running diffusion at full resolution is prohibitively expensive. Consider a 512x512 RGB image: the denoising network operates on a $512 \times 512 \times 3 = 786{,}432$-dimensional input at every denoising step. The self-attention layers in the U-Net scale quadratically with spatial dimension.

At a 64x64 internal feature map (reached after 3 downsampling operations), the attention matrix has $4{,}096 \times 4{,}096 = 16{,}777{,}216$ elements per head. This is manageable but expensive. At the 32x32 feature map, it is $1{,}048{,}576$ per head - fine. But computing this at 512x512 input before any downsampling is completely infeasible.

The ADM model (Dhariwal and Nichol 2021) trained at 256x256 on ImageNet using a 554M parameter U-Net. Extending to 512x512 would have required roughly 4x more compute per step. At 1000 steps, training became financially prohibitive for academic labs.

But here is the deeper issue: most information in a natural image is perceptually redundant. Adjacent pixels are correlated. An image can be reconstructed from a much lower-dimensional representation with no perceptual loss, if the encoding is learned rather than hand-crafted. The question was: what kind of encoding?

PCA or DCT compression destroys texture at high compression ratios. What is needed is a learned perceptual compression that:

1. Preserves what the human visual system cares about: edges, textures, semantic content
2. Discards what it does not: statistical redundancy at the pixel level
3. Is differentiable, so the diffusion model can train on the compressed representation

This is exactly what a KL-regularized VAE trained with perceptual and adversarial losses provides.

---

Historical Context - Separating Rate from Distortion

The idea of separating compression from generation comes from information theory. Shannon's rate-distortion theory establishes that there is a fundamental tradeoff between compression ratio (rate) and reconstruction error (distortion). For images, perceptual quality correlates poorly with pixel-level metrics like PSNR - the human visual system cares about structure, not individual pixel values.

The deep learning community rediscovered this insight with VQ-VAE (van den Oord et al. 2017), which learned discrete codebook-based image representations. VQGAN (Esser et al. 2021) combined vector quantization with a perceptual-adversarial loss to achieve high-quality image compression at 16x spatial downsampling - and showed that a transformer trained in this compressed space could generate high-quality images autoregressively.

Rombach et al. (2022) made the critical connection: instead of using the compressed space for autoregressive generation, use it for diffusion. Train a high-quality autoencoder once, then train diffusion in the compressed latent space. The autoencoder handles perceptual fidelity. The diffusion model handles semantics. The result: 48x fewer dimensions for the diffusion process on a 512x512 image (from 786K to 16K), similar perceptual quality, and the ability to train on consumer-class hardware.

---

1. The Two-Stage Training Pipeline

Stage 1: Perceptual Compression - The Autoencoder

Train an encoder-decoder pair $(\mathcal{E}, \mathcal{D})$ on images:

$z = \mathcal{E}(x), \quad \hat{x} = \mathcal{D}(z), \quad z \in \mathbb{R}^{h \times w \times c}$

For a spatial downsampling factor $f$, an input image $x \in \mathbb{R}^{H \times W \times 3}$ maps to $z \in \mathbb{R}^{H/f \times W/f \times c}$.

Downsampling	Input	Latent shape	Spatial reduction	Total compression
$f=4$	256x256	64x64x3	16x	~5x
$f=8$	512x512	64x64x4	64x	48x
$f=8$	256x256	32x32x4	64x	48x
$f=16$	256x256	16x16x4	256x	192x

Stable Diffusion 1.x uses $f=8$: a 512x512x3 image maps to a 64x64x4 latent. The diffusion process operates on this 16,384-dimensional space. Compare to pixel space: 786,432 dimensions. That is 48x fewer dimensions - the key computational advantage.

The autoencoder is trained once and frozen before diffusion training begins.

Stage 2: Semantic Compression - Diffusion in Latent Space

Train a diffusion model in the latent space $\mathcal{Z}$:

$L_{LDM} = \mathbb{E}_{\mathcal{E}(x), \varepsilon \sim \mathcal{N}(0,I), t}\!\left[\|\varepsilon - \varepsilon_\theta(z_t, t, \tau_\theta(y))\|^2\right]$

where $z_t = \sqrt{\bar\alpha_t}z_0 + \sqrt{1-\bar\alpha_t}\varepsilon$ is the noised latent, $z_0 = \mathcal{E}(x)$ is the encoded clean image, and $\tau_\theta(y)$ is the conditioning signal (text embedding, class label, etc.).

Key points: the encoder $\mathcal{E}$ is frozen. At inference, sample $z_0$ via diffusion (starting from Gaussian noise), then decode $\hat{x} = \mathcal{D}(z_0)$ to get the pixel image. The VAE decode step is a single fast forward pass, not part of the diffusion loop.

---

2. Autoencoder Design - Why Plain L2 Fails

The Problem with Pixel-Level Reconstruction Loss

A VAE trained with only pixel-level L2 loss learns to produce the mean of all plausible reconstructions. Given an ambiguous image region (e.g., a blurry background), the optimal L2 solution is the average - which is itself blurry. The encoder learns to compress and the decoder learns to average, resulting in a latent space where reconstructions are always smooth and low-frequency.

When you train diffusion in a blurry latent space, the diffusion model inherits the blur. Worse, the model wastes capacity learning to model something that cannot be sharp - the decoder's blur becomes a hard ceiling on output quality. Every denoising step is predicting "more blur."

The Three-Part Loss Function

The LDM autoencoder uses a principled three-part objective:

$L_{AE} = L_{rec} + \lambda_{perc} \cdot L_{perc} + \lambda_{adv} \cdot L_{adv} + \lambda_{KL} \cdot L_{KL}$

$L_{rec}$ - Pixel reconstruction (L1 preferred over L2):

$L_{rec} = \|x - \mathcal{D}(\mathcal{E}(x))\|_1$

L1 is sharper than L2 because it penalizes absolute deviation equally regardless of magnitude, rather than squaring small errors (which encourages blurring).

$L_{perc}$ - Perceptual loss (LPIPS):

$L_{perc} = \sum_{l \in \{relu_1, relu_2, relu_3\}} \|VGG_l(x) - VGG_l(\hat{x})\|_2^2$

Computes VGG-16 feature distance between real and reconstructed images. VGG features respond to textures, shapes, and semantic content - the things humans actually perceive. Two images can be nearly identical to a human but have large pixel-level L2 distance (e.g., a tiny spatial shift). LPIPS penalizes perceptual differences, not pixel differences.

The frozen VGG-16 used for LPIPS acts as a "perceptual critic" - the autoencoder must fool it into thinking the reconstruction has the same features as the original.

$L_{adv}$ - PatchGAN adversarial loss:

A PatchGAN discriminator operates on $70 \times 70$ pixel patches, classifying each patch as real or generated. This forces local realism at the patch level - blurry or textureless regions that fool pixel-level metrics get penalized by the discriminator. The adaptive weight balancing:

$\lambda_{adv} = \frac{\nabla_{G_{last}}[L_{rec} + \lambda_{perc} L_{perc}]}{\nabla_{G_{last}}[L_{GAN}] + \delta}$

ensures that the adversarial loss does not dominate reconstruction early in training when the discriminator is strong.

$L_{KL}$ - KL regularization (used in SD):

$L_{KL} = D_{KL}(q_\phi(z|x) \| \mathcal{N}(0,I))$

with a very small weight ($\lambda_{KL} \sim 10^{-6}$). The autoencoder is nearly deterministic - the KL penalty is just enough to prevent the latent space from becoming discontinuous (ensuring nearby images map to nearby latents), without forcing the latent space to be exactly Gaussian (which would degrade reconstruction quality).

Why the Combination Works

Each loss component addresses a different failure mode:

• Without $L_{rec}$: the model might learn to fool the perceptual/adversarial losses without actually reconstructing the input
• Without $L_{perc}$: reconstructions are blurry (L1 alone is not sharp enough for textures)
• Without $L_{adv}$: fine details and high-frequency texture are lost (the discriminator forces local sharpness that losses on mean values miss)
• Without $L_{KL}$: the latent space can become fragmented, with no smooth interpolation between nearby images

Together they produce a near-losslessly perceptual autoencoder.

---

3. VQ-VAE vs KL-VAE - Two Regularization Approaches

The original LDM paper tried both VQ-regularized and KL-regularized autoencoders. Stable Diffusion uses KL-VAE.

VQ-VAE (Vector Quantized)

The encoder output is quantized to the nearest code in a learned codebook of $K$ vectors:

$z_q = \text{quantize}(\mathcal{E}(x)) = \arg\min_{e_k \in \mathcal{C}} \|\mathcal{E}(x) - e_k\|_2$

During training, the straight-through estimator passes gradients through the non-differentiable argmin. The commitment loss $\|sg[\mathcal{E}(x)] - z_q\|^2 + \beta\|\mathcal{E}(x) - sg[z_q]\|^2$ trains the encoder to commit to codebook entries.

Advantages: discrete latent space is compatible with autoregressive transformers (like VQGAN + GPT). Codebook provides an implicit prior that can be used for unconditional generation.

Disadvantages: codebook collapse is a training instability - many codebook entries become unused. Quantization introduces discrete artifacts. Less suitable for diffusion models which naturally operate on continuous Gaussian distributions.

KL-VAE (used in Stable Diffusion)

The encoder outputs a Gaussian distribution $q_\phi(z|x) = \mathcal{N}(\mu_\phi(x), \sigma_\phi(x))$. The KL penalty regularizes this toward $\mathcal{N}(0,I)$:

$L_{KL} = \frac{1}{2}\sum_j\left(\mu_j^2 + \sigma_j^2 - 1 - \log\sigma_j^2\right)$

With a tiny KL weight ($10^{-6}$), the encoder is nearly deterministic - it always produces nearly the same $z$ for the same $x$, with negligible variance. The regularization just prevents pathological behavior at the boundaries of the latent space.

Advantages: smooth, continuous latent space. Directly compatible with Gaussian diffusion. No codebook instabilities.

The VAE scaling factor: the KL-VAE produces latents with a particular variance. Empirically, the SD 1.x latents have a standard deviation of approximately 5.5. To normalize them to unit variance (matching the Gaussian noise in diffusion), a scaling factor of $1/5.5 \approx 0.18215$ is applied:

$z_{scaled} = \mathcal{E}(x) \times 0.18215$

This is vae.config.scaling_factor in the HuggingFace diffusers library. The scaling is critical: the diffusion model was trained on these scaled latents. Forgetting it shifts the latent distribution and produces garbage.

---

4. Memory Math - Why Latent Diffusion Fits on Consumer GPUs

Compute Comparison: Pixel Space vs Latent Space

For a 512x512 image with $T=1000$ DDPM steps:

Pixel space U-Net (hypothetical):

• Input per step: $512 \times 512 \times 3 = 786{,}432$ values
• Self-attention at 64x64 feature map: $4{,}096^2 = 16.7M$ ops per head
• If 16 heads: $268M$ attention ops per layer, per step
• Total: prohibitive - requires multi-day training on 256+ V100s

Latent space U-Net (SD 1.5):

• Input per step: $64 \times 64 \times 4 = 16{,}384$ values - 48x fewer
• Self-attention at 8x8 feature map: $64^2 = 4{,}096$ ops per head
• 16 heads: $65K$ attention ops per layer - 4,096x cheaper than pixel-space 64x64 feature
• Total training: feasible on 32 A100s in a few weeks

# Memory comparison: pixel vs latent diffusion

pixel_dims     = 512 * 512 * 3      # 786,432
latent_dims    = 64 * 64 * 4        # 16,384
compression    = pixel_dims / latent_dims  # 48x

# Self-attention at U-Net bottleneck
# Pixel space U-Net: deepest feature map 64x64 (after 3x downsample from 512)
pixel_attn_len   = 64 * 64           # 4,096 tokens
pixel_attn_ops   = pixel_attn_len**2 # 16,777,216 per head

# Latent space U-Net: deepest feature map 8x8 (after 3x downsample from 64)
latent_attn_len  = 8 * 8             # 64 tokens
latent_attn_ops  = latent_attn_len**2 # 4,096 per head

attn_speedup = pixel_attn_ops / latent_attn_ops  # 4,096x

print(f"Input dimensions per step: {pixel_dims:,} → {latent_dims:,} ({compression:.0f}x fewer)")
print(f"Attention ops at bottleneck: {pixel_attn_ops:,} → {latent_attn_ops:,} ({attn_speedup:,.0f}x fewer)")
print()
print("GPU VRAM for 512x512 inference (FP16):")
print("  Pixel-space model (hypothetical): ~40+ GB for U-Net + activations")
print("  Latent-space SD 1.5:               ~4.5 GB total (U-Net + VAE + CLIP)")
print("  Consumer GPU (8GB RTX 3070):       Fits! With attention slicing.")

Why 8x Downsampling Works But 16x Fails

At $f=8$: a 512x512 image maps to 64x64x4. At this compression level, the KL-VAE with perceptual+adversarial training is essentially lossless - humans cannot distinguish $x$ from $\mathcal{D}(\mathcal{E}(x))$ at casual inspection. Textures, fine edges, facial features, and text are all preserved.

At $f=16$: a 512x512 image maps to 32x32x4. At this compression, even with perceptual training, fine spatial details (eyes, hair texture, small text) begin to blur in reconstruction. The diffusion model cannot recover what the VAE has discarded. FID degrades noticeably at $f=16$ compared to $f=8$ for the same U-Net capacity.

Rombach et al.'s ablation confirms: $f=8$ is the sweet spot between compression efficiency and reconstruction quality for natural images.

---

5. Text Conditioning via Cross-Attention

The Conditioning Pipeline

For text-to-image generation, the prompt $y$ is converted to embeddings via a text encoder $\tau_\theta$ (typically CLIP):

$\tau_\theta(y) \in \mathbb{R}^{L \times d_\tau}$

where $L = 77$ is the token sequence length (CLIP's context window) and $d_\tau$ is the embedding dimension (768 for CLIP-L, 1280 for OpenCLIP-bigG).

These embeddings condition the U-Net via cross-attention at every decoder block.

Cross-Attention Mechanics

At each cross-attention layer, the spatial features $\varphi_i(z_t)$ of shape $(N, d_\epsilon)$ act as queries, and the text embeddings $\tau_\theta(y)$ act as keys and values:

$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d}}\right)V$

$Q = W_Q \cdot \varphi_i(z_t), \quad K = W_K \cdot \tau_\theta(y), \quad V = W_V \cdot \tau_\theta(y)$

The matrices $W_Q \in \mathbb{R}^{d \times d_\epsilon}$, $W_K, W_V \in \mathbb{R}^{d \times d_\tau}$ are learned projections. The attention map $A \in \mathbb{R}^{N \times L}$ tells each of the $N$ spatial positions which of the $L$ text tokens to attend to. This is where text-image spatial alignment is mechanistically learned.

Why Cross-Attention, Not Concatenation?

Earlier conditioning methods (e.g., class-conditional diffusion) concatenated or added the conditioning to the spatial features. This works when conditioning is a single vector (a class embedding). For text, it fails:

• Text has variable length (1 to 77 tokens)
• Different words correspond to different spatial regions ("the red ball in the left corner")
• Global concatenation makes all spatial positions respond equally to all words

Cross-attention solves all three problems: spatial positions can selectively attend to relevant tokens, attention is computed over sequences of any length, and the learned $W_Q, W_K, W_V$ projections can transform between text embedding space and spatial feature space.

Research on attention map visualization (Prompt-to-Prompt, Attend-and-Excite) has confirmed empirically that cross-attention maps learn semantically meaningful correspondences - "cat" attends to cat regions, "red" attends to red regions, spatial prepositions influence the spatial distribution of attention.

---

6. CLIP Text Encoder Integration

Why CLIP

CLIP (Contrastive Language-Image Pretraining, Radford et al. 2021) was trained on 400M image-text pairs to align image and text embeddings in a shared space. Its text encoder has seen an enormous variety of descriptive language, making it excellent at encoding prompts into meaningful vectors.

For diffusion models, CLIP's text encoder is used purely as a feature extractor - only the text branch is used (not the image branch). The text encoder is frozen during LDM training; it acts as a fixed "language understanding module" that the U-Net cross-attention learns to read.

CLIP-L vs OpenCLIP-bigG

CLIP ViT-L/14 (used in SD 1.x):

• Context window: 77 tokens
• Embedding dim: 768
• 123M parameters in text encoder
• Trained by OpenAI on proprietary 400M pairs dataset
• Strong at common objects, styles, and scenes
• Weaker at complex spatial relationships and unusual noun phrases

OpenCLIP ViT-bigG/14 (used in SDXL):

• Context window: 77 tokens
• Embedding dim: 1280
• 354M parameters in text encoder
• Trained on open LAION-2B dataset
• Better at rare concepts, fine-grained attributes, longer descriptions
• 1280 vs 768 dimensions provides richer semantic space

SDXL concatenates both encoders: $[768 \mid 1280] = 2048$-dimensional conditioning per token position, providing both OpenAI's and OpenCLIP's text understanding simultaneously.

Token Budget and Long Prompts

CLIP tokenizes with a BPE tokenizer. Common words = 1 token. Rare words or compound words = multiple tokens. The 77-token limit means complex prompts are silently truncated. The model attends only to the first 77 tokens.

Some implementations handle long prompts by:

1. Token chunking: split into 77-token windows, run CLIP separately, average embeddings - loses cross-token dependencies
2. T5/LLM text encoder: SD3 and FLUX use T5-XXL with a 256-token limit and much richer semantic encoding

---

7. Architecture Diagram

---

8. Stable Diffusion Versions - What Changed

SD 1.x (2022)

• VAE: KL-VAE, $f=8$, scaling factor 0.18215 - 512px → 64x64x4
• U-Net: 860M parameters, 4 resolution levels, cross-attention at every decoder block
• Text encoder: OpenAI CLIP ViT-L/14, 77 tokens × 768 dims
• Training data: LAION-5B filtered subset, ~5 billion image-text pairs
• Native resolution: 512x512 (distorted results outside this resolution)
• Community adoption: massive - thousands of fine-tuned checkpoints, LoRAs, ControlNets

SD 2.x (2022-2023)

• VAE: new VAE with improved reconstruction at 768px
• Text encoder: OpenCLIP ViT-H/14 (open weights), 77 tokens × 1024 dims
• U-Net: similar to 1.5 but wider; optimized for 768px
• Notable change: NSFW filter changed training distribution - SD 1.5 LoRAs do not transfer cleanly to SD 2.x
• Community outcome: less adopted than SD 1.5 due to limited ecosystem compatibility

SDXL (2023)

• Two text encoders: CLIP-L (768-dim) + OpenCLIP-bigG (1280-dim) → concatenated 2048-dim conditioning
• U-Net: 2.6B parameters (3x larger than SD 1.5), wider channels at all resolutions
• Refiner model: second 1.3B U-Net for the last 200 timesteps - adds fine-grained detail to the base model's composition
• Native resolution: 1024x1024 with aspect-ratio conditioning
• Aesthetic conditioning: Fourier-embedded image dimensions and crop coordinates prevent artifacts from mixed-resolution training
• Quality jump: substantially better text-image alignment, anatomy, and typography than SD 1.5

SDXL aesthetic conditioning encodes three extra signals alongside timestep $t$:

original_size:   (H_orig, W_orig) - original image resolution before any resizing
crops_coords:    (crop_y, crop_x) - random crop location during training data prep
target_size:     (H_out, W_out)  - desired output resolution

Each signal: Fourier embedding → MLP → added to timestep embedding
Purpose: Allows model to know "this image was upscaled/cropped" vs "native resolution"
         → prevents the model from learning artifacts from low-quality training data

SD3 and FLUX (2024)

• Architecture: replaces U-Net with DiT (Diffusion Transformer) - MM-DiT applies joint attention between image patches and text tokens at every layer
• Three text encoders: CLIP-L + OpenCLIP-bigG + T5-XXL (4096-dim, 256 tokens) - dramatically better long-prompt understanding
• Noise schedule: Rectified Flow - straight-line ODE paths between noise and data, requiring fewer steps
• FLUX.1: state-of-the-art photorealism and text rendering; surpasses Midjourney v6 on multiple benchmarks
• Why DiT: transformers scale better with compute and data; global attention at every layer eliminates the bottleneck architecture's global coherence limitations

---

9. Complete PyTorch Inference Pipeline

"""
Complete Stable Diffusion inference pipeline.
All shapes annotated throughout.
"""

import torch
import numpy as np
from PIL import Image
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import (
    AutoencoderKL,
    UNet2DConditionModel,
    DDIMScheduler,
)


class ManualSDPipeline:
    """
    Transparent Stable Diffusion inference pipeline.
    Shows every component, shape, and design decision.
    Use diffusers.StableDiffusionPipeline for production;
    use this to understand the architecture.
    """

    def __init__(self, model_id: str = "runwayml/stable-diffusion-v1-5"):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.float16 if self.device == "cuda" else torch.float32

        print("Loading components...")
        self.vae = AutoencoderKL.from_pretrained(
            model_id, subfolder="vae"
        ).to(self.device, dtype=dtype)

        # Key constants:
        # vae_scale_factor = 8 (f=8 spatial downsampling)
        # latent_scale = 0.18215 (normalizes latent variance to ~1)
        self.vae_scale_factor = 8
        self.latent_scale = self.vae.config.scaling_factor  # 0.18215

        self.tokenizer = CLIPTokenizer.from_pretrained(
            model_id, subfolder="tokenizer"
        )
        self.text_encoder = CLIPTextModel.from_pretrained(
            model_id, subfolder="text_encoder"
        ).to(self.device, dtype=dtype)

        self.unet = UNet2DConditionModel.from_pretrained(
            model_id, subfolder="unet"
        ).to(self.device, dtype=dtype)

        self.scheduler = DDIMScheduler.from_pretrained(
            model_id, subfolder="scheduler"
        )

        print(f"Pipeline loaded on {self.device}")
        print(f"  VAE: f={self.vae_scale_factor}, scaling_factor={self.latent_scale}")
        print(f"  U-Net in_channels: {self.unet.config.in_channels}")  # 4
        print(f"  CLIP context length: {self.tokenizer.model_max_length}")  # 77

    @torch.no_grad()
    def encode_text(
        self,
        prompt: str,
        negative_prompt: str = ""
    ) -> torch.Tensor:
        """
        Tokenize and encode text with CLIP.

        Input:  string prompt
        Output: (2, 77, 768) - conditional + unconditional embeddings
        """
        # Tokenize both prompts together as a batch of 2
        inputs = self.tokenizer(
            [prompt, negative_prompt],
            padding="max_length",
            max_length=self.tokenizer.model_max_length,  # 77
            truncation=True,
            return_tensors="pt"
        ).to(self.device)

        # CLIP text encoder: (2, 77) token IDs → (2, 77, 768) embeddings
        embeddings = self.text_encoder(inputs.input_ids)[0]

        print(f"  Text embeddings: {embeddings.shape}")  # (2, 77, 768)
        # embeddings[0] = conditional (prompt)
        # embeddings[1] = unconditional (negative/empty)
        return embeddings

    @torch.no_grad()
    def encode_image_to_latent(
        self,
        image_tensor: torch.Tensor
    ) -> torch.Tensor:
        """
        Encode a pixel image to VAE latent space.
        Used for img2img; not needed for text-to-image.

        Input:  (B, 3, 512, 512) pixels in [-1, 1]
        Output: (B, 4, 64, 64)  scaled latents
        """
        # VAE encode: pixel → Gaussian distribution over latent
        posterior = self.vae.encode(image_tensor).latent_dist

        # Sample from the posterior (nearly deterministic with small KL weight)
        z0 = posterior.sample()

        # Apply scaling factor: normalizes latent variance to ~1
        # CRITICAL: forget this and diffusion fails
        z0_scaled = z0 * self.latent_scale

        print(f"  Encoded latent: {z0_scaled.shape}")  # (B, 4, 64, 64)
        print(f"  Latent std (should be ~1.0): {z0_scaled.std():.3f}")
        return z0_scaled

    @torch.no_grad()
    def decode_latent_to_image(
        self,
        latents: torch.Tensor
    ) -> Image.Image:
        """
        Decode VAE latent to pixel image.

        Input:  (1, 4, 64, 64) latents (scaled)
        Output: PIL Image 512x512
        """
        # Reverse the scaling factor before decoding
        latents_unscaled = latents / self.latent_scale

        # VAE decode: (1, 4, 64, 64) → (1, 3, 512, 512) in [-1, 1]
        decoded = self.vae.decode(latents_unscaled).sample
        print(f"  Decoded tensor: {decoded.shape}")  # (1, 3, 512, 512)

        # Convert [-1, 1] → [0, 255] PIL image
        decoded = (decoded / 2 + 0.5).clamp(0, 1)  # [-1,1] → [0,1]
        decoded = decoded.cpu().float().permute(0, 2, 3, 1).numpy()  # BCHW → BHWC
        image = Image.fromarray((decoded[0] * 255).round().astype(np.uint8))
        return image

    @torch.no_grad()
    def generate(
        self,
        prompt: str,
        negative_prompt: str = "blurry, low quality, ugly",
        guidance_scale: float = 7.5,
        num_inference_steps: int = 50,
        seed: int = 42,
        height: int = 512,
        width: int = 512,
    ) -> Image.Image:
        """
        Full text-to-image generation pipeline.

        Shape progression:
            text_embeddings:   (2, 77, 768)
            initial latents:   (1, 4, 64, 64)
            U-Net input (CFG): (2, 4, 64, 64)
            U-Net output:      (2, 4, 64, 64)
            guided prediction: (1, 4, 64, 64)
            final z_0:         (1, 4, 64, 64)
            decoded image:     (1, 3, 512, 512)
        """
        generator = torch.Generator(self.device).manual_seed(seed)
        self.scheduler.set_timesteps(num_inference_steps)

        # === 1. Encode text ===
        print("1. Encoding text with CLIP...")
        text_embeddings = self.encode_text(prompt, negative_prompt)
        # text_embeddings[0] = conditional, text_embeddings[1] = unconditional

        # === 2. Initialize latent noise ===
        print("2. Initializing latent noise...")
        latent_h = height // self.vae_scale_factor  # 512/8 = 64
        latent_w = width // self.vae_scale_factor    # 512/8 = 64
        n_channels = self.unet.config.in_channels    # 4 for SD 1.x

        # Sample pure Gaussian noise in latent space
        latents = torch.randn(
            (1, n_channels, latent_h, latent_w),
            generator=generator,
            device=self.device,
            dtype=text_embeddings.dtype
        )
        print(f"  Initial noise latents: {latents.shape}")  # (1, 4, 64, 64)

        # DDIM scheduler requires scaling by initial noise magnitude
        latents = latents * self.scheduler.init_noise_sigma

        # === 3. DDIM denoising loop ===
        print(f"3. Denoising ({num_inference_steps} DDIM steps)...")
        for i, timestep in enumerate(self.scheduler.timesteps):
            # CFG: duplicate latents → batch size 2 for one forward pass
            # Much more efficient than two separate forward passes
            latent_input = torch.cat([latents] * 2)  # (2, 4, 64, 64)

            # Scheduler scales the input (timestep-dependent)
            latent_input = self.scheduler.scale_model_input(latent_input, timestep)

            # U-Net forward: predicts noise for both conditional and unconditional
            # Input:  (2, 4, 64, 64) latents + timestep + (2, 77, 768) text
            # Output: (2, 4, 64, 64) noise predictions
            noise_pred = self.unet(
                latent_input,
                timestep,
                encoder_hidden_states=text_embeddings  # cross-attention keys/values
            ).sample

            # Split CFG predictions: [unconditional, conditional]
            noise_uncond = noise_pred[1]   # unconditional (negative prompt)
            noise_cond   = noise_pred[0]   # conditional (positive prompt)

            # CFG: guided = uncond + scale * (cond - uncond)
            # Higher guidance_scale → stronger prompt adherence
            noise_guided = noise_uncond + guidance_scale * (
                noise_cond - noise_uncond
            )

            # DDIM step: x_t → x_{t-1} using guided noise prediction
            latents = self.scheduler.step(
                noise_guided, timestep, latents
            ).prev_sample

            if (i + 1) % 10 == 0:
                print(f"  Step {i+1}/{num_inference_steps}")

        print(f"  Final latents: {latents.shape}")  # (1, 4, 64, 64)

        # === 4. Decode ===
        print("4. Decoding with VAE...")
        return self.decode_latent_to_image(latents)


# ============================================================
# Compute savings analysis
# ============================================================

def compute_savings_analysis():
    """
    Quantifies the compute savings from latent vs pixel diffusion.
    """
    print("Compute Savings: Latent vs Pixel Diffusion")
    print("=" * 60)
    print()

    # Dimensions
    pixel_dims = 512 * 512 * 3    # 786,432
    latent_dims = 64 * 64 * 4    # 16,384
    dim_ratio = pixel_dims / latent_dims

    print(f"Pixel space per step:    {pixel_dims:>10,} values")
    print(f"Latent space per step:   {latent_dims:>10,} values")
    print(f"Dimension reduction:     {dim_ratio:>10.0f}x")
    print()

    # Self-attention at U-Net bottleneck
    # Pixel U-Net: 3 downsamplings from 512 → 64x64
    pixel_token_count = 64 * 64     # 4,096
    pixel_attn = pixel_token_count ** 2

    # Latent U-Net: 3 downsamplings from 64 → 8x8
    latent_token_count = 8 * 8      # 64
    latent_attn = latent_token_count ** 2

    print("Self-attention at U-Net bottleneck:")
    print(f"  Pixel space (64×64 features): {pixel_attn:>12,} ops/head")
    print(f"  Latent space (8×8 features):  {latent_attn:>12,} ops/head")
    print(f"  Attention savings:             {pixel_attn/latent_attn:>12,.0f}x")
    print()

    # Training time estimate (rough)
    # Assume: 1 step in pixel space = 900ms on V100
    # 1 step in latent space = 40ms on V100 (20x+ faster due to smaller input)
    pixel_step_ms = 900
    latent_step_ms = 40
    T = 1000

    pixel_train_step_s = pixel_step_ms / 1000
    latent_train_step_s = latent_step_ms / 1000

    print("Per-step inference time (single V100, approx):")
    print(f"  Pixel-space DDPM 1000-step:  {pixel_train_step_s * T:.0f}s = {pixel_train_step_s * T / 60:.0f} min")
    print(f"  Latent LDM, DDIM 50-step:    {latent_step_ms * 50 / 1000:.1f}s")
    print()

    # VRAM for 512x512 inference
    print("Approximate VRAM at 512×512 (FP16):")
    print("  Pixel-space DDPM (hypothetical): ~40+ GB - impossible on consumer GPU")
    print("  SD 1.5 LDM:                       ~4.5 GB - runs on RTX 3060 (12GB)")
    print("  SD 1.5 with CPU offload:           ~2.0 GB GPU - runs on 4GB GPU")


# ============================================================
# SDXL dual-encoder conditioning
# ============================================================

def sdxl_conditioning_breakdown():
    """
    SDXL text conditioning structure.
    Shows how two encoders combine.
    """
    print("SDXL Dual Text Encoder Conditioning")
    print("=" * 50)
    print()
    print("Encoder 1: CLIP ViT-L/14 (OpenAI)")
    print("  Token IDs  →  (B, 77, 768)")
    print("  Captures: common objects, artistic styles, basic attributes")
    print()
    print("Encoder 2: OpenCLIP ViT-bigG/14")
    print("  Token IDs  →  (B, 77, 1280)")
    print("  Captures: rare concepts, fine-grained semantics, complex descriptions")
    print()
    print("Concatenation:")
    print("  (B, 77, 768) + (B, 77, 1280) → concat on dim=-1 → (B, 77, 2048)")
    print("  All 77 positions have 2048-dim conditioning")
    print("  U-Net cross-attention keys/values: shape (B, 77, 2048)")
    print()
    print("Aesthetic conditioning (added to timestep embedding):")
    print("  original_size:       (H, W) e.g. (1024, 1024)")
    print("  crops_coords:        (top, left) e.g. (0, 0)")
    print("  target_size:         (H, W) e.g. (1024, 1024)")
    print("  Each: 2 values → Fourier embed → MLP → scalar added to t_emb")
    print("  Effect: model knows 'generate at target resolution, not a crop'")
    print()
    print("SDXL Refiner:")
    print("  Second U-Net (1.3B params) handles timesteps 0-200 (low noise)")
    print("  Base model handles 200-1000 (high noise, composition)")
    print("  Refiner sharpens fine details that base model sets up")


if __name__ == "__main__":
    compute_savings_analysis()
    print()
    sdxl_conditioning_breakdown()

---

10. YouTube Resources

Video	Channel	What You Learn
Latent Diffusion Models Paper Explained	Yannic Kilcher	Complete LDM paper walkthrough with all derivations
Stable Diffusion Architecture Deep Dive	Fast.ai	VAE, U-Net, CLIP shapes annotated with live code
How Stable Diffusion Works	Computerphile	Accessible visual pipeline explanation
SDXL Architecture Explained	AI Coffee Break	Dual encoders, refiner, aesthetic conditioning changes
Stable Diffusion From Scratch (PyTorch)	Umar Jamil	Build the complete pipeline in PyTorch - all shapes

---

11. Production Deployment Notes

:::tip Pre-encoding training images saves significant compute For datasets you train on repeatedly, pre-encode all images to latents and cache them. VAE encoding is not fast - for 1M images at 512x512, encoding takes several GPU-hours. Pre-encoded latents can be loaded directly, saving roughly 100ms per training step. This allows larger effective batch sizes or faster training without quality loss. :::

:::note TAESD for fast preview TAESD (Tiny AutoEncoder for Stable Diffusion) is a distilled version of the SD VAE with ~600K parameters instead of 83M. Decoding with TAESD takes ~5ms vs ~200ms for the full VAE - 40x faster. Quality is slightly lower (mild blurring) but acceptable for real-time preview use cases. Pattern: generate preview every 5-10 steps using TAESD, final output using full VAE. :::

:::note FlashAttention is essential for SDXL SDXL's U-Net (2.6B params) has large attention heads due to 1024x1024 input. Without FlashAttention, the attention matrix at the 128x128 feature map (after 3x downsample from 1024) is $16{,}384 \times 16{,}384$ - 268M elements per head. Standard PyTorch attention materializes this full matrix in VRAM. FlashAttention tiles the computation, reducing peak VRAM by ~30% and providing 2-4x speedup. Not optional for production SDXL deployments. :::

---

12. Common Mistakes

:::danger Forgetting the VAE latent scaling factor The SD VAE uses scaling_factor = 0.18215 (accessible via vae.config.scaling_factor). Multiply encoded latents by this before starting diffusion. Divide by it before decoding. This normalizes the latent distribution to approximately unit variance, matching the Gaussian noise used in diffusion. Forgetting to multiply: diffusion runs on latents with ~5.5x too large values - the noise schedule is calibrated for unit-variance latents, so all denoising predictions will be wildly miscalibrated. Output: pure noise or blank images. This is one of the most common bugs in from-scratch SD implementations. :::

:::warning Generating off-resolution images with SD 1.x SD 1.x was trained at 512x512. The U-Net has convolutional bias and positional encoding behaviors calibrated for this resolution. Generating at 256x256 produces tiled/repeated artifacts. Generating at 1024x1024 produces distorted compositions (the model does not know how to compose a full scene at higher resolution). For non-512 sizes with SD 1.x, use multi-diffusion (tiled inference with blending) or SDXL. SDXL was trained at 1024x1024 with aspect-ratio conditioning and handles varied resolutions correctly. :::

:::warning CFG scale is model-specific A guidance scale of 7.5 for SD 1.5 will produce different results on SDXL or SD3 - often oversaturation on SDXL (where 5-7 is more appropriate). The "optimal" guidance scale depends on training data, noise schedule, and U-Net capacity. Always re-tune guidance scale when switching base models. SDXL typically works best at 5-8, SD 1.5 at 7-12, SD3/FLUX at 3-5 (rectified flow needs less guidance). :::

:::danger Indexing CFG noise predictions in wrong order In the standard implementation, the unconditional and conditional branches are concatenated as torch.cat([latents, latents]) and the text embeddings as torch.cat([uncond, cond]). After the U-Net forward pass, noise_pred.chunk(2) gives [uncond, cond] in that order. Swapping them inverts the guidance: guided = cond + scale * (uncond - cond) = guided away from the prompt. The model generates images that maximally avoid the text description. Common mistake when adapting example code - always verify the ordering. :::

---

13. Interview Q&A

Q1: Why does compressing to latent space not reduce image quality? How does the VAE preserve perceptual information?

The quality preservation is entirely due to the training objective. A VAE trained with L2 reconstruction loss learns to blur - the optimal solution to MSE is the mean of all plausible values, which is always smooth. Adding LPIPS perceptual loss (VGG feature distance) forces the decoder to preserve structures perceptible to humans: textures, edges, fine detail. Adding a PatchGAN discriminator forces local realism - the discriminator penalizes blurry patches even when pixel error is small, because real patches are never smooth in the frequency domain. Together, these losses make the VAE effectively lossless at $f=8$: humans cannot distinguish original from reconstructed at casual inspection. The diffusion model then operates on semantically meaningful representations, not pixel noise.

Q2: How does cross-attention achieve text-image alignment in the U-Net?

Cross-attention allows each spatial position in the U-Net feature map to selectively attend to different text tokens. The spatial feature at position $(h,w)$ computes a query vector $q = W_Q \varphi(h,w)$. This attends over all 77 text embeddings to compute an attention distribution $\text{softmax}(q K^T / \sqrt{d})$ of shape $(1, 77)$ - which words matter for this spatial position. The output is $\sum_l a_l v_l$ - a weighted combination of text value vectors. The learned $W_Q, W_K, W_V$ learn to route different words to different spatial regions. Research on attention maps (Prompt-to-Prompt) confirms that "cat" has high attention in cat regions, "red" in red regions, etc.

Q3: What are the key differences between SD 1.5 and SDXL, and when would you use each?

SD 1.5: 860M U-Net, CLIP-L text encoder (768-dim), $f=8$ VAE, 512px native. Massive community ecosystem - thousands of fine-tuned checkpoints, LoRAs, ControlNets. Runs on 4GB VRAM. Best for: leveraging community models, fast inference, resource-constrained deployment, or cases where SD 1.5 community fine-tunes cover the specific style needed.

SDXL: 2.6B U-Net, dual text encoders (CLIP-L + OpenCLIP-bigG, 2048-dim), improved VAE, 1024px native. Substantially better text-image alignment on complex prompts, anatomy, and typography. Requires 8GB+ VRAM. Best for: highest quality without a specific community fine-tune, better prompt adherence on complex descriptions, 1024px resolution needs.

In production: SD 1.5 for real-time/resource-constrained systems; SDXL for maximum quality; SD3/FLUX for state-of-the-art photorealism or text rendering.

Q4: What is MM-DiT in SD3 and FLUX, and why is it replacing U-Nets?

MM-DiT (Multi-Modal Diffusion Transformer) patchifies the latent image (e.g., 2x2 patches → 32x32 = 1024 tokens) and applies full transformer attention over image patches and text tokens simultaneously at every layer. Unlike the U-Net's hierarchical encoder-decoder, every DiT layer has global receptive field.

Why it is replacing U-Nets: (1) Transformers scale better with compute and data - the same scaling laws that produced GPT-3 apply. (2) No resolution bias - no convolutional inductive bias tied to training resolution. (3) Better global coherence - attention at every layer (not just the bottleneck) produces more globally consistent images. (4) Three text encoders in SD3 (CLIP-L + OpenCLIP-G + T5-XXL) give far richer semantic conditioning. The T5-XXL encoder handles 256 tokens with much better understanding of complex, long prompts.

Q5: Walk through the complete LDM training recipe in order.

Stage 1 - Autoencoder training (done once, frozen afterward): Initialize KL-VAE encoder-decoder. Train with: pixel reconstruction (L1) + perceptual loss (LPIPS, frozen VGG-16) + PatchGAN adversarial loss + tiny KL penalty ($\lambda_{KL} \sim 10^{-6}$). Scale: 256-512px images, batch 128, 50-100K iterations. Evaluate: PSNR, SSIM, LPIPS on held-out images. Freeze when LPIPS is satisfactory (typically $<0.05$).

Stage 2 - Diffusion model training: Pre-encode or on-the-fly encode training images to latents using frozen $\mathcal{E}$. Apply scaling_factor. Per step: sample image → encode to $z_0$ → sample $t \sim \mathcal{U}[0,T]$ and $\varepsilon \sim \mathcal{N}(0,I)$ → compute $z_t = \sqrt{\bar\alpha_t}z_0 + \sqrt{1-\bar\alpha_t}\varepsilon$ → encode text with frozen CLIP → predict $\varepsilon$ from $z_t$, $t$, and text embedding via cross-attention → minimize $\|\varepsilon - \varepsilon_\theta(z_t, t, \tau_\theta(y))\|^2$. Apply CFG dropout: 10-15% of steps use null/empty text conditioning. Use AdamW, cosine LR, EMA with decay 0.9999.

Q6: Why is the VAE scaling factor 0.18215 and what happens without it?

The scaling factor normalizes the variance of the encoded latent distribution to approximately 1.0. The KL-VAE is trained with a very small KL coefficient ($10^{-6}$), so it is nearly deterministic and the latent space can develop arbitrary variance. Empirically, the SD 1.x latents have a standard deviation of roughly 5.5 before scaling. Multiplying by $1/5.5 \approx 0.18215$ gives unit-variance latents.

The diffusion model is trained on these unit-variance latents. Its noise schedule (e.g., cosine schedule) is calibrated for unit-variance signals: $\bar\alpha_T \approx 0$ means pure noise at step $T$, but "pure noise" means $\mathcal{N}(0,I)$ - unit variance. If the latents have variance 5.5, then at $\bar\alpha_T = 0.001$, the "noisy" latent is $\sqrt{0.001} \cdot 5.5 \cdot z_0 + \sqrt{0.999} \cdot \varepsilon$ - the signal term is not negligible, so the prior is not pure noise. The model was never trained to start from this distribution. Result: generated images will either be blank (pure noise propagated through decoder) or incoherent.

---

14. Inpainting, Outpainting, and img2img with LDMs

Inpainting Architecture

Inpainting (filling in a masked image region) extends the LDM with minimal modifications. The U-Net's input is extended from 4 channels to 9 channels:

• $z_t$: noisy latent (4 channels) - standard diffusion input
• $m$: downsampled binary mask (1 channel, 64x64) - 1 = preserve original, 0 = inpaint
• $z_{masked}$: VAE latent of the source image with masked region zeroed (4 channels)

Total: 9 channels. The U-Net's initial conv layer is widened from 4 to 9 input channels by adding zero-initialized weights for the 5 new channels. The pretrained model's behavior is preserved; only the new channels are trained during fine-tuning.

At inference: provide the source image, a mask, and a text prompt. The model denoises the masked region while the unmasked region guides coherent blending via the $z_{masked}$ signal.

Training recipe: generate random (image, mask) pairs. Masks should be a mixture of rectangular boxes, circular regions, irregular polygons, and brush-stroke shapes - diversity prevents the model from overfitting to one mask type. Set $z_{masked}$ by zeroing the masked region of the encoded latent before feeding to the U-Net.

img2img - Noise-and-Denoise Editing

img2img (image-to-image) editing uses the forward diffusion process to add partial noise to a real image, then denoises with a new text prompt. The workflow:

1. Encode the source image: $z_0 = \mathcal{E}(x_{source}) \times \text{scaling\_factor}$
2. Add noise to an intermediate timestep $t_{start}$: $z_{t_{start}} = \sqrt{\bar\alpha_{t_{start}}}z_0 + \sqrt{1-\bar\alpha_{t_{start}}}\varepsilon$
3. Denoise from $t_{start}$ back to $t=0$ with the new text prompt

The denoising strength parameter (typically 0.0 to 1.0) controls $t_{start}$ as a fraction of $T$:

• 0.0: no noise added, no denoising - returns source image unchanged
• 0.5: moderate noise, denoising reshapes style/details while preserving structure
• 1.0: full noise (equivalent to text-to-image generation), source image has no influence

This provides a continuous control over how much the output resembles the source. Strength 0.3-0.5 is typical for style transfer; 0.7-0.9 for significant content changes.

Outpainting

Outpainting (extending an image beyond its boundaries) uses inpainting with an expanded canvas. Place the original 512px image in the center of a 768px canvas. Fill the surrounding area with noise or zeros. Set the mask to cover the extension area. The model generates a plausible continuation that matches the source image's perspective, lighting, and style.

Quality is best when the extension direction has natural context - extending sky or floor textures is easier than extending complex scenes. Multi-step outpainting (extend by 128px, then use the result as a new source and extend again) improves quality over single large extensions.

---

15. The VAE Latent Space - Structure and Properties

What Do the 4 Latent Channels Encode?

The SD VAE produces a 4-channel latent $z \in \mathbb{R}^{64 \times 64 \times 4}$. Unlike the 3 RGB channels of pixel space, the 4 latent channels do not have direct human-interpretable meaning. They are learned representations shaped by the combination of reconstruction loss, perceptual loss, adversarial loss, and KL regularization.

Empirically, researchers have found via ablation and visualization:

• The 4 channels encode color/luminance information in a distributed, entangled way
• Different channels respond to different frequency bands of the image
• No single channel corresponds to "edges" or "texture" in isolation - information is holistically distributed
• The 64x64 spatial structure maps roughly to semantic regions (face region activates together, background activates separately) but is not precisely spatially aligned at the pixel level

Latent Space Interpolation

Because the VAE maps similar images to nearby latents (enforced by the smooth perceptual + KL training), interpolating between two latents produces a smooth visual morphing sequence:

$z_{interp}(\alpha) = (1-\alpha)\,\mathcal{E}(x_1) + \alpha\,\mathcal{E}(x_2), \quad \alpha \in [0, 1]$

Then $\mathcal{D}(z_{interp}(\alpha))$ produces a smooth transition from image $x_1$ to image $x_2$. This works because the perceptual+adversarial training ensures the latent space is smooth - the decoder must produce perceptually natural images for all points along the interpolation path, not just the endpoints.

Applications: video morphing (keyframe images → smooth interpolated sequence), progressive style transfer, creative exploration of variations around a seed image.

Latent Arithmetic and Semantic Editing

Inspired by word2vec's semantic arithmetic (king - man + woman ≈ queen), researchers have explored latent arithmetic in LDMs. For two pairs of images $(x_A, x_{A'})$ differing in attribute $a$ (e.g., day/night), the difference vector $\delta = \mathcal{E}(x_{A'}) - \mathcal{E}(x_A)$ can be applied to a new image: $\mathcal{D}(\mathcal{E}(x_B) + \delta) \approx x_{B'}$ with attribute $a$ applied.

This works reliably for global attribute edits (color grading, time of day, weather). It works less reliably for local edits (moving a specific object) because the VAE latent space encodes spatial information locally - adding a global delta affects all spatial regions.

For precise semantic editing in production: DDIM inversion + Prompt-to-Prompt attention manipulation is more reliable than latent arithmetic for LDMs. Latent arithmetic is better suited for global style changes.

Checking Latent Distribution Health

When debugging an LDM pipeline, it is useful to verify the latent distribution statistics at every stage:

def check_latent_health(vae, images: torch.Tensor, scaling_factor: float):
    """
    Diagnostic check for correct latent distribution.
    Expected values when pipeline is correctly configured:
      - Pre-scaling std:  ~4.0–6.5 (model-dependent)
      - Post-scaling std: ~0.9–1.1 (should be near 1.0)
      - Reconstructed LPIPS: < 0.05 (imperceptible difference)
    """
    with torch.no_grad():
        posterior = vae.encode(images).latent_dist
        z_raw = posterior.mode()                   # use mode (no sampling)
        z_scaled = z_raw * scaling_factor

        print(f"Raw latent: mean={z_raw.mean():.3f}, std={z_raw.std():.3f}")
        print(f"Scaled latent: mean={z_scaled.mean():.3f}, std={z_scaled.std():.3f}")

        # Reconstruction sanity check
        z_decoded = vae.decode(z_scaled / scaling_factor).sample
        recon_range = (z_decoded.min().item(), z_decoded.max().item())
        print(f"Reconstructed pixel range: {recon_range}")

    if abs(z_scaled.std().item() - 1.0) > 0.2:
        print("WARNING: Scaled latent std is far from 1.0 - check scaling_factor!")
    else:
        print("OK: Latent distribution looks healthy.")

Running this check before diffusion training or inference on a new model configuration catches the most common pipeline bugs: wrong scaling factor, wrong normalization range, or an incompatible VAE checkpoint.

---

This lesson is part of the Diffusion Models module. Next: Classifier-Free Guidance - Steering Diffusion with Text.

:::tip 🎮 Interactive Playground

Visualize this concept: Try the Diffusion Process (DDPM) demo on the EngineersOfAI Playground - no code required.

:::