Entrenamos un Modelo de Texto a Imagen en 24 Horas (Receta Completa)

Introducción: El Reto de 24 Horas

Antes, entrenar un modelo de difusión de texto a imagen competitivo costaba millones de dólares y semanas de GPU. Los tiempos han cambiado. En este post, compartimos la receta completa de un speedrun de 24 horas usando 32 GPUs H200 (aproximadamente $1500 en total). El objetivo fue apilar todos los trucos que realmente funcionan y ver hasta dónde podemos llegar con un presupuesto restringido.

Publicamos todo el código en el repositorio PRX para que puedas reproducir, modificar y extender los experimentos. Este post resume las principales decisiones arquitectónicas y configuraciones de entrenamiento. Para un análisis detallado de cada componente, revisa nuestros posts anteriores (Parte 1 y Parte 2).

Si te interesan las herramientas de IA que simplifican el debugging, echa un vistazo a nuestra guía sobre cómo construir un asistente de troubleshooting para Kubernetes.

La Receta de Entrenamiento: Un Stack de Técnicas Probadas

1. Entrenamiento en Pixel Space con X-Prediction

Adoptamos la formulación x-prediction del artículo Back to Basics: Let Denoising Generative Models Denoise (Li & He, 2025). A diferencia de los modelos de difusión latente que dependen de un VAE, este enfoque entrena directamente en el espacio de píxeles. Esto simplifica el pipeline y elimina el cuello de botella del VAE.

• Tamaño del patch: 32
• Dimensión del bottleneck: 256 en la capa de proyección inicial de tokens
• Escala de resolución: Empezar en 512px → fine-tune en 1024px (saltando 256px)

En 512px, la secuencia tiene aproximadamente 256 tokens. En 1024px, se convierte en 1024 tokens. Con hardware moderno, esto es totalmente manejable.

2. Pérdidas Perceptuales (LPIPS + DINOv2)

Una gran ventaja del entrenamiento en pixel space es que puedes aplicar directamente pérdidas perceptuales clásicas. Agregamos dos objetivos auxiliares a la pérdida estándar de flow-matching:

• Pérdida LPIPS (peso: 0.1) — captura similitud perceptual de bajo nivel.
• Pérdida perceptual basada en DINOv2 (peso: 0.01) — proporciona un alineamiento semántico más fuerte.

Aplicamos estas pérdidas en imágenes completas pooled (no patch-wise) en todos los niveles de ruido. Esto dio resultados consistentemente mejores que el enfoque del artículo original.

3. Enrutamiento de Tokens con TREAD

Para reducir el costo computacional por paso, usamos TREAD (Krause et al., 2025). Selecciona aleatoriamente el 50% de los tokens y los enruta desde el 2º bloque transformer hasta el penúltimo bloque, saltándose las capas intermedias.

Elegimos TREAD sobre SPRINT por su simplicidad. Para compensar la caída de calidad con CFG vanilla, implementamos un esquema de self-guidance que guía usando predicciones condicionales densas vs. enrutadas.

4. Alineación de Representación con REPA + DINOv3

REPA (Yu et al., 2024) alinea las representaciones intermedias del modelo con un teacher preentrenado. Usamos DINOv3 como teacher y aplicamos la pérdida de alineación en el 8º bloque transformer (peso: 0.5).

Como combinamos REPA con enrutamiento TREAD, calculamos la pérdida de alineación solo en los tokens no enrutados para mantener la señal consistente.

5. Optimizador: Muon

Cambiamos Adam por Muon para todos los parámetros 2D (matrices). Los biases, norms y embeddings se quedaron con Adam. La configuración:

Escala de Entrenamiento y Datos

Usamos tres datasets sintéticos:

• Flux generated (1.7M muestras)
• FLUX-Reason-6M (6M muestras)
• Midjourney v6 (1M muestras, re-captionadas con Gemini 1.5 para consistencia)

La escala:

• 512px: 100k pasos, batch size 1024
• 1024px: 20k pasos, batch size 512 (REPA apagado)

Mantuvimos una EMA de los pesos (smoothing=0.999, update cada 10 pasos, inicio en el paso 0).

