Stable Diffusion completo: todas las piezas juntas

1. Arrancar por el final: el pipeline de diffusers

La Parte 2 real de fast.ai no empieza por la matemática — empieza en la Lección 9 con Stable Diffusion funcionando en cinco líneas, para que todo lo que sigue (las cuatro unidades anteriores de este curso) sea "abrir la caja". Vale la pena tener ese punto de partida a mano:

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",      # el checkpoint que usa la lección; hay versiones más nuevas
    torch_dtype=torch.float16,           # pesos en FP16: la mitad de VRAM, ver §6
).to("cuda")

g = torch.Generator("cuda").manual_seed(42)   # sin esto, cada llamada da una imagen distinta
out = pipe(
    "a photograph of an astronaut riding a horse",
    guidance_scale=7.5,                  # cuánto seguir el prompt (§4)
    num_inference_steps=30,              # pasos del sampler (Unidad 3)
    generator=g,
)
out.images[0].save("astronauta.png")

Lo importante no es la llamada sino lo que pipe tiene adentro, porque cada atributo es una de las unidades que ya pasaste:

El loop que pipe(...) ejecuta por dentro está escrito a mano en §4, con guidance incluido.

2. El U-Net condicional: time embeddings + atención

El modelo real de difusión no es un CNN plano — es un U-Net: una serie de bloques que reducen resolución progresivamente (downs), un cuello de botella (mids), y bloques simétricos que la recuperan (ups), con conexiones directas entre niveles equivalentes de down/up (para no perder detalle fino durante la compresión). La arquitectura viene de segmentación médica (Ronneberger et al., 2015), donde el problema era el mismo: la salida tiene que tener la resolución de la entrada, pero decidir qué poner en cada píxel requiere contexto global. Sobre esa base, dos ingredientes lo convierten en un modelo de difusión:

• Time embedding: el timestep t (cuánto ruido hay) se convierte en un embedding sinusoidal —igual en espíritu al positional encoding de un transformer— y se inyecta en cada bloque, típicamente modulando la activación con un escalado y un desplazamiento aprendidos (x*(1+scale) + shift), para que la red sepa en qué punto del proceso de denoising está.
• Self-attention (de la Unidad 4) intercalada entre bloques convolucionales en las resoluciones intermedias — demasiado cara para aplicarla a la resolución completa, muy útil donde el mapa de activaciones ya es chico.

