Diffusion Transformers & Rectified Flow

Lesson 10 built a DDPM with a U-Net denoiser. That recipe dominated 2020 to 2023: U-Net, beta schedule, noise-prediction loss. It produced Stable Diffusion 1.5, 2.1, and DALL-E 2.

Every 2026 state-of-the-art text-to-image model has moved past it. Stable Diffusion 3, FLUX, SD4, Z-Image, Qwen-Image, and Hunyuan-Image: none of them use a U-Net. They use Diffusion Transformers (DiT) instead. SD3 and FLUX also swap the DDPM noise schedule for rectified flow, which straightens the path from noise to data and enables 1-4 step inference with consistency or distilled variants.

The shift matters because it's the reason diffusion-based image generation got controllable, prompt-accurate (SD3 and SD4 finally solved text rendering), and fast enough to ship. Understand DiT and rectified flow, and you understand the 2026 generative-image stack.

From U-Net to transformer

Four architectures mark this shift, each building on the one before it:

• DiT (Peebles & Xie, 2023): replaces the U-Net with a ViT-like transformer running on latent patches, conditioned via adaptive layer norm (AdaLN).
• MMDiT (SD3, Esser et al., 2024): two streams with separate weights for text and image tokens that share a joint attention layer.
• FLUX (Black Forest Labs, 2024): the first N blocks run double-stream like SD3, then later blocks concatenate and share weights (single-stream) for efficiency at greater depth.
• Z-Image (2025): an efficient single-stream DiT at 6B parameters that challenges the "scale at all costs" assumption.

Rectified flow in one paragraph

DDPM defines the forward process as a noisy SDE (stochastic differential equation: a random process defined one small step at a time) where x_t gets progressively corrupted. The learned reverse process is a second SDE, solved by 1000 small steps.

Rectified flow replaces that with a straight-line interpolation between clean data and pure noise:

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

Train a network to predict the velocity v_theta(x_t, t) = epsilon - x_0: the forward direction along that straight-line path from clean data to noise (dx_t/dt). During sampling, you integrate this velocity backward, stepping from noise toward data. The resulting ODE sits much closer to a straight line than DDPM's curved one, so sampling needs far fewer integration steps.

Analogy: DDPM is a hiker who only trusts short steps and re-checks the map after every one, so reaching the destination takes a thousand cautious steps along a wandering trail. Rectified flow is a hiker who first draws a straight line to the destination on the map, then just walks it: far fewer steps, because the direction never needs correcting.

