ترنسفورمرهای دیفیوژن (DiT) و جریان تصحیح‌شده (Rectified Flow)

اگر پیگیر جریان تولید تصویر با هوش مصنوعی باشید، حتماً دیده‌اید که معماری قدیمی U-Net که سال‌ها ستون فقرات Stable Diffusion 1.5 و مدل‌های DDPM اولیه بود، تقریباً کنار رفته است. مدل‌های پیشرو مانند FLUX و Stable Diffusion 3 موتور نویززدای قدیمی را با ترنسفورمرهای دیفیوژن (DiT) تعویض کرده‌اند و به جای صدها گام دیفیوژن تصادفی، از جریان تصحیح‌شده (Rectified Flow) استفاده می‌کنند تا تصاویر را در 20 گام یا حتی کمتر به شکلی پایدار تولید کنند.

درک این معماری‌ها و مبانی ریاضی جریان تصحیح‌شده به شما اجازه می‌دهد از سطح کاربری صرف کتابخانه‌ها فراتر بروید و معماری مدرن مدل‌های چندرسانه‌ای را دقیقاً بشناسید. در این درس پیاده‌سازی بلوک‌های DiT با تکنیک نرمال‌سازی لایه‌ای انطباقی (AdaLN)، حلقه‌ی آموزش پیش‌بینی سرعت و نمونه‌برداری سریع اویلر را به شکل کاملاً عملی بررسی می‌کنیم.

کدهای آموزشی و پایه‌ای این درس برای اجرا روی سیستم محلی و حتی CPU یا پردازنده‌های سازگار شخصی طراحی شده‌اند و نیازی به GPUهای سنگین ابری ندارند. با این حال، توجه داشته باشید که دانلود وزن مدل‌های بزرگی مانند FLUX از مخازنی مثل Hugging Face گاهی به دلیل اختلالات شبکه با کندی مواجه می‌شود که در صورت لزوم می‌توانید از میرورهای موجود کمک بگیرید.

راز دیفیوژن در U-Net نیست. کافی است آن را با یک ترنسفورمر جایگزین کنید و برنامه‌ی نویز را به یک جریان خط مستقیم (straight-line flow) تغییر دهید؛ ناگهان SD3، FLUX و همه‌ی مدل‌های تبدیل متن به تصویر سال 2026 پیش روی شما خواهند بود.

اهداف یادگیری

• ردیابی مسیر تکامل از DDPM مبتنی بر U-Net (درس 10) تا ترنسفورمر دیفیوژن (Diffusion Transformer یا DiT)، معماری MMDiT در SD3 و ساختار تک‌جریانی + دوجریانی در FLUX
• تشریح جریان تصحیح‌شده (rectified flow): این‌که چرا مسیر خط مستقیم میان نویز و داده به مدل‌ها اجازه می‌دهد به جای 1000 گام، نمونه‌گیری را در 20 گام انجام دهند
• پیاده‌سازی یک بلوک DiT جمع‌وجور و یک حلقه‌ی آموزش جریان تصحیح‌شده، هر دو در کمتر از 100 خط کد
• تمایز میان نسخه‌های مختلف مدل‌ها (SD3، FLUX.1-dev، FLUX.1-schnell، Z-Image و Qwen-Image) بر اساس معماری، تعداد پارامترها و نوع مجوز

صورت مسئله

در درس 10 یک مدل DDPM را به همراه نویززدای U-Net پیاده کردیم. این ترکیب در سال‌های 2020 تا 2023 بر دنیای هوش مصنوعی حکمرانی می‌کرد: U-Net + برنامه‌ی بتای نویز (beta schedule) + تابع زیان (loss) پیش‌بینی نویز. خروجی این فرمول مدل‌های موفقی مثل Stable Diffusion 1.5 و 2.1 و DALL-E 2 بود.

اما همه‌ی مدل‌های پیشروی تبدیل متن به تصویر در سال 2026 از این فرمول عبور کرده‌اند. Stable Diffusion 3، FLUX، SD4، Z-Image، Qwen-Image، Hunyuan-Image — هیچ‌کدام از U-Net استفاده نمی‌کنند؛ بلکه بر پایه‌ی ترنسفورمرهای دیفیوژن (DiT) ساخته شده‌اند. علاوه بر این، SD3 و FLUX برنامه‌ی نویز سنتی DDPM را با جریان تصحیح‌شده (rectified flow) جایگزین کرده‌اند؛ رویکردی که مسیر میان نویز و داده را مستقیم می‌کند و استنتاج را در نسخه‌های تقطیرشده یا سازگار (consistency models) به 1 تا 4 گام کاهش می‌دهد.

