扩散模型原理与工程实践

q(x_t | x_{t-1}) = N(x_t; sqrt(1-β_t) x_{t-1}, β_t I)

q(x_t | x_0) = N(x_t; sqrt(ᾱ_t) x_0, (1-ᾱ_t) I)
x_t = sqrt(ᾱ_t) x_0 + sqrt(1-ᾱ_t) ε,  ε ~ N(0, I)

import numpy as np

def linear_beta_schedule(T=1000, beta_start=1e-4, beta_end=0.02):
    return np.linspace(beta_start, beta_end, T, dtype=np.float64)

def cosine_alpha_bar(T=1000, s=0.008):
    steps = np.arange(T + 1, dtype=np.float64) / T
    f = np.cos((steps + s) / (1 + s) * np.pi / 2) ** 2
    alpha_bar = f / f[0]
    betas = np.clip(1 - alpha_bar[1:] / alpha_bar[:-1], 1e-8, 0.999)
    return betas.astype(np.float64)

def q_sample(x0, t, alpha_bar, noise=None):
    """x0: (B, C, H, W), t: (B,) 整数索引"""
    if noise is None:
        noise = np.random.randn(*x0.shape)
    ab = alpha_bar[t].reshape(-1, 1, 1, 1)
    return np.sqrt(ab) * x0 + np.sqrt(1 - ab) * noise

μ_θ(x_t, t) = (1/sqrt(α_t)) * (x_t - (β_t / sqrt(1-ᾱ_t)) * ε_θ(x_t, t))

L_simple = E_{t, x0, ε} [ || ε - ε_θ(sqrt(ᾱ_t) x0 + sqrt(1-ᾱ_t) ε, t) ||² ]

import torch
import torch.nn.functional as F

def diffusion_loss(model, x0, alpha_bar, device="cuda"):
    B = x0.shape[0]
    t = torch.randint(0, len(alpha_bar), (B,), device=device)
    noise = torch.randn_like(x0)
    ab = torch.as_tensor(alpha_bar, device=device)[t].view(-1, 1, 1, 1)
    x_t = ab.sqrt() * x0 + (1 - ab).sqrt() * noise
    pred = model(x_t, t)
    return F.mse_loss(pred, noise)

@torch.no_grad()
def ddim_sample(model, shape, alpha_bar, steps=50, eta=0.0, device="cuda"):
    x = torch.randn(shape, device=device)
    ts = torch.linspace(len(alpha_bar) - 1, 0, steps).long().to(device)
    for i in range(len(ts) - 1):
        t, t_prev = ts[i], ts[i + 1]
        ab_t = alpha_bar[t]
        ab_prev = alpha_bar[t_prev]
        eps = model(x, t.expand(shape[0]))
        x0_pred = (x - (1 - ab_t).sqrt() * eps) / ab_t.sqrt()
        x0_pred = x0_pred.clamp(-1, 1)
        sigma = eta * ((1 - ab_prev) / (1 - ab_t)).sqrt() * (1 - ab_t / ab_prev).sqrt()
        dir_xt = (1 - ab_prev - sigma ** 2).sqrt() * eps
        x = ab_prev.sqrt() * x0_pred + dir_xt
        if sigma > 0:
            x = x + sigma * torch.randn_like(x)
    return x

ε_guided = ε_uncond + s * (ε_cond - ε_uncond)

@torch.no_grad()
def cfg_step(model, x, t, cond, guidance_scale=7.5):
    # 一次前向同时算条件与无条件，batch 翻倍
    eps_uncond, eps_cond = model(torch.cat([x, x]), torch.cat([t, t]),
                                 torch.cat([torch.zeros_like(cond), cond]))
    return eps_uncond + guidance_scale * (eps_cond - eps_uncond)

def drop_condition(cond, p_uncond=0.1):
    mask = torch.rand(cond.shape[0], device=cond.device) < p_uncond
    cond = cond.clone()
    cond[mask] = 0  # 空条件 embedding
    return cond

ControlNet:  x -> [冻结主干] -> +  [可训练副本] -> [zero-conv] -> 输出

class DiTBlock(torch.nn.Module):
    def __init__(self, dim, n_heads):
        super().__init__()
        self.norm1 = torch.nn.LayerNorm(dim, elementwise_affine=False)
        self.attn = torch.nn.MultiheadAttention(dim, n_heads, batch_first=True)
        self.norm2 = torch.nn.LayerNorm(dim, elementwise_affine=False)
        self.mlp = torch.nn.Sequential(
            torch.nn.Linear(dim, 4 * dim), torch.nn.GELU(), torch.nn.Linear(4 * dim, dim))
        self.adaLN = torch.nn.Sequential(torch.nn.SiLU(), torch.nn.Linear(dim, 6 * dim))

    def forward(self, x, cond):
        shift1, scale1, gate1, shift2, scale2, gate2 = self.adaLN(cond).chunk(6, dim=-1)
        h = self.norm1(x) * (1 + scale1.unsqueeze(1)) + shift1.unsqueeze(1)
        x = x + gate1.unsqueeze(1) * self.attn(h, h, h, need_weights=False)[0]
        h = self.norm2(x) * (1 + scale2.unsqueeze(1)) + shift2.unsqueeze(1)
        return x + gate2.unsqueeze(1) * self.mlp(h)

class LoRALinear(torch.nn.Module):
    def __init__(self, base, rank=16, alpha=16):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False
        self.A = torch.nn.Parameter(torch.zeros(rank, base.in_features))
        self.B = torch.nn.Parameter(torch.zeros(base.out_features, rank))
        self.scale = alpha / rank

    def forward(self, x):
        return self.base(x) + (x @ self.A.T @ self.B.T) * self.scale

import math
import torch

def lr_lambda(step, warmup=1000, total=100_000, min_ratio=0.1):
    if step < warmup:
        return step / max(1, warmup)
    progress = (step - warmup) / max(1, total - warmup)
    return min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * progress))

scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

class EMA:
    def __init__(self, model, decay=0.9999):
        self.decay = decay
        self.shadow = {k: v.detach().clone().float()
                       for k, v in model.state_dict().items() if v.dtype.is_floating_point}

    @torch.no_grad()
    def update(self, model):
        for k, v in model.state_dict().items():
            if k in self.shadow:
                self.shadow[k].mul_(self.decay).add_(v.float(), alpha=1 - self.decay)

    def copy_to(self, model):
        model.load_state_dict({k: v.to(model.state_dict()[k].dtype)
                               for k, v in self.shadow.items()}, strict=False)

ε = (x_t - sqrt(ᾱ_t) x0) / sqrt(1-ᾱ_t)
v = sqrt(ᾱ_t) ε - sqrt(1-ᾱ_t) x0

import torch
from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
    variant="fp16",
)
pipe.to("cuda")
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead")
pipe.enable_vae_slicing()          # 降低 VAE 峰值显存
pipe.enable_attention_slicing()    # 显存换速度
# 生成
img = pipe(prompt="a lighthouse at dusk", num_inference_steps=25,
           guidance_scale=7.0).images[0]