Flow Matching 与整流模型：下一代生成模型

# Flow Matching 的核心数学：条件概率路径
import torch
import torch.nn as nn

class ConditionalFlowMatcher:
    """最简单的条件流匹配器：线性插值路径"""
    
    def __init__(self, sigma=0.0):
        self.sigma = sigma  # 可选的噪声注入
    
    def sample_location_and_conditional_flow(self, x0, x1):
        """
        采样时间点 t 和条件路径上的点 x_t
        x0: 噪声 (batch, dim)
        x1: 数据 (batch, dim)
        """
        t = torch.rand(x0.shape[0], device=x0.device)
        t = t[:, None]  # (batch, 1)
        
        # 线性插值: x_t = t * x1 + (1 - t) * x0
        # 这定义了一条从 x0 到 x1 的直线
        xt = t * x1 + (1 - t) * x0
        
        # 条件速度场: u_t(x|x1) = x1 - x0
        # 对于线性插值，速度是常数（方向指向目标）
        ut = x1 - x0
        
        return t, xt, ut

# 训练循环骨架
def train_flow_matching(model, data_loader, optimizer, n_epochs=100):
    matcher = ConditionalFlowMatcher()
    
    for epoch in range(n_epochs):
        for x1 in data_loader:  # x1 是真实数据
            x0 = torch.randn_like(x1)  # 从高斯噪声采样
            t, xt, ut = matcher.sample_location_and_conditional_flow(x0, x1)
            
            # 神经网络预测速度场
            vt = model(xt, t)
            
            # Flow Matching 损失: MSE 预测 vs 真实速度
            loss = nn.functional.mse_loss(vt, ut)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
    return model

# 采样：从噪声 ODE 积分到数据
import torch
from torchdiffeq import odeint

@torch.no_grad()
def sample_flow_matching(model, noise_shape, device, 
                         method="dopri5", rtol=1e-5, atol=1e-5):
    """
    使用 ODE 求解器从噪声生成样本
    
    method 选项:
    - "euler": 最简单，步数固定
    - "dopri5": Dormand-Prince 5(4)，自适应步长
    - "midpoint": 二阶 Runge-Kutta
    """
    x0 = torch.randn(noise_shape, device=device)
    
    def velocity_fn(t, x):
        return model(x, t.unsqueeze(0))
    
    # 从 t=0 积分到 t=1
    # dopri5 会自动调整步数以控制误差
    trajectory = odeint(
        velocity_fn, x0,
        torch.tensor([0.0, 1.0], device=device, dtype=torch.float64),
        method=method, rtol=rtol, atol=atol
    )
    
    return trajectory[-1]  # 返回 t=1 的终点

# 对比不同 ODE 求解器的速度-质量权衡
solvers = {
    "euler_10": {"method": "euler", "steps": 10},
    "euler_50": {"method": "euler", "steps": 50},
    "rk4_25": {"method": "midpoint", "steps": 25},
    "dopri5": {"method": "dopri5", "rtol": 1e-5},
}

# FID 结果（示例，ImageNet 256x256）:
# euler_10  → FID=12.3, 时间=50ms
# euler_50  → FID=6.8,  时间=200ms
# rk4_25    → FID=7.1,  时间=120ms
# dopri5    → FID=6.5,  时间=180ms（自适应，步数可变）

# 整流模型的迭代训练过程
import torch
import torch.nn as nn
from copy import deepcopy

