Диффузионные трансформеры и выпрямленный поток

U-Net — не секрет диффузии. Замените его трансформером, поменяйте расписание шума на прямолинейный поток — и у вас внезапно появятся SD3, FLUX и каждая text-to-image-модель 2026 года.

Тип: Изучите + соберите Языки: Python Предварительные требования: фаза 4, урок 10 (диффузия DDPM), фаза 4, урок 14 (ViT), фаза 7, урок 02 (Self-Attention) Время: ~75 минут

Цели обучения

• Проследить эволюцию от U-Net DDPM (урок 10) к Diffusion Transformer (DiT), MMDiT (SD3) и одно- и двухпоточным DiT (FLUX)
• Объяснить выпрямленный поток (rectified flow): почему прямолинейная траектория между шумом и данными позволяет моделям генерировать за 20 шагов вместо 1000
• Реализовать небольшой блок DiT и цикл обучения выпрямленного потока, оба менее чем в 100 строках
• Различать варианты моделей (SD3, FLUX.1-dev, FLUX.1-schnell, Z-Image, Qwen-Image) по архитектуре, числу параметров и лицензированию

Проблема

В уроке 10 был собран DDPM с денойзером U-Net. Этот рецепт доминировал в 2020–2023 годах: U-Net + beta-расписание + функция потерь предсказания шума. Он породил Stable Diffusion 1.5 и 2.1, а также DALL-E 2.

Каждая передовая text-to-image-модель 2026 года ушла дальше. Stable Diffusion 3, FLUX, SD4, Z-Image, Qwen-Image, Hunyuan-Image — ни одна из них не использует U-Net. Они используют Diffusion Transformer (DiT). SD3 и FLUX также заменяют DDPM-расписание шума на выпрямленный поток, который выпрямляет путь от шума к данным и позволяет выполнять вывод за 1–4 шага у вариантов consistency или distilled.

Этот сдвиг важен, потому что именно он сделал генерацию изображений на диффузии управляемой, точно следующей промпту (SD3/SD4 решили отрисовку текста) и достаточно быстрой для продакшена. Понимать DiT + выпрямленный поток — значит понимать стек генеративных изображений 2026 года.

Концепция

От U-Net к трансформеру

• DiT (Peebles & Xie, 2023) — заменяет U-Net трансформером, похожим на ViT, работающим с латентными патчами. Кондиционирование выполняется через адаптивную нормализацию слоёв (adaptive layer norm, AdaLN).
• MMDiT (SD3, Esser et al., 2024) — два потока с отдельными весами для текстовых токенов и токенов изображений, совместно использующие общее внимание.
• FLUX (Black Forest Labs, 2024) — первые N блоков двухпоточные, как в SD3, а поздние блоки конкатенируют данные и совместно используют веса (однопоточный режим), что эффективно при большей глубине.
• Z-Image (2025) — эффективный однопоточный DiT с 6B параметров, ставящий под сомнение принцип «масштабирование любой ценой».

Выпрямленный поток в одном абзаце

DDPM задаёт прямой процесс как зашумлённое SDE, в котором x_t всё сильнее искажается. Изучаемый обратный процесс — второе SDE, решаемое 1000 маленькими шагами.

Выпрямленный поток задаёт прямолинейную интерполяцию между чистыми данными и чистым шумом:

x_t = (1 - t) * x_0 + t * epsilon,     t in [0, 1]

Обучите сеть предсказывать скорость v_theta(x_t, t) = epsilon - x_0 — прямое направление вдоль прямолинейного пути от чистых данных к шуму (dx_t/dt). Во время семплирования вы интегрируете эту скорость в обратном направлении, переходя от шума к данным. Получившееся ODE намного ближе к прямой линии, поэтому для генерации нужно значительно меньше шагов интегрирования.

SD3 называет это Rectified Flow Matching. FLUX, Z-Image и большинство моделей 2026 года используют ту же целевую функцию. Типичный вывод: 20–30 шагов Эйлера (детерминированно) против 50+ шагов DDIM в старом режиме DDPM. Варианты distilled / turbo / schnell / LCM сокращают это до 1–4 шагов.

Кондиционирование AdaLN

DiT кондиционируются по timestep и классу/тексту через адаптивную нормализацию слоя: предсказывайте scale и shift по вектору кондиционирования и применяйте их после LayerNorm. Это гораздо чище, чем модуляция в стиле FiLM в U-Net, и является стандартом в каждом современном DiT.