SD3 calls this Rectified Flow Matching. FLUX, Z-Image, and most 2026 models use the same objective. Typical inference runs 20-30 Euler steps (deterministic), against 50+ steps of DDIM (DDPM's faster deterministic sampler) in the old DDPM regime. Distilled variants (turbo, schnell, LCM) push that down to 1-4 steps.

AdaLN conditioning

DiTs condition on timestep and class/text through adaptive layer norm: predict scale and shift from the conditioning vector, then apply them after LayerNorm. Cleaner than the FiLM-style modulation U-Nets used, and now the default in every modern DiT.

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

Text encoders in SD3 and FLUX

• SD3 uses three text encoders: two CLIP models plus T5-XXL. Their embeddings get concatenated and fed to the image stream as text conditioning.
• FLUX uses one CLIP-L plus T5-XXL.
• Qwen-Image and Z-Image variants use their own in-house text encoders, aligned with their base LLMs.

The text encoder explains a lot of why SD3 and FLUX follow prompts so much better than SD1.5 did. It's also why they can render legible words inside an image at all: T5-XXL's language understanding gives the model enough signal to spell a word out correctly, instead of the scrambled glyphs CLIP-only SD1.5 produced. T5-XXL alone runs 4.7B params.

Classifier-free guidance still holds

Rectified flow changes the sampler, not the conditioning. Classifier-free guidance (drop text with 10% probability during training, mix conditional and unconditional predictions at inference) works identically with rectified flow. Most 2026 models use guidance scale 3.5-5, lower than SD1.5's 7.5, because rectified-flow models follow prompts more tightly by default.

Consistency, Turbo, Schnell, LCM

Four names for the same idea: distil a slow many-step model into a fast few-step model.

• LCM (Latent Consistency Model): trains a student that predicts the final x_0 from any intermediate x_t in one step.
• SDXL Turbo and FLUX schnell: 1-4 step models trained with adversarial diffusion distillation.
• SD Turbo: OpenAI-style Consistency Models adapted to latent diffusion.

Production serving of any new model ships both a "full quality" checkpoint and a "turbo / schnell" variant. Schnell ("fast" in German, Black Forest Labs' naming convention) runs in 1-4 steps and fits real-time pipelines.

Model landscape in 2026

FLUX.1-schnell is the 2026 open-source default. Z-Image is the efficiency leader. FLUX.2 and SD4 are the current quality tips.

Why this phase shift matters

DDPM + U-Net worked. DiT + rectified flow works better, runs faster, and scales more cleanly. The transition parallels the one from RNNs to transformers in NLP: both architectures solved the same problem, but transformers scaled and now dominate. Every 2026 paper on image, video, or 3D generation uses a DiT-shaped denoiser and usually a rectified flow objective. U-Net DDPM is now primarily pedagogical (Lesson 10).

Step 1: A DiT block with AdaLN

The block below implements the scale/shift/gate conditioning described above, with one addition: the modulation MLP starts at zero, so the block begins as an identity function and only drifts away from it as training proceeds.

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 starts as an identity mapping because its MLP weights are initialised to zero. Training nudges the block away from identity; this stabilises deep transformer diffusion models dramatically.

Step 2: A tiny DiT

With the block defined, wiring a full model is just plumbing: patchify the image, add a timestep embedding as conditioning, run the result through depth DiT blocks, then unpatchify back to pixels.

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

Step 3: Rectified flow training

Training needs nothing exotic: sample a random point along the straight line between a real image and noise, ask the model to predict the velocity at that point, and take the mean-squared error against the true velocity.

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()

Compare with DDPM's noise-prediction loss (Lesson 10): same structure, different target. Instead of predicting the noise epsilon, we predict the velocity epsilon - x_0, which points from data to noise along the straight-line interpolation.

Step 4: Euler sampler

Rectified flow is an ODE. Euler's method is the simplest and, for a well-trained rectified-flow model, nearly as accurate as higher-order solvers at 20+ steps.

@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 steps. On a trained model this produces samples comparable to 1000-step DDPM.

Step 5: End-to-end 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)

Train a TinyDiT on this with rectified flow. After 500 steps, sampled outputs should look like faint blobs of colour.

For real image generation with FLUX, SD3, or Z-Image, diffusers ships all three with a unified 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")

Three lines. FLUX.1-schnell in four steps. Swap the model id for black-forest-labs/FLUX.1-dev for higher quality at 20-30 steps with CFG.

For 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]

This lesson produces:

• outputs/prompt-dit-model-picker.md: picks between SD3, FLUX.1-dev, FLUX.1-schnell, Z-Image, SD4 Turbo given quality, latency, and license constraints.
• outputs/skill-rectified-flow-trainer.md: writes a complete training loop for rectified flow with AdaLN DiT and Euler sampling.

1. (Easy) Train the TinyDiT above on the synthetic blob dataset for 500 steps. Compare samples produced with 10, 20, and 50 Euler steps.
2. (Medium) Add text conditioning by concatenating a learned class embedding to the time embedding (10 blob "classes" by colour). Sample with class 0, 5, and 9 and verify colours match.
3. (Hard) Compute the Fréchet distance (FID proxy) between generated samples from rectified-flow and DDPM versions of the same-size network trained on the same data for the same number of steps. Report which converges faster.

• Scalable Diffusion Models with Transformers (Peebles & Xie, 2023): the DiT paper
• Scaling Rectified Flow Transformers (Esser et al., SD3 paper): MMDiT and rectified flow at scale
• FLUX.1 model card and technical report (Black Forest Labs): double-stream plus single-stream details
• Z-Image: Efficient Image Generation Foundation Model (2025): single-stream DiT at 6B
• Elucidating the Design Space of Diffusion (Karras et al., 2022): the reference for every diffusion design trade-off
• Latent Consistency Models (Luo et al., 2023): how LCM-LoRA gives you 4-step inference