class RectifiedFlow:
    def __init__(self, velocity_net, n_rectify_rounds=3):
        self.velocity_net = velocity_net
        self.n_rectify_rounds = n_rectify_rounds
    
    def sample_xt_ut(self, x0, x1, t):
        """线性插值路径"""
        t = t[:, None, None, None]  # 广播到图像维度
        xt = t * x1 + (1 - t) * x0
        ut = x1 - x0  # 常速度
        return xt, ut
    
    def train_round(self, data_loader, optimizer, n_epochs):
        """一轮整流训练"""
        for epoch in range(n_epochs):
            for x1 in data_loader:
                x0 = torch.randn_like(x1)
                t = torch.rand(x0.size(0), device=x0.device)
                
                xt, ut = self.sample_xt_ut(x0, x1, t)
                vt = self.velocity_net(xt, t)
                
                loss = nn.functional.mse_loss(vt, ut)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
    
    @torch.no_grad()
    def generate(self, x0, n_steps=50):
        """ODE 采样"""
        dt = 1.0 / n_steps
        xt = x0.clone()
        for i in range(n_steps):
            t = torch.full((x0.size(0),), i * dt, device=x0.device)
            vt = self.velocity_net(xt, t)
            xt = xt + dt * vt  # Euler 积分
        return xt
    
    def iterative_rectification(self, data_loader, optimizer_fn, 
                                 n_rounds=3, epochs_per_round=50):
        """
        迭代整流：每一轮用上一轮的生成结果重新训练
        """
        # 缓存数据集用于重生成
        original_data = []
        for batch in data_loader:
            original_data.append(batch)
        
        current_data = [x.clone() for x in original_data]
        
        for round_idx in range(n_rounds):
            print(f"=== 整流轮次 {round_idx + 1}/{n_rounds} ===")
            
            # 用当前数据对训练
            optimizer = optimizer_fn(self.velocity_net.parameters())
            self.train_round(current_data, optimizer, epochs_per_round)
            
            # 如果不是最后一轮，重新生成数据对
            if round_idx < n_rectify_rounds - 1:
                new_data = []
                for x0 in [torch.randn_like(x) for x in original_data]:
                    x1_rectified = self.generate(x0, n_steps=100)
                    new_data.append(x1_rectified)
                current_data = new_data
                print(f"  已重新生成 {len(current_data)} 个样本对")
        
        return self.velocity_net

# Score-based Flow Matching
import torch
import torch.nn as nn

class ScoreBasedFlowMatcher(nn.Module):
    """
    结合 Score Matching 和 Flow Matching
    不需要成对 (x0, x1) 数据
    """
    def __init__(self, velocity_net, score_net):
        super().__init__()
        self.velocity_net = velocity_net
        self.score_net = score_net
    
    def forward(self, xt, t, x_data=None):
        """
        统一损失：速度场匹配 + 分数匹配
        L = ||v_θ - u||² + λ ||s_θ - ∇log p_t||²
        """
        if x_data is not None:
            # 有数据时：用条件流匹配
            x0 = torch.randn_like(x_data)
            t_exp = t[:, None, None, None]
            xt_cond = t_exp * x_data + (1 - t_exp) * x0
            ut = x_data - x0
            vt = self.velocity_net(xt_cond, t)
            flow_loss = nn.functional.mse_loss(vt, ut)
        else:
            flow_loss = torch.tensor(0.0, device=xt.device)
        
        # 分数匹配损失（始终可用）
        score_pred = self.score_net(xt, t)
        # 目标分数: ∇_x log p_t(x) ≈ -ε/σ_t (扩散模型中)
        # 或用 Hutchinson 估计器无偏估计
        score_loss = self.compute_score_loss(xt, t)
        
        return flow_loss + 0.1 * score_loss
    
    def compute_score_loss(self, x, t):
        """用 Hutchinson 估计器计算分数匹配损失"""
        z = torch.randn_like(x)
        x.requires_grad_(True)
        
        # 估计 log p_t(x) 的梯度
        log_p = self.estimate_log_density(x, t)
        grad_log_p = torch.autograd.grad(
            log_p.sum(), x, create_graph=True
        )[0]
        
        # Hutchinson 估计: E[z^T ∇²f z] = Tr(∇²f)
        hutch = (z * grad_log_p).sum()
        return hutch ** 2

# SiT: Scaling-Insensitive Flow Matching
import torch

class SiTFlowMatcher:
    """
    尺度不敏感的 Flow Matching
    通过自适应时间缩放和特征归一化提升训练稳定性
    """
    def __init__(self, data_dim, n_bins=100):
        self.data_dim = data_dim
        self.n_bins = n_bins
        # 学习每个特征通道的尺度
        self.feature_scales = nn.Parameter(torch.ones(data_dim))
        # 学习时间缩放函数
        self.time_warp = nn.Sequential(
            nn.Linear(1, 32),
            nn.SiLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()  # 映射到 (0, 1)
        )
    
    def warp_time(self, t):
        """非线性时间缩放：让模型在关键区域花更多"时间""""
        return self.time_warp(t.unsqueeze(-1)).squeeze(-1)
    
    def normalize_features(self, x):
        """特征归一化：使不同维度具有可比尺度"""
        return x * self.feature_scales
    
    def sample_path(self, x0, x1):
        t = torch.rand(x0.size(0), device=x0.device)
        t_warped = self.warp_time(t)
        
        # 归一化后的线性插值
        x0_norm = self.normalize_features(x0)
        x1_norm = self.normalize_features(x1)
        
        t_exp = t_warped[:, None, None, None]
        xt = t_exp * x1_norm + (1 - t_exp) * x0_norm
        ut = x1_norm - x0_norm
        
        return t_warped, xt, ut
    
    # ImageNet 256x256 结果（论文数据）:
    # 标准 FM:   FID=2.37  (1000 steps)
    # SiT-XL/2: FID=1.77  (1000 steps)
    # SiT-XL/2: FID=2.06  (250 steps)
    # DiT-XL/2: FID=2.27  (250 steps)  # 对比基线