Ejemplo de Código: Configurando el Loop de Entrenamiento

A continuación un snippet Python simplificado que muestra cómo configurar los componentes clave. El código completo está en el repositorio PRX.

# config.py (simplificado)
import torch
from prx import Trainer, TREADRouter, REPAAlignment, PerceptualLoss

# Configuración del modelo
model_config = {
    "patch_size": 32,
    "bottleneck_dim": 256,
    "num_layers": 24,
    "hidden_dim": 1024,
}

# Componentes de entrenamiento
router = TREADRouter(
    route_ratio=0.5,
    start_block=2,
    end_block=-2  # bloque penúltimo
)

repa = REPAAlignment(
    teacher_name="dinov3",
    loss_weight=0.5,
    apply_at_block=8
)

perceptual_loss = PerceptualLoss(
    lpips_weight=0.1,
    dino_weight=0.01,
    apply_at_all_noise_levels=True
)

# Optimizador
param_groups = [
    {"params": model.muon_params, "optimizer": "muon", "lr": 1e-4, "momentum": 0.95},
    {"params": model.adam_params, "optimizer": "adam", "lr": 1e-4, "betas": (0.9, 0.95)}
]

trainer = Trainer(
    model=model,
    router=router,
    repa=repa,
    perceptual_loss=perceptual_loss,
    param_groups=param_groups,
    batch_size=1024,
    steps_512=100000,
    steps_1024=20000
)

trainer.run()

Resultados y Limitaciones

Después de 24 horas, el modelo produce imágenes coherentes y estéticamente agradables con un fuerte seguimiento del prompt. El fine-tune en 1024px afila los detalles sin romper la composición. Sin embargo, persisten algunos artefactos:

• Fallos de textura ocasionales
• Anatomía extraña en prompts complejos
• El rendimiento disminuye en prompts muy difíciles o fuera de la distribución

Estos problemas son consistentes con subentrenamiento y diversidad limitada de datos, no un fallo estructural en la receta. Con más computación y una cobertura de datos más amplia, la misma configuración debería seguir mejorando de manera predecible.

Lo Que Esto Significa para el Campo

Este speedrun demuestra que el entrenamiento de difusión se ha vuelto notablemente accesible. Combinando entrenamiento en pixel space, enrutamiento eficiente, alineación de representación y guía perceptual ligera, un modelo significativo es alcanzable en un solo día con un presupuesto que habría sido inimaginable hace unos años.

Limitaciones y Precauciones

• Sesgo de datos sintéticos: Todos los datos de entrenamiento fueron generados por IA. Las imágenes del mundo real podrían comportarse de manera diferente.
• Dependencia de hardware: La receta asume GPUs H200 con gran memoria. Los resultados en hardware de consumo varían significativamente.
• Sin filtro de seguridad: El modelo no fue ajustado para seguridad o mitigación de sesgos. Úsalo con cuidado en producción.

Próximos Pasos

Esta carrera de 24 horas es solo el comienzo. El trabajo futuro se centrará en:

1. Escalar con más computación y conjuntos de datos más amplios
2. Mejorar la calidad y diversidad de los captions
3. Explorar más esquemas de self-guidance

Invitamos a la comunidad a construir sobre este trabajo. Todo el código es open-source en el repositorio PRX. Si estás construyendo tu propio pipeline de difusión, también puede ser útil nuestra guía sobre API nativa de popover para tooltips para crear mejores interacciones de UI en tus apps de demostración.

Conclusión

Entrenar un modelo de difusión de texto a imagen competitivo desde cero en 24 horas ya no es ciencia ficción. La combinación de entrenamiento en pixel space, enrutamiento TREAD, alineación REPA y optimizador Muon funciona sinérgicamente para entregar resultados sólidos con un presupuesto modesto.

La lección principal: la investigación en difusión se ha democratizado. No necesitas millones de dólares para contribuir con trabajo significativo. Comienza con el código de PRX, ajusta los botones y descubre lo que puedes lograr.

Gracias por leer. Únete a la discusión de la comunidad en Discord y mantente atento a la próxima ronda de experimentos. 🚀