""" ldm_lab.py —— 潜空间扩散模型（LDM）的零件：自编码器、潜空间里的条件去噪网络、训练与 DDIM 采样（第 12、14 讲） ===================================================================== 【这个文件要回答的问题】 slides（Rombach et al., 2022）：扩散模型通常在像素空间里工作，训练要几百个 GPU 天，推理也慢。 LDM 的想法：先训练一个自编码器，学一个"感知上和像素空间等价"、但维度低得多的潜空间； 把自编码器固定住，在潜空间里训练扩散模型（z_0 = encoder(x)）；生成时在潜空间里去噪，最后用解码器还原成图像。 这个文件提供： KLAE slides 说的第一种自编码器：像 VAE 一样把潜变量往高斯拉（KL 正则，权重很小） 28×28 的图 → 4×7×7 = 196 维的潜变量（像素 784 维） CondDenoiser 潜空间里的去噪网络：几层 3×3 卷积 + 残差 + GroupNorm，时间步和条件（描述的嵌入） 加在每一层上；cond_dim 个输入通道可以是任何"条件向量"（第 14 讲换成交叉注意力） AttrEncoder 最简单的"文字编码器"：描述里 4 个关键词各查一个嵌入、加起来（词袋） train_ae / train_latent_diffusion / ddim_latent / load_or_train """ import os import sys import time import torch import torch.nn as nn import torch.nn.functional as F _HERE = os.path.dirname(os.path.abspath(__file__)) _p = os.path.join(_HERE, "..", "lecture07") if _p not in sys.path: sys.path.insert(0, _p) from diffusion_lab import Schedule, time_embedding # noqa: E402 CKPT_DIR = os.path.join(_HERE, "checkpoints") class KLAE(nn.Module): """ 带 KL 正则的自编码器（LDM 论文的 "KL-reg" 版本）。 编码器：像素重排 (1, 28, 28) → (4, 14, 14) → 3×3 卷积 → ReLU → 步长 2 的卷积 → (c, 7, 7) → ReLU → 1×1 卷积 → (2·zc, 7, 7)，前一半是均值 μ、后一半是对数方差 解码器：结构和第 9 讲 VQ-VAE 的解码器相同，只是输入是连续的 z 【参数】zc：潜变量的通道数（默认 4，潜空间 4×7×7 = 196 维）；c：中间通道数 【为什么要一点点 KL】不加的话，潜变量的尺度和形状可以很"怪"（某些维度特别大、分布很尖），扩散模型不好学； LDM 用很小的 KL 权重（比如 1e-6 量级）只是把它轻轻拉向标准高斯，不影响重建。 """ def __init__(self, zc=4, c=32): super().__init__() self.zc = zc self.enc = nn.Sequential(nn.Conv2d(4, c, 3, padding=1), nn.ReLU(), nn.Conv2d(c, c, 4, 2, 1), nn.ReLU(), nn.Conv2d(c, 2 * zc, 1)) self.dec = nn.Sequential(nn.Conv2d(zc, c, 1), nn.ReLU(), nn.ConvTranspose2d(c, c, 4, 2, 1), nn.ReLU(), nn.Conv2d(c, 4, 3, padding=1)) def encode(self, x): """【返回】(μ, log σ²)，形状都是 (N, zc, 7, 7)""" mu, logvar = self.enc(F.pixel_unshuffle(x, 2)).chunk(2, dim=1) return mu, logvar def decode(self, z): return torch.sigmoid(F.pixel_shuffle(self.dec(z), 2)) def train_ae(ae, x, steps, kl_weight=1e-4, batch=64, lr=1e-3, seed=0): """训练自编码器：重建误差（每张图的平方误差之和）+ kl_weight × KL。【返回】最后 200 步的平均重建误差""" torch.manual_seed(seed) gen = torch.Generator().manual_seed(seed) opt = torch.optim.Adam(ae.parameters(), lr=lr) hist = [] for _ in range(steps): xb = x[torch.randint(len(x), (batch,), generator=gen)] mu, logvar = ae.encode(xb) z = mu + torch.randn(mu.shape, generator=gen) * (0.5 * logvar).exp() rec = ((ae.decode(z) - xb) ** 2).flatten(1).sum(1).mean() kl = 0.5 * (mu ** 2 + logvar.exp() - 1 - logvar).flatten(1).sum(1).mean() loss = rec + kl_weight * kl opt.zero_grad() loss.backward() opt.step() hist.append(rec.item()) ae.eval() return sum(hist[-200:]) / len(hist[-200:]) @torch.no_grad() def encode_all(ae, x, batch=1024): """所有图的潜变量均值 μ（扩散模型的训练数据 z_0 用均值，不再加采样噪声）。""" return torch.cat([ae.encode(x[i:i + batch])[0] for i in range(0, len(x), batch)]) class AttrEncoder(nn.Module): """ 最简单的"文字编码器"：描述 "a big circle at the top left" 里决定内容的 4 个词（形状、大小、上下、左右） 各查一个嵌入、加起来——相当于一个只看关键词、不看词序的"词袋"模型。 【参数】dim：输出维度；null=True 时每个属性多一个"空"的取值（第 14 讲的无分类器引导要用） 【输入】attrs (N, 4) long；取值为 -1 表示"不给这个属性"（用空嵌入） """ def __init__(self, dim=64, null=True): super().__init__() self.null = null self.embs = nn.ModuleList([nn.Embedding(k + (1 if null else 0), dim) for k in (4, 2, 2, 2)]) def forward(self, attrs): out = 0 for i, emb in enumerate(self.embs): a = attrs[:, i] if self.null: a = torch.where(a < 0, torch.full_like(a, emb.num_embeddings - 1), a) out = out + emb(a) return out class CondDenoiser(nn.Module): """ 潜空间里的去噪网络 ε_θ(z_t, t, c)：输入卷积 → n 个残差块（GroupNorm → 加上"时间 + 条件"的嵌入 → SiLU → 3×3 卷积） → 输出卷积。潜空间只有 7×7，所以不需要 U-Net 那样的下采样和上采样。 【参数】zc：潜变量通道数；c：中间通道数；n：残差块个数；t_dim：时间 / 条件嵌入的维度 【输入】z (N, zc, 7, 7)、t (N,)、cond (N, t_dim)（AttrEncoder 的输出；None 表示无条件） """ def __init__(self, zc=4, c=64, n=4, t_dim=64): super().__init__() self.t_dim = t_dim self.inp = nn.Conv2d(zc, c, 3, padding=1) self.convs = nn.ModuleList([nn.Conv2d(c, c, 3, padding=1) for _ in range(n)]) self.embs = nn.ModuleList([nn.Linear(t_dim, c) for _ in range(n)]) self.norms = nn.ModuleList([nn.GroupNorm(8, c) for _ in range(n)]) self.out = nn.Conv2d(c, zc, 3, padding=1) def forward(self, z, t, cond=None): e = time_embedding(t, self.t_dim) if cond is not None: e = e + cond h = self.inp(z) for conv, lin, norm in zip(self.convs, self.embs, self.norms): h = h + conv(F.silu(norm(h) + lin(e)[:, :, None, None])) return self.out(F.silu(h)) def train_latent_diffusion(den, enc, z0, attrs, sched, steps, batch=64, lr=1e-3, seed=0, p_drop=0.0, ema=0.995, warmup=200): """ 在潜空间里训练条件扩散模型（预测噪声 ε，第 7 讲的选项 C）： t ~ 均匀(1..T)，ε ~ N(0, I)，z_t = √ᾱ_t z_0 + √(1−ᾱ_t) ε，损失 ‖ε − ε_θ(z_t, t, enc(属性))‖² 【参数】den：CondDenoiser；enc：AttrEncoder；z0：(N, zc, 7, 7) 已归一化的潜变量；attrs (N, 4) p_drop：以这个概率把条件整个换成"空"（第 14 讲的无分类器引导需要，这里默认 0） ema、warmup：和第 8 讲训练像素空间 U-Net 时一样——学习率在前 warmup 步线性升上去；训练结束后用参数的 指数滑动平均（EMA，θ_ema ← ema·θ_ema + (1 − ema)·θ）代替最后一步的参数（采样质量通常更好） 【返回】最后 200 步的平均损失 """ torch.manual_seed(seed) gen = torch.Generator().manual_seed(seed) params = list(den.parameters()) + list(enc.parameters()) opt = torch.optim.Adam(params, lr=lr) shadow = [p.detach().clone() for p in params] hist = [] for step in range(1, steps + 1): for gr in opt.param_groups: gr["lr"] = lr * min(1.0, step / warmup) idx = torch.randint(len(z0), (batch,), generator=gen) x0, a = z0[idx], attrs[idx].clone() if p_drop: a[torch.rand(batch, generator=gen) < p_drop] = -1 t = torch.randint(1, sched.T + 1, (batch,), generator=gen) eps = torch.randn(x0.shape, generator=gen) ab = sched.alpha_bar[t].view(-1, 1, 1, 1) zt = ab.sqrt() * x0 + (1 - ab).sqrt() * eps loss = ((den(zt, t, enc(a)) - eps) ** 2).mean() opt.zero_grad() loss.backward() opt.step() hist.append(loss.item()) if ema: with torch.no_grad(): for sp, p in zip(shadow, params): sp.mul_(ema).add_(p.detach(), alpha=1 - ema) if ema: with torch.no_grad(): for sp, p in zip(shadow, params): p.copy_(sp) den.eval() enc.eval() return sum(hist[-200:]) / len(hist[-200:]) @torch.no_grad() def ddim_latent(den, enc, attrs, sched, steps, gen, zc=4, guidance=None): """ DDIM（η = 0，第 8 讲）在潜空间里采样：从 z_T ~ N(0, I) 出发，按均匀间隔的 steps 个时间步去噪。 【参数】guidance：None 表示普通的条件采样；一个数 w 表示无分类器引导 ε = ε_空 + w·(ε_条件 − ε_空)（第 14 讲） 【返回】归一化尺度下的 z_0 估计 (N, zc, 7, 7) """ n = len(attrs) ts = torch.linspace(sched.T, 1, steps).round().long().tolist() z = torch.randn(n, zc, 7, 7, generator=gen) cond = enc(attrs) null = enc(torch.full_like(attrs, -1)) if guidance is not None else None for i, t in enumerate(ts): ab = sched.alpha_bar[t] ab_prev = sched.alpha_bar[ts[i + 1]] if i + 1 < len(ts) else torch.tensor(1.0) tt = torch.full((n,), t) eps = den(z, tt, cond) if guidance is not None: eps_null = den(z, tt, null) eps = eps_null + guidance * (eps - eps_null) x0 = (z - (1 - ab).sqrt() * eps) / ab.sqrt() z = ab_prev.sqrt() * x0 + (1 - ab_prev).sqrt() * eps return z def load_or_train(name, build, train_fn, config, log=print): """ 有缓存就读取 checkpoints/.pt（配置一致时），否则 build() → train_fn(model) → 保存。 【返回】(model, info)：info 是 train_fn 返回的字典（含用时） """ path = os.path.join(CKPT_DIR, f"{name}.pt") # 【可复现】PyTorch 每次启动时的默认随机种子是随机的；构造模型（随机初始化）之前一定要先固定种子 torch.manual_seed(0) model = build() if os.path.exists(path): blob = torch.load(path, weights_only=True) if blob["config"] == config: model.load_state_dict(blob["state_dict"]) if log: log(f" （读取缓存：checkpoints/{name}.pt）") return model, blob["info"] t0 = time.perf_counter() info = train_fn(model) info["seconds"] = time.perf_counter() - t0 os.makedirs(CKPT_DIR, exist_ok=True) torch.save({"config": config, "state_dict": model.state_dict(), "info": info}, path) if log: log(f" （已保存到 checkpoints/{name}.pt）") return model, info class LatentModel(nn.Module): """把去噪网络和文字编码器打包成一个模块，方便一起保存。""" def __init__(self, c=64, n=4, t_dim=64): super().__init__() self.den = CondDenoiser(c=c, n=n, t_dim=t_dim) self.enc = AttrEncoder(t_dim) __all__ = ["KLAE", "AttrEncoder", "CondDenoiser", "LatentModel", "Schedule", "train_ae", "encode_all", "train_latent_diffusion", "ddim_latent", "load_or_train"]