# SD3 / FLUX 风格的 Flow Matching Transformer
import torch
import torch.nn as nn
import math

class FlowMatchingDiT(nn.Module):
    """
    简化版 Flow Matching + Transformer 架构
    模仿 SD3 的 MM-DiT 设计
    """
    def __init__(self, hidden_dim=1024, num_layers=24, num_heads=16,
                 patch_size=2, in_channels=16):
        super().__init__()
        self.patch_embed = nn.Conv2d(in_channels, hidden_dim, 
                                      kernel_size=patch_size, stride=patch_size)
        
        # 时间步编码
        self.t_embed = nn.Sequential(
            nn.Linear(1, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
        )
        
        # 文本条件编码
        self.text_embed = nn.Linear(4096, hidden_dim)  # T5-XXL
        
        # Transformer 层
        self.blocks = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=hidden_dim,
                nhead=num_heads,
                dim_feedforward=hidden_dim * 4,
                batch_first=True
            ) for _ in range(num_layers)
        ])
        
        # 输出头：预测速度场
        self.head = nn.Sequential(
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, patch_size * patch_size * in_channels),
        )
    
    def forward(self, x, t, text_cond):
        """
        x: (B, C, H, W) - latent 空间输入
        t: (B,) - 时间步 [0, 1]
        text_cond: (B, L, 4096) - 文本条件
        返回: (B, C, H, W) - 预测的速度场
        """
        B = x.size(0)
        
        # Patch embedding
        patches = self.patch_embed(x)  # (B, hidden, h, w)
        B, C, h, w = patches.shape
        seq = patches.flatten(2).transpose(1, 2)  # (B, h*w, hidden)
        
        # 时间编码 + 文本编码
        t_emb = self.t_embed(t.unsqueeze(-1))  # (B, hidden)
        text_emb = self.text_embed(text_cond)  # (B, L, hidden)
        
        # 将条件注入序列
        t_tokens = t_emb.unsqueeze(1).expand(-1, seq.size(1), -1)
        cond_seq = seq + t_tokens  # 简单加法条件注入
        # 实际 SD3 使用更复杂的跨注意力机制
        
        # Transformer
        for block in self.blocks:
            cond_seq = block(cond_seq)
        
        # 解码回速度场
        velocity = self.head(cond_seq)  # (B, h*w, patch²*C)
        velocity = velocity.transpose(1, 2).view(B, C, h, w)
        
        return velocity

# FLUX [schnell] 风格：4 步采样
import torch

@torch.no_grad()
def flux_schnell_sample(model, noise_shape, text_cond, n_steps=4):
    """
    FLUX schnell 的 4 步 Flow Matching 采样
    通过深度整流，轨迹接近直线，只需极少步数
    """
    x = torch.randn(noise_shape, device=text_cond.device)
    dt = 1.0 / n_steps
    
    for i in range(n_steps):
        t = torch.full((noise_shape[0],), i * dt, device=x.device)
        v = model(x, t, text_cond)
        x = x + dt * v  # Euler 积分
    
    return x

# 性能对比（A100 GPU, 1024x1024 图像）:
# SDXL + DDIM (50步): ~3.2s, FID=6.2
# SD3 + FM (25步):   ~1.8s, FID=4.1
# FLUX schnell (4步): ~0.5s, FID=4.8
# FLUX dev (25步):    ~2.1s, FID=3.2

# 完整的 Flow Matching 训练循环（生产级）
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import CosineAnnealingLR
from tqdm import tqdm
import copy

