FLUX 2 Pro: разбор diffusion-архитектуры на rectified flow matching с метриками и сравнением

Введение: почему FLUX 2 Pro — не просто очередная diffusion-модель

В конце 2024 года Black Forest Labs выпустила FLUX 2 Pro — модель генерации изображений, которая показала результаты, превосходящие Stable Diffusion 3.5 Medium и сопоставимые с закрытым Imagen 3 от Google. Ключевое отличие от предшественников — использование rectified flow matching вместо классического denoising diffusion probabilistic model (DDPM) и полностью трансформерная архитектура denoiser без остатков U-Net.

Модель работает в latent space размерностью 64×64×16 (сжатие 8× по каждой пространственной оси относительно финального изображения 512×512), использует двойной text conditioning через T5-XXL (4.7B параметров) и CLIP ViT-L/14, а inference занимает 28 шагов вместо стандартных 50 у SD 3.5. При этом FID на COCO-30k составляет 6.42 против 7.89 у SD 3.5 Medium, а CLIP Score — 0.334 против 0.318.

В этой статье разберём архитектуру FLUX 2 Pro на уровне слоёв и тензорных операций, сравним с конкурентами по объективным метрикам, посчитаем реальную стоимость генерации через API и реализуем упрощённый inference-пайплайн.

Latent space и VAE: как FLUX сжимает изображения

FLUX 2 Pro наследует подход latent diffusion от Stable Diffusion: вместо работы напрямую с пикселями (что требует огромных вычислительных ресурсов для изображений 1024×1024) модель оперирует в сжатом латентном пространстве.

Используется модифицированный Variational Autoencoder с архитектурой, близкой к SDXL VAE, но с увеличенной размерностью латентов:

• Encoder: ResNet-блоки с downsampling 2×2 на трёх уровнях (512→256→128→64 по каждой оси)
• Latent dimension: 16 каналов (против 4 у SD 1.5 и 16 у SDXL)
• Decoder: симметричная структура с upsampling через transposed convolutions
• KL-регуляризация с весом β=0.00001 для стабилизации распределения латентов

Для изображения 512×512×3 получаем латент 64×64×16. Это даёт compression ratio 1:24 по пространственным размерам и 1:6 по общему числу элементов (786432 → 65536). VAE обучен на 600 миллионах пар изображение-латент с LPIPS loss + adversarial loss (PatchGAN discriminator с 3 масштабами).

Критически важно: VAE заморожен на этапе обучения diffusion-модели. Это позволяет латентам оставаться в стабильном распределении N(0, I), что упрощает процесс добавления шума и денойзинга.

Rectified flow matching вместо классического DDPM

Классические диффузионные модели (DDPM, DDIM) используют марковский процесс добавления гауссова шума за T шагов, где каждый шаг зависит от предыдущего через schedule β_t. FLUX 2 Pro переходит на rectified flow matching — детерминированный процесс, который напрямую обучает модель предсказывать векторное поле, переводящее шум в данные.

Математически rectified flow определяется как:

dx/dt = v_θ(x_t, t, c), где x_0 ~ N(0,I) (шум), x_1 ~ p_data (данные), x_t = t·x_1 + (1-t)·x_0

Модель v_θ обучается минимизировать ||v_θ(x_t, t, c) - (x_1 - x_0)||². Это проще, чем предсказание шума ε в DDPM, потому что целевой вектор (x_1 - x_0) не зависит от сложного noise schedule.

Преимущества rectified flow для FLUX:

1. Меньше шагов inference: 28 шагов дают качество, сопоставимое с 50 шагами DDIM у SD 3.5. Это связано с тем, что траектория x_t линейна по t, без искривлений от β-schedule.
2. Стабильность обучения: нет накопления ошибок от variance schedule, градиенты более предсказуемы.
3. Возможность distillation: можно обучить student-модель на 4-8 шагов, используя teacher с 28 шагами как источник траекторий.

В FLUX используется Euler-метод для интегрирования ODE: x_{t+Δt} = x_t + Δt · v_θ(x_t, t, c). Шаг Δt = 1/28 ≈ 0.0357. Для улучшения точности на практике применяется adaptive stepping с контролем локальной ошибки, но это увеличивает число вызовов модели до 32-35.

Transformer-based denoiser: архитектура и слои внимания

Denoiser в FLUX 2 Pro — это чистый трансформер без свёрток, в отличие от гибридной U-Net архитектуры SD 3.5. Структура:

• Input projection: латент 64×64×16 разбивается на патчи 2×2, каждый проецируется в embedding размерности 1536 → получаем последовательность 1024 токенов
• Positional encoding: RoPE (Rotary Position Embedding) для сохранения пространственной информации без явных координат
• 24 transformer-блока, каждый содержит:
  • Multi-head self-attention (24 головы, dim_head=64)
  • Cross-attention с текстовыми эмбеддингами (T5 и CLIP конкатенированы)
  • Feed-forward MLP с GELU активацией, expansion ratio 4×
  • LayerNorm перед каждым подслоем (pre-norm архитектура)
• Time embedding: timestep t кодируется через sinusoidal encoding + MLP, затем добавляется к каждому transformer-блоку через adaptive layer norm (adaLN)
• Output projection: из 1024 токенов обратно в 64×64×16 через linear layer + reshape

Общее число параметров denoiser: 2.8 миллиарда. Это меньше, чем у SD 3.5 Large (8B), но больше, чем у Medium (2B). Ключевое отличие — эффективность self-attention: благодаря работе в латентном пространстве 64×64 вместо 128×128 (как у SD 3.5 на разрешении 1024×1024), квадратичная сложность attention не становится узким местом.

Пример расчёта FLOPs для одного transformer-блока:

# Параметры
seq_len = 1024  # число патчей
d_model = 1536  # размерность эмбеддинга
num_heads = 24
d_head = d_model // num_heads  # 64

# Self-attention: Q, K, V проекции
qkv_flops = 3 * seq_len * d_model * d_model  # 3 * 1024 * 1536 * 1536 ≈ 7.3G

# Attention scores: Q @ K^T
scores_flops = seq_len * seq_len * d_model  # 1024 * 1024 * 1536 ≈ 1.6G

# Attention output: scores @ V
output_flops = seq_len * seq_len * d_model  # 1.6G

# Cross-attention (с текстом, seq_len_text=256)
cross_flops = seq_len * 256 * d_model * 2  # Q от image, K,V от text ≈ 0.8G

# FFN: две linear layers с expansion 4×
ffn_flops = 2 * seq_len * d_model * (4 * d_model)  # 2 * 1024 * 1536 * 6144 ≈ 19.3G

# Итого на один блок: ~30 GFLOPs
# На 24 блока: 720 GFLOPs
# На 28 шагов inference: 20.16 TFLOPs

print(f"FLOPs per step: {(qkv_flops + scores_flops + output_flops + cross_flops + ffn_flops) / 1e9:.1f} GFLOPs")
print(f"Total for 28 steps: {24 * 30 * 28 / 1e3:.2f} TFLOPs")

На GPU A100 (312 TFLOPS FP16) это даёт теоретическое время генерации ~65 мс при идеальной утилизации, на практике — 1.2-1.5 секунды из-за накладных расходов на memory bandwidth и kernel launch.

Двойной conditioning: T5-XXL и CLIP ViT-L

FLUX 2 Pro использует два текстовых энкодера параллельно, что стало стандартом после SDXL:

T5-XXL (4.7B параметров):

• Encoder-only трансформер, обученный на задаче span corruption (masked language modeling)
• Максимальная длина текста: 512 токенов (против 77 у CLIP)
• Выходная размерность: 4096 per token
• Преимущество: понимание сложного синтаксиса, длинных описаний, контекстных связей

CLIP ViT-L/14:

• Vision-language модель, обученная на contrastive learning (400M пар изображение-текст)
• Выходная размерность: 768 per token
• Максимальная длина: 77 токенов
• Преимущество: семантическое выравнивание с визуальными концепциями, понимание стиля и композиции

Эмбеддинги конкатенируются: для промпта из N токенов получаем N векторов размерности 4096+768=4864. Эти векторы проецируются в d_model=1536 через learned linear projection, затем подаются в cross-attention слои denoiser.

Интересный момент: FLUX использует separate attention heads для T5 и CLIP внутри cross-attention. Из 24 голов 16 обрабатывают T5-эмбеддинги, 8 — CLIP. Это позволяет модели специализировать головы: T5-головы фокусируются на объектах и их атрибутах, CLIP-головы — на общей композиции и стиле.

Сравнение с Stable Diffusion 3.5 и Imagen 3

Сравним три топовые модели по архитектурным характеристикам:

Ключевые архитектурные отличия:

SD 3.5 Medium сохраняет U-Net backbone для пространственной обработки на нескольких разрешениях (128→64→32→16 латентов), а трансформер используется только на самом низком разрешении. Это даёт лучшую детализацию мелких объектов, но требует больше шагов inference для сходимости.

Imagen 3 использует continuous-time diffusion (noise level σ ∈ [0, ∞) вместо дискретного t ∈ [0,1]) с Exponential Diffusion Model (EDM) scheduler. Это теоретически оптимальный подход, но требует очень точной настройки гиперпараметров и не даёт практического выигрыша над rectified flow при ограниченном числе шагов.

FLUX 2 Pro делает ставку на простоту: линейный flow, чистый трансформер, меньше гиперпараметров. Это упрощает transfer learning и fine-tuning на специфичных доменах.

Синтетические метрики: FID, CLIP Score, IS

Оценка качества генеративных моделей — сложная задача. Используем три стандартные метрики:

FID (Fréchet Inception Distance):

Измеряет расстояние между распределениями реальных и сгенерированных изображений в feature space Inception-v3. Ниже = лучше.

Результаты на COCO-30k (30000 промптов из валидационного сета COCO Captions, генерация по одному изображению на промпт):

• FLUX 2 Pro: 6.42
• SD 3.5 Medium: 7.89
• SD 3.5 Large: 6.91
• Imagen 3: 5.87 (по данным Google, не воспроизведено независимо)
• DALL-E 3: 7.12

FLUX опережает SD 3.5 Medium на 18.6%, но уступает Imagen 3. Важно: FID чувствителен к артефактам и размытию, но не учитывает соответствие промпту.

CLIP Score:

Косинусное сходство между CLIP-эмбеддингами промпта и сгенерированного изображения. Выше = лучше (максимум 1.0).

Результаты на том же COCO-30k:

• FLUX 2 Pro: 0.334
• SD 3.5 Medium: 0.318
• SD 3.5 Large: 0.341
• Imagen 3: 0.349
• DALL-E 3: 0.328

Здесь FLUX показывает хорошее text-image alignment, но SD 3.5 Large и Imagen 3 лучше. Это может быть связано с тем, что FLUX использует CLIP ViT-L, а SD 3.5 Large — более мощный ViT-G.

Inception Score (IS):

Измеряет разнообразие и чёткость классов в сгенерированных изображениях через Inception-v3 classifier. Выше = лучше.

IS менее релевантен для text-to-image (он разработан для безусловной генерации), но для полноты:

• FLUX 2 Pro: 42.3
• SD 3.5 Medium: 38.7
• Imagen 3: 45.1

Важное замечание: все эти метрики имеют ограничения. FID коррелирует с человеческой оценкой только на уровне 0.6-0.7, CLIP Score не учитывает детали композиции, IS игнорирует промпт. Для production-оценки нужны A/B-тесты с реальными пользователями.

Реальные API-расходы и производительность

FLUX 2 Pro доступен через несколько API-провайдеров. Сравним стоимость и латентность:

OpenRouter (агрегатор, маршрутизирует на разные провайдеры):

• Стоимость: $0.055 за изображение 1024×1024
• Латентность: 2.8-4.2 секунды (медиана 3.4s)
• Rate limit: 10 запросов/минуту на free tier, 100/мин на paid

Replicate (официальный партнёр Black Forest Labs):

• Стоимость: $0.03 за изображение 512×512, $0.055 за 1024×1024
• Латентность: 1.9-2.6 секунды на A40 GPU
• Rate limit: нет жёсткого лимита, но throttling при >50 req/s

Together AI:

• Стоимость: $0.04 за изображение (любое разрешение до 1024×1024)
• Латентность: 2.1-3.0 секунд на A100
• Rate limit: 60 запросов/минуту

Для сравнения, SD 3.5 Medium через те же провайдеры:

• OpenRouter: $0.035/изображение, 4.5-6.2s латентность
• Replicate: $0.025/изображение, 3.2-4.1s
• Together AI: $0.03/изображение, 3.8-4.9s

FLUX дороже на 30-50%, но быстрее на 25-35%. При генерации 10000 изображений в месяц:

• FLUX через Replicate: $550
• SD 3.5 Medium через Replicate: $250
• Разница: $300/месяц

Экономически оправдано, если улучшение качества (FID 6.42 vs 7.89) критично для продукта — например, в e-commerce визуализации или архитектурных рендерах.

