Phase 08 - Lesson 07

Latent Diffusion & Stable Diffusion

La difusión en el espacio de píxeles en imágenes de 512×512 es un crimen de guerra computacional. Rombach et al. (2022) notaron que no se necesitan todas las 786k dimensiones para generar una imagen — se necesita lo suficiente para capturar la estructura semántica y un decodificador separado para el resto. Ejecuta la difusión dentro del espacio latente de un VAE. Esa única idea es Stable Diffusion.

Tipo: Build Lenguajes: Python Prerrequisitos: Phase 8 · 02 (VAE), Phase 8 · 06 (DDPM), Phase 7 · 09 (ViT) Tiempo: ~75 minutos

El Problema

La difusión en el espacio de píxeles a 512² significa que la U-Net se ejecuta en tensores de forma [B, 3, 512, 512]. Cada paso de muestreo es de ~100 GFLOPS para una U-Net de 500M de parámetros. Cincuenta pasos equivalen a 5 TFLOPS por imagen. Entrenar con mil millones de imágenes hace que la factura de cómputo sea absurda.

La mayoría de esos FLOPs se destinan a empujar detalles perceptualmente irrelevantes a través de la red — la textura de alta frecuencia que un VAE con pérdidas (lossy VAE) podría comprimir. La idea de Rombach: entrenar un VAE una vez (la primera etapa), congelarlo y ejecutar la difusión por completo en el espacio latente de 4 canales de 64×64 (la segunda etapa). La misma U-Net. 1/16 de los píxeles. ~64 veces menos FLOPs para una calidad comparable.

Esta es la receta de Stable Diffusion. SD 1.x / 2.x usó una U-Net de 860M sobre latentes de 64×64×4, SDXL usó una U-Net de 2.6B sobre 128×128×4, SD3 cambió la U-Net por un Diffusion Transformer (DiT) con flow matching. Flux.1-dev (Black Forest Labs, 2024) incluye un DiT-MMDiT de 12B de parámetros. Todos se ejecutan sobre el mismo sustrato de dos etapas.

El Concepto

Dos etapas, entrenadas por separado.

1. Etapa 1 — VAE. Codificador E(x) → z, decodificador D(z) → x. Compresión objetivo: submuestreo de 8× en cada eje espacial + ajuste de canales para que el tamaño latente total sea ~1/16 del conteo de píxeles. Loss = reconstrucción (L1 + LPIPS perceptual) + KL (peso pequeño para que z no sea forzado a ser demasiado gaussiano, ya que no necesitamos un muestreo exacto de z). A menudo se entrena con una loss adversarial para que las imágenes decodificadas sean nítidas.

2. Etapa 2 — difusión en z. Trata z = E(x_real) como los datos. Entrena una U-Net (or DiT) para eliminar el ruido de z_t. En la inferencia: muestrea z_0 mediante difusión, luego x = D(z_0).

Condicionamiento de texto. Dos componentes adicionales. Un codificador de texto congelado (CLIP-L para SD 1.x, CLIP-L+OpenCLIP-G para SD 2/XL, T5-XXL para SD3 and Flux). Una inyección de atención cruzada (cross-attention): cada bloque de la U-Net toma [Q = image features, K = V = text tokens] y los mezcla. Los tokens son la única forma en que el texto influye en la imagen.

La función de pérdida es idéntica a la de la Lección 06. Mismo MSE de DDPM / flow matching sobre el ruido. Solo cambias el dominio de los datos.

Variantes de arquitectura

La tendencia: reemplazar la U-Net con DiT (transformer sobre parches latentes), escalar el codificador de texto (T5 supera a CLIP en adherencia al prompt), aumentar los canales latentes (4 → 16 ofrece más margen para detalles).

Build It

code/main.py apila un "VAE" unidimensional de juguete (codificador + decodificador de identidad, para demostración; un VAE real sería una red convolucional) sobre el DDPM de la Lección 06 y añade condicionamiento de clase con classifier-free guidance. Muestra que la misma loss de difusión funciona ya sea que se ejecute sobre valores unidimensionales puros o sobre valores codificados — el insight fundamental.

Paso 1: codificador/decodificador

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

Un VAE real tiene pesos entrenados. Con fines pedagógicos, este mapa lineal es suficiente para mostrar que la difusión opera sobre z sin preocuparse por el espacio de datos original.

Paso 2: difusión en el espacio z

Mismo DDPM de la Lección 06. Los datos que ve la red son z = E(x). Después de muestrear z_0, decodifica con D(z_0).

Paso 3: classifier-free guidance

Durante el entrenamiento, descarta la etiqueta de clase el 10% del tiempo (reemplázala con un token nulo). En la inferencia, calcula tanto ε_cond como ε_uncond, luego:

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