اهمیت این تحول در این است که تولید تصویر مبتنی بر دیفیوژن دقیقاً به همین دلیل کنترل‌پذیرتر شده، پرامپت‌ها را با دقت بسیار بالاتری اجرا می‌کند (مشکل رندر متن در SD3 و SD4 حل شد) و سرعت آن به حد مناسب برای محیط‌های پروداکشن رسیده است. درک DiT و جریان تصحیح‌شده، کلید شناخت استک تولید تصویر در سال 2026 است.

مفاهیم پایه

از U-Net تا ترنسفورمر

• DiT (پیبلز و شی، 2023) — جایگزینی U-Net با یک ترنسفورمر شبیه به ViT روی پچ‌های نهفته (latent patches). اعمال شرایط (conditioning) از طریق نرمال‌سازی لایه‌ای انطباقی (adaptive layer norm یا AdaLN) انجام می‌شود.
• MMDiT (مدل SD3، اسِر و همکاران، 2024) — دارای دو جریان با وزن‌های مجزا برای توکن‌های متن و تصویر که در یک سازوکار اتنشن (attention) مشترک سهیم هستند.
• FLUX (آزمایشگاه بلک فارست، 2024) — در N بلوک اول مثل SD3 از معماری دوجریانی استفاده می‌کند؛ اما در بلوک‌های بعدی برای کارایی بهتر در عمق‌های بالاتر، توکن‌ها را به هم متصل کرده و وزن‌ها را به اشتراک می‌گذارد (تک‌جریانی).
• Z-Image (2025) — یک مدل DiT تک‌جریانی کارآمد با 6 میلیارد پارامتر که ایده‌ی «بزرگ‌تر کردن مدل به هر قیمتی» را به چالش می‌کشد.

جریان تصحیح‌شده در یک پاراگراف

مدل DDPM فرایند رفت را به شکل یک معادله‌ی دیفرانسیل تصادفی (SDE) نویزدار تعریف می‌کند که در آن x_t به‌مرور تخریب می‌شود. فرایند برگشتِ یادگرفته‌شده نیز یک SDE دوم است که با 1000 گام ریز حل می‌شود.

در مقابل، جریان تصحیح‌شده یک درون‌یابی خط مستقیم بین داده‌ی تمیز و نویز خالص تعریف می‌کند:

$$x_t = (1 - t) \, x_0 + t \, \epsilon, \qquad 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 گام اویلر (Euler - قطعی) انجام می‌شود، در حالی که در چارچوب قدیمی DDPM به بیش از 50 گام DDIM نیاز بود. نسخه‌های تقطیرشده (distilled)، توربو (turbo)، اشنل (schnell) و LCM حتی این عدد را به 1 تا 4 گام رسانده‌اند.

شرطی‌سازی به روش AdaLN

مدل‌های DiT از طریق نرمال‌سازی لایه‌ای انطباقی بر اساس گام زمانی و کلاس یا متن شرطی می‌شوند: مقادیر 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 انکودرهای متن اختصاصی خود را به کار می‌گیرند که با مدل‌های زبانی پایه‌ی آن‌ها هماهنگ شده است.

بخش عمده‌ای از برتری چشمگیر SD3 و FLUX در درک و تحلیل پرامپت‌ها نسبت به SD1.5، به همین انکودر متن برمی‌گردد. مدل T5-XXL به‌تنهایی 4.7 میلیارد پارامتر دارد.

هدایت بدون طبقه‌بندی‌کننده کماکان پابرجاست

جریان تصحیح‌شده نمونه‌بردار (sampler) را تغییر می‌دهد، نه سازوکار شرطی‌سازی را. هدایت بدون طبقه‌بندی‌کننده (classifier-free guidance — یعنی حذف متن با احتمال 10٪ در زمان آموزش و ترکیب پیش‌بینی‌های شرطی و غیرشرطی در زمان استنتاج) در جریان تصحیح‌شده نیز دقیقاً مثل قبل کار می‌کند. اکثر مدل‌های سال 2026 از مقیاس هدایت (guidance scale) بین 3.5 تا 5 استفاده می‌کنند؛ مقداری کمتر از عدد 7.5 در SD1.5، چرا که مدل‌های مبتنی بر جریان تصحیح‌شده ذاتاً پرامپت‌ها را با دقت بسیار بالاتری دنبال می‌کنند.

سازگاری، توربو، اشنل و LCM

این چهار عنوان در واقع نام‌های متفاوتی برای یک ایده‌ی واحد هستند: تقطیر (distil) یک مدل کند و چندگامی به یک مدل سریع با گام‌های اندک.

• LCM (مدل سازگاری نهفته) — آموزش یک مدل شاگرد (student) که می‌تواند از هر نقطه‌ی میانی x_t، خروجی نهایی x_0 را تنها در یک گام پیش‌بینی کند.
• SDXL Turbo / FLUX schnell — مدل‌های 1 تا 4 گامی که با تقطیر دیفیوژن رقابتی (adversarial diffusion distillation) آموزش دیده‌اند.
• SD Turbo — مدل‌های سازگاری (Consistency Models) سبک OpenAI که برای دیفیوژن در فضای نهفته بهینه‌سازی شده‌اند.