cond -> MLP -> (scale, shift, gate)
norm(x) * (1 + scale) + shift, then residual add * gate

Текстовые энкодеры в SD3 и FLUX

• SD3 использует три текстовых энкодера: две модели CLIP + T5-XXL. Эмбеддинги конкатенируются и подаются в поток изображений как текстовое кондиционирование.
• FLUX использует один CLIP-L + T5-XXL.
• Варианты Qwen-Image / Z-Image используют собственные текстовые энкодеры, согласованные с их базовыми LLM.

Текстовый энкодер — большая часть причины того, почему SD3/FLUX рассуждают о промптах намного лучше, чем SD1.5. Один только T5-XXL содержит 4.7B параметров.

Classifier-free guidance по-прежнему актуален

Выпрямленный поток меняет семплер, а не кондиционирование. Classifier-free guidance (отбрасывание текста с вероятностью 10% при обучении, смешивание условных и безусловных предсказаний при выводе) работает с выпрямленным потоком совершенно так же. Большинство моделей 2026 года используют масштаб guidance 3.5–5 — меньше, чем 7.5 у SD1.5, потому что модели с выпрямленным потоком изначально точнее следуют промптам.

Consistency, Turbo, Schnell, LCM

Четыре названия одной идеи: дистиллировать медленную многошаговую модель в быструю модель с несколькими шагами.

• LCM (Latent Consistency Model) — обучите ученика предсказывать финальный x_0 из любого промежуточного x_t за один шаг.
• SDXL Turbo / FLUX schnell — модели на 1–4 шага, обученные adversarial diffusion distillation.
• SD Turbo — адаптированные к латентной диффузии Consistency Models в стиле OpenAI.

Продакшен-сервинг любой новой модели поставляет и чекпойнт «полного качества», и вариант «turbo / schnell». Schnell («быстро» по-немецки, соглашение Black Forest Labs) работает за 1–4 шага и подходит для конвейеров реального времени.

Ландшафт моделей в 2026 году

Модель	Размер	Архитектура	Лицензия
Stable Diffusion 3 Medium	2B	MMDiT	SAI Community
Stable Diffusion 3.5 Large	8B	MMDiT	SAI Community
FLUX.1-dev	12B	Double + Single Stream DiT	non-commercial
FLUX.1-schnell	12B	same, distilled	Apache 2.0
FLUX.2	—	iterated FLUX.1	mixed
Z-Image	6B	S3-DiT (Scalable Single-Stream)	permissive
Qwen-Image	~20B	DiT + Qwen text tower	Apache 2.0
Hunyuan-Image-3.0	~80B	DiT	research
SD4 Turbo	3B	DiT + distillation	SAI Commercial

FLUX.1-schnell — открытый стандартный выбор 2026 года. Z-Image — лидер по эффективности. FLUX.2 и SD4 — нынешние вершины качества.

Почему этот фазовый сдвиг важен

DDPM + U-Net работали. DiT + выпрямленный поток работают лучше, быстрее и чище масштабируются. Этот переход параллелен переходу от RNN к трансформерам в NLP: обе архитектуры решали одну задачу, но трансформеры масштабировались и теперь доминируют. Каждая статья 2026 года о генерации изображений, видео или 3D использует денойзер формы DiT и обычно целевую функцию выпрямленного потока. U-Net DDPM теперь прежде всего учебный материал (урок 10).

Соберите это

Шаг 1: блок DiT с AdaLN

import torch
import torch.nn as nn


class AdaLNZero(nn.Module):
    """
    Adaptive LayerNorm with a gate. Predicts (scale, shift, gate) from the conditioning.
    Init such that the whole block starts as identity ("zero init").
    """

    def __init__(self, dim, cond_dim):
        super().__init__()
        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
        self.mlp = nn.Linear(cond_dim, dim * 3)
        nn.init.zeros_(self.mlp.weight)
        nn.init.zeros_(self.mlp.bias)

    def forward(self, x, cond):
        scale, shift, gate = self.mlp(cond).chunk(3, dim=-1)
        h = self.norm(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
        return h, gate.unsqueeze(1)


class DiTBlock(nn.Module):
    def __init__(self, dim=192, heads=3, mlp_ratio=4, cond_dim=192):
        super().__init__()
        self.adaln1 = AdaLNZero(dim, cond_dim)
        self.attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.adaln2 = AdaLNZero(dim, cond_dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, dim * mlp_ratio),
            nn.GELU(),
            nn.Linear(dim * mlp_ratio, dim),
        )

    def forward(self, x, cond):
        h, gate1 = self.adaln1(x, cond)
        a, _ = self.attn(h, h, h, need_weights=False)
        x = x + gate1 * a
        h, gate2 = self.adaln2(x, cond)
        x = x + gate2 * self.mlp(h)
        return x