def train_flow_matching_production(
    model, dataset, batch_size=64, n_epochs=500,
    lr=1e-4, ema_decay=0.9999, device="cuda"
):
    """生产级 Flow Matching 训练"""
    model = model.to(device)
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=lr, weight_decay=1e-2
    )
    
    # 学习率预热 + 余弦衰减
    warmup_steps = int(0.02 * n_epochs * len(dataset) / batch_size)
    total_steps = n_epochs * len(dataset) // batch_size
    
    def lr_lambda(step):
        if step < warmup_steps:
            return step / warmup_steps
        else:
            return 0.5 * (1 + math.cos(math.pi * (step - warmup_steps) / (total_steps - warmup_steps)))
    
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
    
    # EMA 模型
    ema_model = copy.deepcopy(model)
    ema_model.eval()
    
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    global_step = 0
    
    for epoch in range(n_epochs):
        pbar = tqdm(dataloader, desc=f"Epoch {epoch}")
        for x1 in pbar:
            x1 = x1.to(device)
            x0 = torch.randn_like(x1)
            t = torch.rand(x1.size(0), device=device)
            
            # 非线性时间采样：在两端增加密度
            # Beta 分布可以让 t 在 0 和 1 附近更密集
            t = torch.beta(t, torch.full_like(t, 0.8)).to(device)
            
            # 线性插值路径
            t_exp = t[:, None, None, None]
            xt = t_exp * x1 + (1 - t_exp) * x0
            ut = x1 - x0
            
            # 前向传播
            vt = model(xt, t)
            loss = nn.functional.mse_loss(vt, ut)
            
            # 梯度裁剪
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            optimizer.zero_grad()
            scheduler.step()
            
            # 更新 EMA
            with torch.no_grad():
                for ema_param, param in zip(ema_model.parameters(), model.parameters()):
                    ema_param.mul_(ema_decay).add_(param, alpha=1 - ema_decay)
            
            global_step += 1
            pbar.set_postfix(loss=loss.item(), lr=scheduler.get_last_lr()[0])
    
    return model, ema_model

# ODE 求解器对比实验
import torch
from torchdiffeq import odeint
import time

def benchmark_ode_solvers(model, x0, text_cond):
    """对比不同 ODE 求解器的速度和质量"""
    results = {}
    
    solvers = {
        "euler_10": lambda: odeint(
            lambda t, x: model(x, t.unsqueeze(0), text_cond),
            x0, torch.tensor([0., 1.]),
            method="euler", options={"step_size": 0.1}
        ),
        "euler_25": lambda: odeint(
            lambda t, x: model(x, t.unsqueeze(0), text_cond),
            x0, torch.tensor([0., 1.]),
            method="euler", options={"step_size": 0.04}
        ),
        "dopri5": lambda: odeint(
            lambda t, x: model(x, t.unsqueeze(0), text_cond),
            x0, torch.tensor([0., 1.]),
            method="dopri5", rtol=1e-5, atol=1e-5
        ),
        "heun_10": lambda: odeint(
            lambda t, x: model(x, t.unsqueeze(0), text_cond),
            x0, torch.tensor([0., 1.]),
            method="heun3", options={"step_size": 0.1}
        ),
    }
    
    for name, solver_fn in solvers.items():
        torch.cuda.synchronize()
        start = time.perf_counter()
        with torch.no_grad():
            trajectory = solver_fn()
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - start
        
        n_steps = trajectory.size(0) - 1
        results[name] = {
            "time_ms": elapsed * 1000,
            "n_steps": n_steps,
            "output": trajectory[-1]
        }
        print(f"{name}: {elapsed*1000:.1f}ms, {n_steps} steps")
    
    return results

# 从 Diffusion 迁移到 Flow Matching 的最小改动示例
# 假设你已经有扩散模型代码，以下是需要修改的部分

# ===== 修改 1：训练损失 =====
# Diffusion (旧):
# noise = torch.randn_like(x0)
# xt = sqrt_alpha_bar * x0 + sqrt(1 - alpha_bar) * noise
# noise_pred = model(xt, t)
# loss = MSE(noise_pred, noise)

# Flow Matching (新):
x0 = torch.randn_like(x1)  # x1 是真实数据
t = torch.rand(x1.size(0), device=x1.device)
t_exp = t[:, None, None, None]
xt = t_exp * x1 + (1 - t_exp) * x0  # 线性插值
ut = x1 - x0  # 真实速度场
vt = model(xt, t)  # 预测速度场
loss = nn.functional.mse_loss(vt, ut)  # 速度场 MSE

# ===== 修改 2：采样循环 =====
# Diffusion (旧):
# for t in reversed(range(T)):
#     noise_pred = model(xt, t)
#     xt = denoise_step(xt, noise_pred, t)

# Flow Matching (新):
x = torch.randn(shape, device=device)
dt = 1.0 / n_steps
for i in range(n_steps):
    t = torch.full((shape[0],), i * dt, device=device)
    v = model(x, t)
    x = x + dt * v  # Euler 积分

# 就是这么简单！架构、优化器、数据加载都不需要改。