در سرویس‌دهی پروداکشن برای هر مدل جدید، معمولاً دو نسخه‌ی مجزا عرضه می‌شود: یک چک‌پوینت «با کیفیت کامل» و یک نسخه‌ی «turbo / schnell». مدل Schnell (در زبان آلمانی به معنای «سریع»؛ قرارداد نام‌گذاری آزمایشگاه بلک فارست) در 1 تا 4 گام اجرا می‌شود و مناسب پایپ‌لاین‌های بلادرنگ (real-time) است.

چشم‌انداز مدل‌ها در سال 2026

در سال 2026، مدل FLUX.1-schnell گزینه‌ی پیش‌فرض متن‌باز است. Z-Image پیشتاز کارایی منابع به شمار می‌رود و FLUX.2 و SD4 پرچم‌داران فعلی کیفیت هستند.

چرا این دگرگونی بنیادین اهمیت دارد؟

ترکیب DDPM + U-Net کار می‌کرد؛ اما فرمول DiT + جریان تصحیح‌شده کیفیت بالاتر و سرعت بیشتری دارد و مقیاس‌پذیری آن بسیار تمیزتر است. این گذار، مشابه گذار از RNNها به ترنسفورمرها در پردازش زبان طبیعی است: هر دو معماری صورت‌مسئله‌ی یکسانی را حل می‌کردند، اما ترنسفورمرها مقیاس‌پذیر بودند و اکنون بر فضا مسلط شده‌اند. تمام مقالات سال 2026 در حوزه‌ی تولید تصویر، ویدیو یا مدل‌های سه‌بعدی از نویززداهایی با معماری 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 کار خود را به عنوان یک تبدیل همانی (identity mapping) شروع می‌کند، زیرا وزن‌های لایه‌ی 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: نمونه‌بردار اویلر (Euler Sampler)

جریان تصحیح‌شده یک 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 در چهار گام. برای رسیدن به کیفیت بالاتر در 20 تا 30 گام به همراه CFG، کافی است شناسه را به black-forest-labs/FLUX.1-dev تغییر دهید.

برای 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 — پیاده‌سازی یک حلقه‌ی آموزش کامل برای جریان تصحیح‌شده با استفاده از DiT مجهز به AdaLN و نمونه‌برداری اویلر.

تمرین‌ها

1. (ساده) مدل TinyDiT بالا را به مدت 500 گام روی مجموعه‌داده‌ی ساختگی لکه‌ها آموزش دهید. نمونه‌های تولیدشده با 10، 20 و 50 گام اویلر را مقایسه کنید.
2. (متوسط) با متصل کردن یک امبدینگ کلاس یادگرفتنی به امبدینگ زمان، شرطی‌سازی متنی را اضافه کنید (10 «کلاس» لکه بر اساس رنگ). با کلاس‌های 0، 5 و 9 نمونه‌گیری کنید و بررسی کنید که آیا رنگ‌ها همخوانی دارند یا خیر.
3. (سخت) فاصله‌ی فرشه (معیاری جایگزین برای FID) را میان نمونه‌های تولیدشده از نسخه‌ی جریان تصحیح‌شده و نسخه‌ی DDPM برای مدلی با ابعاد یکسان محاسبه کنید که روی داده‌های یکسان و به تعداد گام‌های برابر آموزش دیده‌اند. گزارش دهید که کدام مدل سریع‌تر همگرا می‌شود.

اصطلاحات کلیدی

برای مطالعه‌ی بیشتر

• مدل‌های دیفیوژن مقیاس‌پذیر با ترنسفورمرها (پیبلز و شی، 2023) — مقاله‌ی اصلی DiT
• مقیاس‌پذیری ترنسفورمرهای جریان تصحیح‌شده (اسِر و همکاران، مقاله‌ی SD3) — بررسی MMDiT و جریان تصحیح‌شده در ابعاد بزرگ
• شناسنامه و گزارش فنی مدل FLUX.1 (بلک فارست لبز) — جزئیات معماری‌های دوجریانی و تک‌جریانی
• مدل Z-Image: مدل پایه‌ی کارآمد تولید تصویر (2025) — بررسی DiT تک‌جریانی با 6 میلیارد پارامتر
• روشن‌سازی فضای طراحی مدل‌های دیفیوژن (کاراس و همکاران، 2022) — مرجع اصلی موازنه‌های طراحی در مدل‌های دیفیوژن
• مدل‌های سازگاری نهفته (لو و همکاران، 2023) — چگونگی استنتاج در 4 گام با تکنیک LCM-LoRA