Diffusers integra FLUX.2: nueva arquitectura para generación de imágenes | Keryc
FLUX.2 llega como la siguiente generación de modelos de imagen de Black Forest Labs, y Diffusers lo incorpora con soporte técnico para inferencia y fine-tuning. ¿Qué cambia respecto a Flux.1 y por qué debería importarte si trabajas con generación multimodal? Aquí te lo explico paso a paso, con ejemplos prácticos para correrlo en hardware real.
Qué es FLUX.2
FLUX.2 es una serie nueva de modelos de generación de imágenes entrenada desde cero con una arquitectura distinta a Flux.1. Es multimodal: funciona como text-to-image y image-to-image, y además acepta hasta 10 imágenes de referencia por prompt, combinando referencias por índice o por lenguaje natural.
A nivel práctico, esto significa que puedes pedirle al modelo que mezcle dos fotos distintas y describir qué elementos tomar de cada una. Pero ojo: cada imagen adicional aumenta el uso de VRAM.
Cambios arquitecturales clave
FLUX.2 conserva la idea general de un multimodal diffusion transformer (MM-DiT) y bloques DiT en paralelo, pero introduce varias optimizaciones y decisiones de diseño importantes.
Usa un solo codificador de texto: Mistral Small 3.1. Esto simplifica el cálculo de embeddings y permite una max_sequence_length de 512.
Los bloques MM-DiT mantienen la separación inicial entre latentes de imagen y texto (doble-stream) y luego los procesan juntos en bloques single-stream. Esto sigue la filosofía de Flux.1 pero con distintos balances entre ambos tipos de bloques.
Se comparte la información de tiempo y guidance (como parámetros de AdaLayerNorm-Zero) entre todos los bloques doble-stream y single-stream, en lugar de tener parámetros de modulación por bloque. Menos parámetros redundantes y más coherencia en la dinámica de difusión.
No hay parámetros de bias en las capas del modelo: ni en atención ni en bloques feedforward. Esto es una decisión de diseño que influye en estabilidad numérica y tamaño del checkpoint.
Los bloques single-stream fusionan proyecciones: además de fusionar la proyección de salida de atención con la del FF, ahora se fusionan las proyecciones QKV de atención con la proyección de entrada del FF, creando un bloque transformer más paralelo.
FLUX.2 usa activación tipo SwiGLU en el MLP en lugar de GELU.
Hay mucho más peso en bloques single-stream que en doble-stream: por ejemplo, la configuración comparativa muestra 8 bloques double-stream vs 48 single-stream en FLUX.2, y una proporción de parámetros en double-stream que cae de ~54% (Flux.1) a ~24% (FLUX.2[dev]-32B).
Nuevo autoencoder y mejor manejo de schedules de timestep dependientes de resolución.
En resumen: mismos principios multimodales, pero rediseño para paralelismo, menor redundancia y mejores trade-offs en compute y memoria.
Inferencia práctica con Diffusers
El modelo combinado (DiT grande + Mistral3 Small) puede superar los 80 GB de VRAM si lo cargas sin optimizaciones. Afortunadamente Diffusers y varias técnicas de engineering permiten ejecutar FLUX.2 en hardware más modesto.
Ejemplo básico con CPU offloading
Este flujo fue probado en H100. Es casi obligatorio activar offloading para que quepa:
from diffusers import Flux2Pipeline
import torch
repo_id = "black-forest-labs/FLUX.2-dev"
pipe = Flux2Pipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload()
image = pipe(
prompt="dog dancing near the sun",
num_inference_steps=50,
guidance_scale=2.5,
height=1024,
width=1024,
).images[0]
Con CPU offload activado, ese setup toma alrededor de 62 GB en una H100. ¿No tienes H100? Sigue leyendo.
Flash Attention 3 para GPUs Hopper
Si tienes una GPU Hopper (por ejemplo H100 o equivalente) puedes acelerar la atención con Flash Attention 3:
Diffusers expone varios backends de atención; busca el que mejor aproveche tu hardware.
Cargar modelos en 4-bit con bitsandbytes (NF4)
Con bitsandbytes puedes cargar transformer y text encoder en 4-bit para ejecutar FLUX.2 en GPUs de ~24 GB o incluso 20 GB VRAM con optimizaciones.
from diffusers import Flux2Pipeline, Flux2Transformer2DModel
from transformers import Mistral3ForConditionalGeneration
import torch
repo_id = "diffusers/FLUX.2-dev-bnb-4bit"
# Cargar transformer y text encoder en 4-bit y usar CPU offload
El ejemplo del blog muestra cómo generar imágenes realistas con un prompt largo y guardar el resultado.
Local + remoto: externalizar el text encoder
Por diseño modular, puedes desplegar el text encoder en un endpoint remoto y mantener el DiT y el VAE locales. Esto reduce mucho el uso de VRAM local.
Idea: el pipeline local recibe prompt_embeds que provienen de un servicio remoto. Ideal si tienes 1 GPU con 18-20 GB de VRAM.
Otra opción es group_offloading, que permite ejecutar en GPUs con 8 GB de VRAM libres, a costa de RAM del sistema (recomiendan 32 GB RAM, o usar low_cpu_mem_usage=True para bajar a 10 GB RAM pero más lento).
Quantizaciones y backends
NF4 (a través de bitsandbytes) es recomendado para inferencia en 4-bit.
Diffusers documenta varios backends de quantización y técnicas para probar cómo afectan la calidad; cuentan con un Space para experimentar interactivamente.
Si quieres comparar resultados visuales, prueba el playground de cuantización que acompaña la release.
Fine-tuning LoRA: cómo hacerlo en GPUs limitadas
FLUX.2 es excelente para fine-tuning LoRA (text-to-image e image-to-image). Pero entrenar requiere la misma preocupación por memoria que la inferencia.
Técnicas que usan en el post y que puedes combinar:
Remote Text Encoder durante el training: pasa --remote_text_encoder para ahorrarte VRAM local de texto.
CPU offloading: --offload para mover VAE y text encoder a CPU y solo traerlos cuando se necesitan.
Latent caching: pre-encodea las imágenes con el VAE y elimina el VAE durante el entrenamiento activo con --cache_latents.
QLoRA / NF4 con bitsandbytes: define un config.json con load_in_4bit y bnb_4bit_quant_type: nf4 y pásalo con --bnb_quantization_config_path.
FP8 training con aceleradores que soporten FP8 (se necesita compute capability >= 8.9) para aprovechar núcleos FP8 y reducir memoria.
Gradient checkpointing y gradient accumulation: reducen memoria a costa de tiempo. Usa --gradient_checkpointing y ajusta --gradient_accumulation_steps.
Optimizadores en 8-bit: --use_8bit_adam con bitsandbytes para reducir memoria del optimizador.
Herramientas alternativas: SimpleTuner o Ostris AI Toolkit si tu GPU es más antigua o quieres entrenadores más ligeros.
Si tu hardware no soporta FP8, usa QLoRA con bitsandbytes y la config JSON para 4-bit NF4.
Buenas prácticas y consejos rápidos
Empieza con prompts y resoluciones moderadas (512 o 768) para iterar rápido y medir consumo.
Haz pruebas con num_inference_steps entre 28 y 50; 28 suele ser buen trade-off calidad/velocidad.
Si trabajas en producción, externalizar el text encoder es una forma efectiva de escalar y reducir costos por instancia GPU.
Combina quantización + offloading + group offload según tu memoria GPU/CPU; algunas técnicas son mutuamente exclusivas, lee la documentación antes de mezclar todo.
Reflexión final
FLUX.2 es un salto interesante: mantiene la idea multimodal de Flux.1, pero reequilibra la arquitectura hacia bloques más paralelos y menos redundancia. Para la comunidad significa más potencia creativa y más opciones de despliegue gracias a las técnicas de offloading y quantización. ¿Tienes 24 GB de GPU o menos? No es una barrera definitiva: con bitsandbytes, remote encoders y group offload puedes experimentar y afinar modelos con recursos accesibles.
Si te interesa, podemos revisar un pipeline concreto para tu hardware y tu caso de uso: imagen comercial, arte generativo o dataset específico. ¿Cuál quieres optimizar primero?