w = 0 = sin guía (diversidad total), w = 3 = predeterminado, w = 7+ = saturado / demasiado nítido.

Paso 4: condicionamiento de texto (concepto, no código)

Reemplaza la etiqueta de clase con la salida de un codificador de texto congelado. Alimenta el embedding de texto a la U-Net a través de cross-attention:

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

Esta es la única diferencia sustantiva entre un modelo de difusión condicional de clase y Stable Diffusion.

Trampas

• Desajuste de escala del VAE. Los VAE de SD 1.x tienen una constante de escala (scaling_factor ≈ 0.18215) que se aplica después de la codificación. Olvidar esto hace que la U-Net se entrene con latentes que tienen una varianza totalmente incorrecta. Cada checkpoint incluye una.
• Codificador de texto silenciosamente incorrecto. SD3 necesita T5-XXL con >=128 tokens, y la alternativa de usar solo CLIP tiene pérdidas. Siempre verifica use_t5=True o la fidelidad al prompt se desplomará.
• Mezcla de espacios latentes. SDXL, SD3 y Flux usan VAE diferente. Un LoRA entrenado en latentes de SDXL no funcionará en SD3. Hugging Face diffusers 0.30+ se niega a cargar checkpoints que no coincidan.
• CFG demasiado alto. w > 10 produce imágenes saturadas, con aspecto aceitoso, y sobreajusta el prompt a costa de la diversidad. El punto ideal es w = 3-7.
• Fuga de prompts negativos. Un prompt negativo vacío se convierte en el token nulo; un prompt negativo lleno se convierte en el ε_uncond. No son lo mismo; algunos pipelines asumen silenciosamente el nulo por defecto.

Use It

Stacks de producción en 2026:

Ship It

Guarda outputs/skill-sd-prompter.md. La skill toma un prompt de texto + estilo objetivo y genera: modelo + checkpoint, escala CFG, sampler, prompt negativo, resolución, combinación opcional de ControlNet/IP-Adapter y una lista de verificación de QA paso a paso.

Ejercicios

1. Fácil. Ejecuta code/main.py con guía w ∈ {0, 1, 3, 7, 15}. Registra la media de las muestras por clase. ¿En qué w las medias de las clases divergen más allá de las medias de los datos reales?
2. Medio. Cambia el codificador lineal de juguete por un par de codificador/decodificador tanh-MLP con una loss de reconstrucción. Vuelve a entrenar la difusión en los nuevos latentes. ¿Cambia la calidad de la muestra?
3. Difícil. Configura una inferencia real de Stable Diffusion con diffusers: carga sdxl-base, ejecuta 30 pasos de Euler con CFG=7 y toma el tiempo. Ahora cambia a sdxl-turbo con 4 pasos y CFG=0. Mismo tema, diferente calidad — describe qué cambió y por qué.

Términos Clave

Nota de producción: ejecutar Flux-12B en una GPU de consumo de 8GB

la integración de referencia de Flux es la receta canónica para "tengo una GPU de consumo, ¿puedo implementar esto?". El truque es la misma receta de tres perillas que la literatura de inferencia en producción enumera aplicada a un DiT de difusión:

1. Carga escalonada. Flux tiene tres redes que nunca necesitan coexistir en VRAM: el codificador de texto T5-XXL (~10 GB en fp32), CLIP-L (pequeño), el MMDiT de 12B y el VAE. Codifica el prompt primero, elimina los codificadores, carga el DiT, elimina el ruido, elimina el DiT, carga el VAE, decodifica. Las GPU de consumo de 8GB solo admiten una etapa a la vez.
2. Cuantización de 4 bits a través de bitsandbytes. BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16) tanto en el codificador T5 como en el DiT. Reduce la memoria en 8×, y la caída de calidad es imperceptible para texto a imagen según los benchmarks de Aritra (enlazados en el notebook).
3. CPU offload. pipe.enable_model_cpu_offload() intercambia módulos automáticamente entre CPU y GPU a medida que avanza cada paso del forward pass. Añade un 10-20% de latencia pero hace que el pipeline funcione de todos modos.

La contabilidad de memoria es: 10 GB T5 / 8 = 1.25 GB cuantizado, 12 B params × 0.5 bytes = ~6 GB DiT cuantizado, más activaciones. En términos de stas00, este es el extremo de la inferencia TP=1: sin paralelismo de modelo, cuantización máxima. Para producción, ejecutarías TP=2 o TP=4 en H100; para una sola laptop de desarrollo, esta es la receta.

Lectura Adicional

• 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 — Familia Flux.1.
• Hugging Face Diffusers docs — implementación de referencia para cada checkpoint de arriba.