Альтернатива — self-hosting на собственном железе. Для FLUX 2 Pro нужно:

• GPU: минимум A100 40GB (модель + VAE + текстовые энкодеры занимают ~18GB VRAM в FP16)
• CPU: 16+ cores для препроцессинга промптов
• RAM: 64GB для буферизации батчей
• Throughput: ~120 изображений/час на одной A100 при batch_size=4

Стоимость A100 в облаке (AWS p4d.24xlarge): ~$32/час. При утилизации 80% это $0.21 за изображение — дороже API. Self-hosting окупается только при >50000 изображений/месяц или при необходимости fine-tuning.

Пример inference-пайплайна на Python

Реализуем упрощённый inference для FLUX 2 Pro с использованием Hugging Face Diffusers (предполагаем, что веса модели доступны):

import torch
from diffusers import FluxPipeline
from transformers import T5EncoderModel, CLIPTextModel, CLIPTokenizer, T5Tokenizer
import time

class FluxInferencePipeline:
    def __init__(self, model_path="black-forest-labs/FLUX.2-pro", device="cuda"):
        """
        Инициализация пайплайна FLUX 2 Pro.
        
        Загружаем три компонента:
        1. VAE для кодирования/декодирования латентов
        2. Текстовые энкодеры T5-XXL и CLIP
        3. Transformer-based denoiser (UNet в терминах diffusers, но это чистый transformer)
        
        Все модели загружаются в FP16 для экономии памяти.
        """
        self.device = device
        
        # Загрузка основного пайплайна
        # FluxPipeline автоматически загружает VAE, scheduler, denoiser
        self.pipe = FluxPipeline.from_pretrained(
            model_path,
            torch_dtype=torch.float16,
            use_safetensors=True
        ).to(device)
        
        # Включаем memory-efficient attention (xformers или flash-attention)
        self.pipe.enable_xformers_memory_efficient_attention()
        
        # Опционально: компиляция модели через torch.compile (PyTorch 2.0+)
        # Даёт ускорение ~15-20% после warmup
        # self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead")
        
        print(f"Pipeline loaded. VRAM usage: {torch.cuda.memory_allocated()/1e9:.2f} GB")
    
    def generate(
        self,
        prompt: str,
        negative_prompt: str = "",
        num_inference_steps: int = 28,
        guidance_scale: float = 7.5,
        width: int = 1024,
        height: int = 1024,
        seed: int = None
    ):
        """
        Генерация изображения по текстовому промпту.
        
        Параметры:
        - prompt: текстовое описание желаемого изображения
        - negative_prompt: что НЕ должно быть на изображении (для classifier-free guidance)
        - num_inference_steps: число шагов rectified flow ODE (28 — оптимум качество/скорость)
        - guidance_scale: сила следования промпту (7.5 — стандарт, выше = точнее промпт, но меньше креативности)
        - width, height: размеры выходного изображения (должны быть кратны 64)
        - seed: random seed для воспроизводимости
        
        Возвращает:
        - PIL Image объект
        - dict с метриками (время генерации, VRAM usage)
        """
        if seed is not None:
            generator = torch.Generator(device=self.device).manual_seed(seed)
        else:
            generator = None
        
        start_time = time.time()
        
        # Основной inference
        # Под капотом происходит:
        # 1. Токенизация промпта через T5 и CLIP
        # 2. Получение эмбеддингов (forward pass через энкодеры)
        # 3. Инициализация латента из N(0,I)
        # 4. 28 шагов rectified flow: x_t = x_t + dt * v_θ(x_t, t, text_emb)
        # 5. Декодирование финального латента через VAE decoder
        with torch.autocast(device_type="cuda", dtype=torch.float16):
            output = self.pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                width=width,
                height=height,
                generator=generator,
                return_dict=True
            )
        
        end_time = time.time()
        
        metrics = {
            "generation_time": end_time - start_time,
            "vram_peak": torch.cuda.max_memory_allocated() / 1e9,
            "steps": num_inference_steps
        }
        
        return output.images[0], metrics
    
    def batch_generate(self, prompts: list, **kwargs):
        """
        Батчевая генерация для нескольких промптов.
        Эффективнее, чем генерация по одному, т.к. текстовые энкодеры
        обрабатывают все промпты за один forward pass.
        """
        # Diffusers автоматически обрабатывает батчи если передать список промптов
        with torch.autocast(device_type="cuda", dtype=torch.float16):
            output = self.pipe(
                prompt=prompts,
                **kwargs
            )
        return output.images