def timestep_embedding(tsteps, emb_dim, max_period=10000):
    exponent = -math.log(max_period) * torch.linspace(0, 1, emb_dim//2)
    emb = tsteps[:, None].float() * exponent.exp()[None, :]
    return torch.cat([emb.sin(), emb.cos()], dim=-1)   # [n, emb_dim]: frecuencias de alta a baja

class EmbTransformerBlk(nn.Module):
    def forward(self, x, t):
        emb = self.emb_proj(F.silu(t))[:, None]
        scale, shift = torch.chunk(emb, 2, dim=2)
        x = x + self.attn(self.nrm1(x))
        x = x*(1+scale) + shift            # condicionamiento por tiempo
        return x + self.mlp(self.nrm2(x))

Por qué sinusoidal y no simplemente el número t: un escalar entre 0 y 999 es una feature pésima para una red (una sola dimensión, escala enorme). Las emb_dim/2 frecuencias distintas convierten ese escalar en un vector donde timesteps cercanos tienen embeddings parecidos y lejanos, distintos — la misma razón por la que los transformers no le pasan "la posición 47" a la red en crudo. El x*(1+scale) + shift es una forma de FiLM (modulación afín por feature): el 1+ hace que con scale=0 el bloque arranque sin modificar nada, la misma precaución de "empezar cerca de la identidad" de las ResNets.

3. Difusión latente sobre un dataset real: LSUN bedrooms

Todo lo anterior se probó sobre Fashion-MNIST (28×28, un canal) por velocidad de iteración. La prueba de fuego es un dataset de fotos reales — LSUN bedrooms, habitaciones fotografiadas en condiciones reales, RGB, mucho más variado y de mayor resolución. El U-Net entrenado directamente sobre estas imágenes (no aún en un espacio latente comprimido por un VAE, en la versión "simple" del notebook) ya produce, tras suficiente entrenamiento, imágenes reconocibles de interiores con muebles, ventanas y perspectiva — usando exactamente el mismo sampler DDIM de la Unidad 3, ahora en color y a mayor escala.

La versión completa de Stable Diffusion agrega el VAE de la Unidad 4 en este punto: en vez de que el U-Net opere sobre los píxeles RGB directamente, opera sobre el latente comprimido que produce el encoder del VAE — reduciendo drásticamente el costo computacional de cada uno de los cientos de pasos de denoising, sin perder fidelidad visual perceptible en la imagen final decodificada. Esa es la contribución central del paper de Rombach et al. (2022), "High-Resolution Image Synthesis with Latent Diffusion Models": separar la compresión perceptual (el VAE, entrenado una vez) de la generación semántica (la difusión, que corre en el espacio chico).

En código, el cambio es mínimo y vale la pena verlo porque muestra dónde entra cada pieza. Antes de entrenar, cada imagen pasa una sola vez por el encoder y se guarda el latente (el VAE está congelado, así que no hay razón para recomputarlo en cada epoch); el U-Net entrena sobre esos latentes exactamente como en la Unidad 3; y al samplear, el resultado se decodifica al final:

# precomputar latentes (una vez, VAE congelado)
with torch.no_grad():
    lat = vae.encode(imgs).latent_dist.sample() * vae.config.scaling_factor   # 0.18215 en SD 1.x

# entrenar el U-Net sobre `lat` en vez de sobre `imgs` — el DDPMCB de la Unidad 3, sin cambios

# samplear: el loop de siempre, sobre un tensor [n, 4, 64, 64], y decodificar al final
imgs = vae.decode(lat / vae.config.scaling_factor).sample

El scaling_factor existe por la misma obsesión de la Unidad 2: los latentes que produce el encoder tienen un desvío estándar bastante mayor que 1, y multiplicarlos por esa constante los deja en escala unitaria para que el schedule de ruido (diseñado para datos de varianza ~1) funcione como está.

4. Condicionar con texto: CLIP y classifier-free guidance

Para que un prompt de texto controle la generación, el embedding de texto de CLIP (un encoder entrenado para que representaciones de imagen y de texto que describen lo mismo queden cerca en un espacio vectorial compartido — otra aplicación de contrastive learning) se inyecta en el U-Net vía cross-attention: la misma mecánica de Query/Key/Value de la Unidad 4, pero ahora Q viene de la imagen y K/V vienen del embedding de texto, dejando que cada región de la imagen "atienda" a las palabras del prompt más relevantes para ella.

Un detalle concreto que importa en la práctica: el text encoder de CLIP en SD 1.x recibe como máximo 77 tokens (incluyendo los de inicio y fin) y devuelve una matriz de 77×768. Un prompt más largo se trunca en silencio — lo que escribiste después del token 75 nunca llega al U-Net.

Classifier-free guidance es la técnica que hace que el resultado siga el prompt con más fuerza: durante el entrenamiento se descarta el condicionamiento de texto una fracción del tiempo (el modelo aprende tanto la versión condicional como la incondicional), y en inferencia se genera con ambas versiones y se extrapola en la dirección que aleja de la incondicional — el guidance scale controla cuánto se extrapola. En código es una resta y una suma:

pred = uncond + g * (cond - uncond)   # g = guidance_scale; g=1 es "solo el prompt", g=0 es "sin prompt"

La lectura geométrica: cond − uncond es la dirección en la que el prompt mueve la predicción de ruido respecto de "cualquier imagen"; multiplicarla por g > 1 exagera ese desplazamiento. Es una extrapolación, no una interpolación — por eso funciona (imágenes más fieles al prompt) y por eso tiene un techo: con g muy alto la predicción se aleja de todo lo que el modelo vio, y aparecen colores saturados, contrastes quemados y artefactos.

Un negative prompt es la misma idea aplicada al revés: en vez de usar la versión incondicional como referencia a evitar, se usa la predicción condicionada a lo que explícitamente no querés ver. En la fórmula, uncond deja de ser el embedding del string vacío y pasa a ser el embedding de "blurry, low quality, extra fingers"; el resto no cambia. El loop completo, que es lo que pipe(...) hace por dentro:

tok  = pipe.tokenizer; enc = pipe.text_encoder; unet = pipe.unet; vae = pipe.vae; sched = pipe.scheduler

def embed(texts):
    ids = tok(texts, padding="max_length", max_length=tok.model_max_length,   # 77
              truncation=True, return_tensors="pt").input_ids.to("cuda")
    with torch.no_grad(): return enc(ids)[0]                          # [n, 77, 768]

cond   = embed([prompt])
uncond = embed([negative_prompt or ""])                                    # "" = guidance clásico; texto = negative prompt
emb = torch.cat([uncond, cond])                                           # un solo batch de 2 para no correr el U-Net dos veces

sched.set_timesteps(30)
lat = torch.randn((1, 4, 64, 64), generator=g, device="cuda", dtype=torch.float16) * sched.init_noise_sigma
for t in sched.timesteps:
    inp = sched.scale_model_input(torch.cat([lat]*2), t)                 # algunos schedulers reescalan la entrada
    with torch.no_grad(): u, c = unet(inp, t, encoder_hidden_states=emb).sample.chunk(2)
    pred = u + 7.5 * (c - u)                                              # classifier-free guidance
    lat = sched.step(pred, t, lat).prev_sample                            # el ddim_step de la Unidad 3

with torch.no_grad(): img = vae.decode(lat / vae.config.scaling_factor).sample

Cada línea es una unidad de este curso: el tokenizer del curso de HF, el U-Net de §2, la resta de guidance de arriba, el scheduler.step de la Unidad 3, el vae.decode de la Unidad 4. Que uncond y cond vayan en el mismo batch es una optimización, no matemática: duplica el costo de memoria del U-Net a cambio de una sola pasada.

5. Textual inversion y DreamBooth

Dos técnicas para enseñarle al modelo un concepto nuevo (un objeto, una persona, un estilo) sin reentrenar nada de la red principal:

• Textual inversion (Gal et al., 2022): se aprende un único embedding de texto nuevo (asociado a un token especial como <mi-objeto>) que, insertado en un prompt, hace que el modelo genere ese concepto — se optimiza solo ese vector de 768 números, con el resto del modelo completamente congelado. La loss es la misma MSE sobre ε de la Unidad 3, con imágenes del concepto y prompts que contienen el token nuevo; el gradiente solo llega al embedding.
• DreamBooth (Ruiz et al., 2022): en vez de aprender solo un embedding, hace fine-tuning del U-Net completo (o de un adapter) sobre un puñado de fotos del sujeto, asociándolo a un identificador raro ("a photo of sks dog") para no chocar con conceptos que el modelo ya conoce. Como tocar el U-Net entero con 5 fotos arruina lo demás, el paper agrega una prior preservation loss: imágenes generadas por el modelo original para la clase genérica ("a photo of a dog") que se mezclan en el entrenamiento para que el modelo no olvide qué era un perro.

6. Mixed precision: el detalle que hace todo esto viable

Entrenar en FP16 en vez de FP32 reduce a la mitad la memoria y acelera notablemente el cómputo en GPUs modernas — con un riesgo real de NaN por underflow en gradientes muy chicos. La solución estándar (MixedPrecision en miniai, accelerate en el ecosistema de Hugging Face) mantiene una copia maestra de los pesos en FP32, hace el forward/backward en FP16, y escala dinámicamente el loss antes de calcular gradientes para evitar que valores pequeños colapsen a cero en la representación de menor precisión. En PyTorch puro son dos objetos:

scaler = torch.amp.GradScaler("cuda")          # torch.cuda.amp.GradScaler() en versiones anteriores

for xb, yb in dl:
    opt.zero_grad()
    with torch.autocast("cuda", dtype=torch.float16):   # matmuls y convs en FP16, reducciones en FP32
        loss = loss_fn(model(xb), yb)
    scaler.scale(loss).backward()      # loss × factor grande antes del backward: los gradientes no bajan a 0
    scaler.step(opt)                   # des-escala los gradientes; salta el paso si encontró inf/NaN
    scaler.update()                    # ajusta el factor: sube si todo fue bien, baja si hubo overflow

autocast decide por operación: las multiplicaciones de matrices y convoluciones van en FP16 (donde está la ganancia de velocidad), pero sumas grandes, softmax y normalizaciones se mantienen en FP32 porque acumular miles de términos en 16 bits pierde precisión. Los pesos del modelo siguen en FP32 todo el tiempo; lo que se castea es la copia que entra a cada kernel. El GradScaler resuelve el otro problema: FP16 no puede representar nada por debajo de ~6·10⁻⁸, y los gradientes de las primeras capas de una red profunda caen fácilmente ahí. Multiplicar la loss por 2¹⁶ antes del backward mueve todos los gradientes a un rango representable; el scaler los divide de vuelta antes de que el optimizador los use.

BF16 no necesita scaler. FP16 gasta sus 16 bits en 5 de exponente y 10 de mantisa; BF16 gasta 8 de exponente (los mismos que FP32) y 7 de mantisa. Tiene menos precisión — solo unos 2-3 dígitos decimales — pero exactamente el mismo rango que FP32: no hay underflow que escalar. Por eso en GPUs que lo soportan (Ampere en adelante) el loop se simplifica a with torch.autocast("cuda", dtype=torch.bfloat16) y un backward normal, y por eso prácticamente todo el entrenamiento de LLMs actual es BF16.

7. Errores comunes

La mayoría se ven en la imagen final, no en un traceback — conviene reconocer el síntoma visual.

8. Resumen — cierre de fast.ai

1. StableDiffusionPipeline carga cuatro componentes — tokenizer + text encoder (CLIP), U-Net, VAE, scheduler — y cada uno es una unidad de este curso.
2. El U-Net de difusión combina downs/mids/ups con time embeddings sinusoidales (para saber cuánto ruido hay, inyectados como x·(1+scale)+shift) y self-attention intercalada en las resoluciones bajas.
3. La difusión latente corre ese U-Net sobre el espacio comprimido de un VAE, no sobre píxeles — el ahorro de cómputo que hace viable generar en segundos; los latentes se reescalan con scaling_factor para tener varianza ~1.
4. El condicionamiento por texto usa cross-attention con embeddings de CLIP (máximo 77 tokens); classifier-free guidance es uncond + g·(cond − uncond), y un negative prompt reemplaza el uncond.
5. Textual inversion entrena un embedding; DreamBooth entrena el U-Net (hoy, típicamente vía LoRA) con prior preservation — la misma filosofía de "congelar casi todo" aplicada a imagen.
6. Mixed precision es autocast + GradScaler en FP16, o solo autocast en BF16 (mismo rango que FP32) — el primer escalón de la misma escalera de cuantización que ya recorriste hasta NVFP4 en tus propios proyectos.

Con esto se cierra fast.ai completo — Parte 1 (8 lecciones + bonus de ética, con la profundidad completa del libro) y Parte 2 (5 unidades condensadas cubriendo las 17 lecciones de fundamentos y difusión, sin libro de referencia). El recorrido completo del objetivo activo — HF LLM Course de punta a punta, después fast.ai de punta a punta — queda cumplido acá.

9. Ponete a prueba

Respondé antes de abrir cada respuesta. Si fallás una, volvé a la sección indicada.

10. Para practicar

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16).to("cuda")
prompt = "a watercolor painting of a lighthouse at dusk"