AdaLNZero начинает как тождественное отображение, потому что веса его MLP инициализируются нулями. Обучение постепенно уводит блок от тождества; это резко стабилизирует глубокие трансформерные диффузионные модели.

Шаг 2: маленький DiT

def timestep_embedding(t, dim):
    import math
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
    args = t[:, None].float() * freqs[None]
    return torch.cat([args.sin(), args.cos()], dim=-1)


class TinyDiT(nn.Module):
    def __init__(self, image_size=16, patch_size=2, in_channels=3, dim=96, depth=4, heads=3):
        super().__init__()
        self.patch_size = patch_size
        self.num_patches = (image_size // patch_size) ** 2
        self.patch = nn.Conv2d(in_channels, dim, kernel_size=patch_size, stride=patch_size)
        self.pos = nn.Parameter(torch.zeros(1, self.num_patches, dim))
        self.time_mlp = nn.Sequential(
            nn.Linear(dim, dim * 2),
            nn.SiLU(),
            nn.Linear(dim * 2, dim),
        )
        self.blocks = nn.ModuleList([DiTBlock(dim, heads, cond_dim=dim) for _ in range(depth)])
        self.norm_out = nn.LayerNorm(dim, elementwise_affine=False)
        self.head = nn.Linear(dim, patch_size * patch_size * in_channels)

    def forward(self, x, t):
        n = x.size(0)
        x = self.patch(x)
        x = x.flatten(2).transpose(1, 2) + self.pos
        t_emb = self.time_mlp(timestep_embedding(t, self.pos.size(-1)))
        for blk in self.blocks:
            x = blk(x, t_emb)
        x = self.norm_out(x)
        x = self.head(x)
        return self._unpatchify(x, n)

    def _unpatchify(self, x, n):
        p = self.patch_size
        h = w = int(self.num_patches ** 0.5)
        x = x.view(n, h, w, p, p, -1).permute(0, 5, 1, 3, 2, 4).reshape(n, -1, h * p, w * p)
        return x

Шаг 3: обучение выпрямленного потока

import torch.nn.functional as F

def rectified_flow_train_step(model, x0, optimizer, device):
    model.train()
    x0 = x0.to(device)
    n = x0.size(0)
    t = torch.rand(n, device=device)
    epsilon = torch.randn_like(x0)
    x_t = (1 - t[:, None, None, None]) * x0 + t[:, None, None, None] * epsilon

    target_velocity = epsilon - x0
    pred_velocity = model(x_t, t)

    loss = F.mse_loss(pred_velocity, target_velocity)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

Сравните с функцией потерь предсказания шума DDPM (урок 10): структура та же, но целевая величина другая. Вместо предсказания шума epsilon мы предсказываем скорость epsilon - x_0, указывающую от данных к шуму вдоль прямолинейной интерполяции.

Шаг 4: семплер Эйлера

Выпрямленный поток — это ODE. Метод Эйлера — самый простой и для хорошо обученной модели выпрямленного потока почти так же точен при 20+ шагах, как солверы более высокого порядка.

@torch.no_grad()
def rectified_flow_sample(model, shape, steps=20, device="cpu"):
    model.eval()
    x = torch.randn(shape, device=device)
    dt = 1.0 / steps
    t = torch.ones(shape[0], device=device)
    for _ in range(steps):
        v = model(x, t)
        x = x - dt * v
        t = t - dt
    return x

20 шагов. На обученной модели это создаёт семплы, сопоставимые с 1000-шаговым DDPM.

Шаг 5: сквозной smoke test

import numpy as np

def synthetic_blobs(num=200, size=16, seed=0):
    rng = np.random.default_rng(seed)
    out = np.zeros((num, 3, size, size), dtype=np.float32)
    yy, xx = np.meshgrid(np.arange(size), np.arange(size), indexing="ij")
    for i in range(num):
        cx, cy = rng.uniform(4, size - 4, size=2)
        r = rng.uniform(2, 4)
        mask = (xx - cx) ** 2 + (yy - cy) ** 2 < r ** 2
        colour = rng.uniform(-1, 1, size=3)
        for c in range(3):
            out[i, c][mask] = colour[c]
    return torch.from_numpy(out)

Обучите TinyDiT на этом наборе синтетических пятен с помощью выпрямленного потока. После 500 шагов результаты семплирования должны напоминать бледные цветные пятна.

Используйте это

Для реальной генерации изображений с FLUX / SD3 / Z-Image пакет diffusers предоставляет единый API для каждой из них:

from diffusers import FluxPipeline, StableDiffusion3Pipeline
import torch

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=torch.bfloat16,
).to("cuda")