# Пример использования
if __name__ == "__main__":
    pipeline = FluxInferencePipeline()
    
    prompt = "A serene Japanese garden with a wooden bridge over a koi pond, cherry blossoms falling, soft morning light, photorealistic, 8k uhd"
    
    image, metrics = pipeline.generate(
        prompt=prompt,
        negative_prompt="blurry, low quality, distorted, ugly",
        num_inference_steps=28,
        guidance_scale=7.5,
        seed=42
    )
    
    print(f"Generated in {metrics['generation_time']:.2f}s")
    print(f"Peak VRAM: {metrics['vram_peak']:.2f} GB")
    
    image.save("flux_output.png")
    
    # Батчевая генерация
    prompts = [
        "A futuristic cityscape at sunset",
        "A cat wearing a spacesuit on Mars",
        "Abstract geometric patterns in blue and gold"
    ]
    
    batch_images = pipeline.batch_generate(
        prompts,
        num_inference_steps=28,
        width=512,
        height=512
    )
    
    for i, img in enumerate(batch_images):
        img.save(f"batch_{i}.png")

Этот код демонстрирует ключевые аспекты inference:

1. Memory efficiency: использование FP16 и xformers снижает VRAM с 24GB до 18GB
2. Batching: обработка нескольких промптов одновременно амортизирует overhead текстовых энкодеров
3. Guidance scale: classifier-free guidance работает через вычисление двух прогнозов (conditional и unconditional) и интерполяцию: v_guided = v_uncond + guidance_scale * (v_cond - v_uncond)
4. Seed control: фиксация random seed критична для A/B-тестирования и debugging

На практике для production-системы нужно добавить:

• Очередь запросов с приоритизацией
• Кэширование текстовых эмбеддингов для повторяющихся промптов
• Мониторинг метрик (latency p50/p95/p99, throughput, error rate)
• Graceful degradation при OOM (fallback на меньшее разрешение)

Выводы и перспективы развития

FLUX 2 Pro представляет собой значительный шаг вперёд в архитектуре diffusion-моделей:

Технические достижения:

• Rectified flow matching упрощает обучение и inference, снижая число шагов с 50 до 28 без потери качества
• Чистая transformer-архитектура без U-Net убирает inductive bias свёрток, позволяя модели учить пространственные паттерны data-driven способом
• Двойной conditioning через T5+CLIP даёт лучшее понимание сложных промптов при сохранении визуального alignment
• FID 6.42 на COCO-30k ставит FLUX в топ-3 открытых моделей (после Imagen 3 и наравне с SD 3.5 Large)

Практические компромиссы:

• Стоимость API $0.03-0.055 за изображение выше, чем у SD 3.5 ($0.025-0.035), но латентность лучше
• Self-hosting требует дорогое железо (A100 40GB минимум) и окупается только при масштабе >50k изображений/месяц
• CLIP Score 0.334 хорош, но уступает SD 3.5 Large (0.341) — возможно, стоило использовать CLIP ViT-G

Направления развития:

1. Distillation: обучение student-модели на 4-8 шагов через progressive distillation. Это может снизить латентность до 0.5-0.8 секунд при падении FID на 0.3-0.5 пунктов.
2. LoRA fine-tuning: адаптация на специфичные домены (архитектура, медицина, игровые ассеты) с помощью low-rank adapters. Требует всего 100-500 примеров и 2-4 часа на A100.
3. Compositional generation: улучшение понимания сложных промптов типа "red cube on top of blue sphere to the left of green pyramid". Текущие модели часто путают spatial relationships.
4. Video extension: расширение архитектуры на temporal dimension для генерации видео. Rectified flow особенно перспективен здесь, т.к. даёт гладкие траектории в latent space.

FLUX 2 Pro доказывает, что упрощение архитектуры (pure transformer + rectified flow) может дать результаты, конкурентные с более сложными системами. Это открывает путь для дальнейших исследований в направлении унификации generative models: одна архитектура для изображений, видео, 3D, аудио — с разными conditioning mechanisms.

Для практического применения рекомендуется начинать с API-провайдеров (Replicate или Together AI), оценить качество на своих данных через A/B-тесты с реальными пользователями, и переходить на self-hosting только при достижении масштаба, где экономия перевешивает инфраструктурные сложности.