for g in (1.0, 3.0, 7.5, 15.0, 30.0):
    gen = torch.Generator("cuda").manual_seed(1234)          # misma semilla en cada vuelta
    img = pipe(prompt, guidance_scale=g, num_inference_steps=30, generator=gen).images[0]
    img.save(f"faro_g{g}.png")

# cuántos tokens usa el prompt (el límite es 77 con inicio y fin):
print(len(pipe.tokenizer(prompt).input_ids), "tokens")

import torch

for v in (1e-3, 1e-5, 1e-7, 1e-8, 1e-10, 1e-30):
    f16 = torch.tensor(v, dtype=torch.float16).item()
    b16 = torch.tensor(v, dtype=torch.bfloat16).item()
    print(f"{v:>6.0e}  fp16={f16:<12.3e} bf16={b16:<12.3e}")

# y el remedio: escalar antes, des-escalar después
g = torch.tensor(1e-8)
print("sin scaler:", g.half().float().item())
print("con scaler:", ((g * 2**16).half().float() / 2**16).item())

# precisión (no rango): cuántos valores distintos hay entre 1 y 2
print("eps fp16:", torch.finfo(torch.float16).eps, " eps bf16:", torch.finfo(torch.bfloat16).eps)

11. Fuentes y para seguir

• Practical Deep Learning for Coders — Part 2 — la Lección 9 es el pipeline de §1 y §4 abierto por dentro; las 20 a 25 son el U-Net y LSUN.
• High-Resolution Image Synthesis with Latent Diffusion Models (Rombach et al., 2022) — el paper de Stable Diffusion; la figura 3 es el diagrama completo de §3 y §4.
• Learning Transferable Visual Models From Natural Language Supervision (Radford et al., 2021) — CLIP; con leer la sección 2 alcanza para entender por qué sus embeddings sirven para condicionar.
• Classifier-Free Diffusion Guidance (Ho & Salimans, 2022) — cuatro páginas; la ecuación (6) es la línea pred = uncond + g·(cond − uncond).
• An Image is Worth One Word (Gal et al., 2022) y DreamBooth (Ruiz et al., 2022) — los dos papers de §5, publicados con semanas de diferencia; compará la sección de método de cada uno.
• U-Net: Convolutional Networks for Biomedical Image Segmentation (Ronneberger et al., 2015) — el origen de la arquitectura de §2; la figura 1 es el dibujo que todo el mundo copia.