out = pipe(
    prompt="a golden retriever surfing a tsunami, hyperrealistic, studio lighting",
    guidance_scale=0.0,           # schnell was trained without CFG
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
out.save("surf.png")

Три строки. FLUX.1-schnell за четыре шага. Замените идентификатор модели на black-forest-labs/FLUX.1-dev, чтобы получить более высокое качество за 20–30 шагов с CFG.

Для SD3:

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-large",
    torch_dtype=torch.bfloat16,
).to("cuda")
out = pipe(prompt, guidance_scale=3.5, num_inference_steps=28).images[0]

Внедрите это

Этот урок создаёт:

• outputs/prompt-dit-model-picker.md — выбирает между SD3, FLUX.1-dev, FLUX.1-schnell, Z-Image и SD4 Turbo с учётом ограничений по качеству, задержке и лицензии.
• outputs/skill-rectified-flow-trainer.md — пишет полный цикл обучения для выпрямленного потока с AdaLN DiT и семплированием Эйлера.

Упражнения

1. (Легко) Обучите приведённый выше TinyDiT на наборе синтетических пятен 500 шагов. Сравните семплы, созданные с 10, 20 и 50 шагами Эйлера.
2. (Средне) Добавьте текстовое кондиционирование, конкатенируя обучаемый эмбеддинг класса с временным эмбеддингом (10 «классов» пятен по цвету). Сгенерируйте для классов 0, 5 и 9 и проверьте соответствие цветов.
3. (Сложно) Вычислите расстояние Фреше (прокси FID) между сгенерированными семплами от версий выпрямленного потока и DDPM сети одинакового размера, обученных на одних данных одинаковое число шагов. Сообщите, какая сходится быстрее.

Ключевые термины

Термин	Как говорят	Что это на самом деле означает
DiT	«Диффузионный трансформер»	Трансформер, заменяющий U-Net как диффузионный денойзер; работает с разбитыми на патчи латентами
AdaLN	«Адаптивная нормализация слоя»	Кондиционирование по timestep/тексту через обучаемые scale, shift, gate после LayerNorm; стандарт каждого современного DiT
MMDiT	«Мультимодальный DiT (SD3)»	Потоки с отдельными весами для текстовых токенов и токенов изображений, совместно использующие общее self-attention
Однопоточный / двухпоточный	«Приём FLUX»	Первые N блоков двухпоточные (отдельные веса по модальности), поздние блоки однопоточные (конкатенация + общие веса) ради эффективности
Выпрямленный поток	«Прямая линия от шума к данным»	Линейная интерполяция между данными и шумом; сеть предсказывает скорость; при выводе нужно меньше шагов ODE
Целевая скорость	«epsilon - x_0»	Регрессионная цель в выпрямленном потоке; указывает от чистых данных к шуму
CFG guidance	«Classifier-free guidance»	Смешивает условные и безусловные предсказания; по-прежнему применяется в моделях выпрямленного потока
Schnell / turbo / LCM	«Дистилляция в 1–4 шага»	Мало-шаговые варианты, дистиллированные из моделей полного качества; подходят для продакшена в реальном времени

Дополнительное чтение

• Scalable Diffusion Models with Transformers (Peebles & Xie, 2023) — статья DiT
• Scaling Rectified Flow Transformers (Esser et al., статья SD3) — MMDiT и выпрямленный поток в масштабе
• Карточка модели и технический отчёт FLUX.1 (Black Forest Labs) — подробности двух- и однопоточной архитектуры
• Z-Image: Efficient Image Generation Foundation Model (2025) — однопоточный DiT с 6B параметров
• Elucidating the Design Space of Diffusion (Karras et al., 2022) — эталонная работа для любого компромисса в проектировании диффузии
• Latent Consistency Models (Luo et al., 2023) — как LCM-LoRA даёт вывод за 4 шага

---

Источник: Diffusion Transformers & Rectified Flow 04.22 — 3D Gaussian Splatting с нуля · Фаза 04 — Компьютерное зрение · 04.24 — SAM 3 и сегментация с открытым словарём · Полный каталог