"""
t2i_lab.py —— 文生图的零件：提示编码器、交叉注意力、带交叉注意力的潜空间去噪网络、扩散 Transformer（DiT）、
无分类器引导、Prompt-to-Prompt 的注意力控制（第 14 讲）
=====================================================================

【这个文件要回答的问题】
    第 12 讲的潜空间扩散模型（LDM）把条件做得很"偷懒"：直接告诉去噪网络 4 个属性的编号，各查一个嵌入加起来。
    真实的 LDM / Stable Diffusion 拿到的是一句话（"orange cat"），要做三件事：
        1. 提示编码器 τ_θ(y)：把一句话变成一串向量（每个词一个）；
        2. 交叉注意力：去噪网络里每个（潜空间的）像素当查询，去"看"提示里的每个词——查询来自图像，键和值来自文字；
        3. 无分类器引导：训练时偶尔把提示换成空的，采样时把"有提示"和"没提示"的预测外推，让图更听话。
    再往后两步：
        4. 扩散 Transformer（DiT）：把 U-Net 换成 ViT 式的 Transformer，用 adaLN-Zero 注入条件；
        5. Prompt-to-Prompt：交叉注意力图决定"哪个词画在哪里"，换词时把原来的注意力图抄过来，就能只改那个词对应的部分。

【本文件提供】
    数据与提示  MAX_TXT、caption_ids、partial_caption、mentioned、sample_captions、load_autoencoder
    模块        PromptEncoder、CrossAttention、CrossLDM（卷积残差块 + 交叉注意力）、DiT（四种注入条件的方式）
    训练与采样  train_t2i、ddim_t2i（可选无分类器引导）、guided_eps
    注意力控制  AttnRecorder（记录交叉注意力图）、align_tokens、P2PController、p2p_edit、switch_edit（对照）、
                ReweightController、reweight_sample
    实验配置    MAIN_CONFIG、PARTIAL_CONFIG、DIT_CONFIG、build_model、latent_data、get_model
    其他        load_or_train、evaluate、per_shape、gflops、patchify / unpatchify、pos_2d

【数据】第 12 讲的带描述的形状世界（caption_world.py）。图像先用第 12 讲缓存的 KL 自编码器压成 4×7×7 的潜变量
    （load_autoencoder；没有缓存时会训练一个，约 1 分钟）。描述有 4 种说法，训练时随机用，测试时默认用第 0 种。
"""

import difflib
import math
import os
import random
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__))
for _d in ("lecture12", "lecture07", "lecture06", "lecture05", "lecture02"):
    _p = os.path.join(_HERE, "..", _d)
    if _p not in sys.path:
        sys.path.insert(0, _p)
from caption_world import (CAPTION_VOCAB, CID, HOLDOUT, HORIZ_WORDS, SHAPE_WORDS, SIZE_WORDS,  # noqa: E402
                           TEMPLATES, VERT_WORDS, attribute_accuracy, caption, make_captioned, read_attributes)
from diffusion_lab import Schedule, time_embedding  # noqa: E402
from ldm_lab import KLAE, encode_all, train_ae  # noqa: E402
from ldm_lab import load_or_train as load_or_train_l12  # noqa: E402
from lm_models import Block, FeedForward, MultiHeadAttention  # noqa: E402
from shape_metrics import frechet_distance  # noqa: E402
from vision_models import patchify  # noqa: E402

CKPT_DIR = os.path.join(_HERE, "checkpoints")
MAX_TXT = 10                         # "<s>" + 最长的说法（"the circle is big and at the top left"，9 个词）
ALL_TEMPLATES = (0, 1, 2, 3)
PAD, BOS = CID["<pad>"], CID["<s>"]
# 第 12 讲 06 号脚本的自编码器配置——一字不差，这样能直接读到它缓存的 checkpoints/ldm_ae.pt
AE_CONFIG = {"zc": 4, "c": 32, "steps": 3000, "kl_weight": 1e-4, "seed": 0}


# =====================================================================
# 1. 数据与提示
# =====================================================================

def caption_ids(captions):
    """
    描述（词列表的列表）→ (N, MAX_TXT) 的编号：每行是 [<s>, 词, 词, …, <pad>, <pad>]。
    【为什么每句前面都加 <s>】空提示（无分类器引导要用的"没有条件"）就是只有一个 <s> 的句子——和 slides 里的
        c_∅ = tokenize("") 一样：真实的分词器把空字符串编码成"开始 + 结束"两个特殊记号。这样交叉注意力的每一行
        至少有一个能看的键，softmax 永远有定义。
    【边界】空列表 [] 表示空提示；超过 MAX_TXT − 1 个词会报错。
    """
    ids = torch.full((len(captions), MAX_TXT), PAD, dtype=torch.long)
    for i, c in enumerate(captions):
        if len(c) > MAX_TXT - 1:
            raise ValueError(f"描述太长（{len(c)} 个词）：{c}")
        ids[i, :len(c) + 1] = torch.tensor([BOS] + [CID[w] for w in c])
    return ids


def partial_caption(attr, rng, p_mention=0.5):
    """
    "不完整的描述"：形状一定说，大小、上下、左右各以 p_mention 的概率说出来，比如 "a circle"、"a big circle at the top"。
    【为什么需要】真实的提示词几乎从不交代一张图的全部内容（"orange cat" 没说猫在哪、多大）——没说的部分由起点噪声决定。
        Prompt-to-Prompt 要解决的正是"固定随机种子、只换一个词，没说的那部分（构图）却变了"；如果每句描述都把位置和
        大小写死，构图就没有变化的余地，这个问题也就看不出来（05 号脚本两种情况都做）。
    【返回】词列表
    """
    s, z, v, h = (SHAPE_WORDS[int(attr[0])], SIZE_WORDS[int(attr[1])], VERT_WORDS[int(attr[2])],
                  HORIZ_WORDS[int(attr[3])])
    say_z, say_v, say_h = (rng.random() < p_mention for _ in range(3))
    words = ["a"] + ([z] if say_z else []) + [s]
    if say_v or say_h:
        words += ["at", "the"] + ([v] if say_v else []) + ([h] if say_h else [])
    return words


def mentioned(words):
    """一句描述提到了哪些属性：[形状, 大小, 上下, 左右] 的 bool 列表。"""
    ws = set(words)
    return [bool(ws & set(g)) for g in (SHAPE_WORDS, SIZE_WORDS, VERT_WORDS, HORIZ_WORDS)]


def sample_captions(attrs, rng, templates=ALL_TEMPLATES, p_partial=0.0):
    """
    每行属性写成一句描述：以概率 p_partial 写成不完整的描述（partial_caption），否则从 templates 里随机挑一种完整的说法。
    【参数】rng：random.Random（可复现）
    """
    return [partial_caption(a, rng) if p_partial and rng.random() < p_partial else caption(a, rng.choice(templates))
            for a in attrs]


def load_autoencoder(x_train, log=print):
    """
    读取第 12 讲缓存的 KL 自编码器（lecture12/checkpoints/ldm_ae.pt；没有就按同样的配置训练一个并存到那里）。
    【返回】(ae, scale)：scale 是训练图潜变量的标准差，扩散模型用的是 z / scale（让潜变量的尺度接近 1，
        LDM 论文的 scale factor）；解码时要乘回去：ae.decode(z * scale)。
    """
    def train_fn(ae):
        return {"rec": train_ae(ae, x_train, AE_CONFIG["steps"], kl_weight=AE_CONFIG["kl_weight"],
                                seed=AE_CONFIG["seed"])}

    ae, _ = load_or_train_l12("ldm_ae", lambda: KLAE(AE_CONFIG["zc"], AE_CONFIG["c"]), train_fn, AE_CONFIG, log=log)
    ae.eval()
    for p in ae.parameters():
        p.requires_grad_(False)
    scale = encode_all(ae, x_train).std().item()
    return ae, scale


# =====================================================================
# 2. 提示编码器与交叉注意力
# =====================================================================

class PromptEncoder(nn.Module):
    """
    slides 的提示模型 τ_θ(y)："就是一个 Transformer"，和去噪网络一起训练。
        词嵌入 + 位置嵌入 → n_layers 个非因果 Transformer 块（不看 <pad>）→ LayerNorm
    【返回】(h, valid)：h (N, T, d) 是每个位置的向量（交叉注意力的键和值从这里来）；valid (N, T) bool，<pad> 处为 False
    【边界】n_layers = 0 时就是"词嵌入 + 位置嵌入"，每个位置只知道自己是什么词——01 号脚本用它做对照。
        Stable Diffusion 用的是冻结的 CLIP 文字编码器（因果的），我们为了简单用一个一起训练的小编码器，和 LDM 论文一致。
    """

    def __init__(self, d=64, n_layers=2, n_heads=4):
        super().__init__()
        self.d = d
        self.tok = nn.Embedding(len(CAPTION_VOCAB), d)
        self.pos = nn.Embedding(MAX_TXT, d)
        self.blocks = nn.ModuleList(Block(d, n_heads, 4 * d, norm="pre", causal=False) for _ in range(n_layers))
        self.norm = nn.LayerNorm(d)

    def forward(self, ids):
        valid = ids != PAD
        h = self.tok(ids) + self.pos(torch.arange(ids.size(1)))
        pad_mask = ((~valid).float() * -1e9)[:, None, None, :]           # (N, 1, 1, T)：键是 <pad> 的地方屏蔽
        for blk in self.blocks:
            blk.attn.extra_mask = pad_mask
            h = blk(h)
            blk.attn.extra_mask = None
        return self.norm(h), valid


def pool(h, valid):
    """对真实的词（非 <pad>）取平均，得到整句的一个向量 (N, d)。DiT 的 adaLN 要一个向量而不是一串。"""
    m = valid.float()[..., None]
    return (h * m).sum(1) / m.sum(1)


class CrossAttention(nn.Module):
    """
    交叉注意力（slides）：查询来自一个序列（图像），键和值来自另一个序列（文字）。
        Q = Y W_q ∈ R^{n×d}（n 个图像位置），K = X W_k ∈ R^{m×d}、V = X W_v ∈ R^{m×d}（m 个提示词）
        A = softmax(Q Kᵀ / √d) ∈ R^{n×m}，输出 = A V ∈ R^{n×d}
    第 t 行 A_t 是"第 t 个（潜空间）像素在提示的 m 个词上的概率分布"——Prompt-to-Prompt 编辑的就是这个矩阵。
    【参数】d_q：图像一侧的维度（也是输出维度）；d_kv：文字一侧的维度；n_heads：头数（每个头一张 n×m 的注意力图）
    【输入】x (B, n, d_q)、ctx (B, m, d_kv)、valid (B, m) bool（False 的键被屏蔽）
    【注意力控制】self.ctrl 不为 None 时，算出 A 之后先交给 ctrl(self.layer_id, A) 处理（记录 / 替换 / 重新加权），
        再用处理后的 A 去乘 V。layer_id 由外面的模型编号（第几个交叉注意力层）。
    【边界】每行至少要有一个能看的键（我们的提示总以 <s> 开头），否则 softmax 全是 −∞，会得到 NaN。
    """

    def __init__(self, d_q, d_kv, n_heads=4):
        super().__init__()
        if d_q % n_heads:
            raise ValueError(f"d_q={d_q} 不能被 n_heads={n_heads} 整除")
        self.h, self.dk = n_heads, d_q // n_heads
        self.W_q = nn.Linear(d_q, d_q)
        self.W_k = nn.Linear(d_kv, d_q)
        self.W_v = nn.Linear(d_kv, d_q)
        self.W_o = nn.Linear(d_q, d_q)
        self.ctrl = None
        self.layer_id = 0

    def forward(self, x, ctx, valid):
        B, n, _ = x.shape
        m = ctx.size(1)
        q = self.W_q(x).view(B, n, self.h, self.dk).transpose(1, 2)        # (B, h, n, dk)
        k = self.W_k(ctx).view(B, m, self.h, self.dk).transpose(1, 2)      # (B, h, m, dk)
        v = self.W_v(ctx).view(B, m, self.h, self.dk).transpose(1, 2)
        s = q @ k.transpose(-2, -1) / math.sqrt(self.dk)                   # (B, h, n, m) 分数
        s = s.masked_fill(~valid[:, None, None, :], float("-inf"))
        A = s.softmax(-1)                                                   # 每行（每个像素）在 m 个词上和为 1
        if self.ctrl is not None:
            A = self.ctrl(self.layer_id, A)
        out = (A @ v).transpose(1, 2).reshape(B, n, self.h * self.dk)
        return self.W_o(out)


def cross_attention_modules(model):
    """模型里所有的 CrossAttention，按出现的顺序（layer_id 也按这个顺序编号）。"""
    return [m for m in model.modules() if isinstance(m, CrossAttention)]


# =====================================================================
# 3. 去噪网络 A：卷积残差块 + 交叉注意力（LDM 的缩小版）
# =====================================================================

class CrossLDM(nn.Module):
    """
    LDM 的去噪网络 ε_θ(z_t, t, τ_θ(y))（缩小版）：
        输入卷积 → n 次 [卷积残差块（GroupNorm + 时间嵌入 → SiLU → 3×3 卷积）→ 交叉注意力块（LayerNorm → 交叉注意力 → 残差）]
        → 输出卷积。
    交叉注意力块把 (c, 7, 7) 的特征图摊平成 49 个 c 维的"像素向量"当查询，提示编码器的输出当键和值——
    slides 说的"Transformer layer inside UNet"里的那一层。时间步仍像第 12 讲那样加在每个卷积块上。
    【参数】zc：潜变量通道数；c：通道数；n：块数；t_dim：时间嵌入维度；d_txt：提示编码器维度；txt_layers：提示编码器层数
    【输入】z (B, zc, 7, 7)、t (B,)、ids (B, MAX_TXT)
    """

    def __init__(self, zc=4, c=64, n=4, t_dim=64, d_txt=64, heads=4, txt_layers=2):
        super().__init__()
        self.t_dim = t_dim
        self.text = PromptEncoder(d_txt, txt_layers)
        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.xnorms = nn.ModuleList([nn.LayerNorm(c) for _ in range(n)])
        self.xattn = nn.ModuleList([CrossAttention(c, d_txt, heads) for _ in range(n)])
        for i, xa in enumerate(self.xattn):
            xa.layer_id = i
        self.out = nn.Conv2d(c, zc, 3, padding=1)

    def forward(self, z, t, ids):
        ctx, valid = self.text(ids)
        e = time_embedding(t, self.t_dim)
        h = self.inp(z)
        B, c, H, W = h.shape
        for conv, lin, norm, xn, xa in zip(self.convs, self.embs, self.norms, self.xnorms, self.xattn):
            h = h + conv(F.silu(norm(h) + lin(e)[:, :, None, None]))
            tok = h.flatten(2).transpose(1, 2)                              # (B, 49, c)：49 个像素当查询
            tok = tok + xa(xn(tok), ctx, valid)
            h = tok.transpose(1, 2).reshape(B, c, H, W)
        return self.out(F.silu(h))


# =====================================================================
# 4. 去噪网络 B：扩散 Transformer（DiT）
# =====================================================================

def unpatchify(tokens, P, C, H):
    """patchify 的逆运算：(B, (H/P)², C·P·P) → (B, C, H, H)（块按行优先排列，和第 5 讲的 patchify 对应）。"""
    B, g = len(tokens), H // P
    return tokens.view(B, g, g, C, P, P).permute(0, 3, 1, 4, 2, 5).reshape(B, C, H, H)


def pos_2d(g, d):
    """
    g×g 网格的二维正弦 / 余弦位置编码 (1, g², d)（DiT 用固定的、不学习的位置编码）：前一半维度编码行号，后一半编码列号，
    每一半都是第 7 讲时间嵌入的同一个公式。
    """
    if d % 4:
        raise ValueError("d 必须是 4 的倍数")
    r, c = torch.meshgrid(torch.arange(g), torch.arange(g), indexing="ij")
    return torch.cat([time_embedding(r.flatten(), d // 2, 100.0), time_embedding(c.flatten(), d // 2, 100.0)], 1)[None]


COND_KINDS = ("in_context", "cross", "adaln", "adaln_zero")


class DiTBlock(nn.Module):
    """
    DiT 块（Peebles & Xie, 2023）：自注意力 + 前馈，四种注入条件 c（= 时间嵌入 + 整句提示的向量）的方式（DiT 论文图 3）：
        in_context  条件当作额外的 token 拼在序列后面，普通的 Transformer 块（像 ViT 的 [CLS]）
        cross       普通块里多一层交叉注意力，键和值 = [时间 token, 提示的每个词]
        adaln       自适应 LayerNorm：LayerNorm 不带可学习的缩放 / 平移，改由一个线性层从 c 回归出 γ、β：
                        LN(x)·(1 + γ(c)) + β(c)
        adaln_zero  adaLN 再加上两个"门" α₁、α₂：x + α·子层(…)；回归 γ、β、α 的线性层初始化为全 0，
                    所以一开始 α = 0，每个块都是恒等映射——深的残差网络从"什么都不做"开始学，训练更稳（和 LoRA 的 B = 0、
                    第 10 讲 adapter 的 W_up = 0 同一个道理）。论文的消融里它最好。
    """

    def __init__(self, d, heads, cond):
        super().__init__()
        if cond not in COND_KINDS:
            raise ValueError(f"cond 只能是 {COND_KINDS}")
        self.cond = cond
        ada = cond in ("adaln", "adaln_zero")
        self.ln1 = nn.LayerNorm(d, elementwise_affine=not ada)
        self.attn = MultiHeadAttention(d, heads, causal=False)
        self.ln2 = nn.LayerNorm(d, elementwise_affine=not ada)
        self.ffn = FeedForward(d, 4 * d)
        if cond == "cross":
            self.ln_x = nn.LayerNorm(d)
            self.xattn = CrossAttention(d, d, heads)
        if ada:
            self.n_mod = 6 if cond == "adaln_zero" else 4
            self.ada = nn.Linear(d, self.n_mod * d)
            if cond == "adaln_zero":
                nn.init.zeros_(self.ada.weight)
                nn.init.zeros_(self.ada.bias)

    def forward(self, x, c, ctx=None, valid=None):
        if self.cond in ("adaln", "adaln_zero"):
            mods = self.ada(F.silu(c))[:, None].chunk(self.n_mod, -1)       # 每个都是 (B, 1, d)，对所有 token 一样
            if self.n_mod == 6:
                b1, g1, a1, b2, g2, a2 = mods
            else:
                (b1, g1, b2, g2), a1, a2 = mods, 1.0, 1.0
            x = x + a1 * self.attn(self.ln1(x) * (1 + g1) + b1)
            return x + a2 * self.ffn(self.ln2(x) * (1 + g2) + b2)
        x = x + self.attn(self.ln1(x))
        if self.cond == "cross":
            x = x + self.xattn(self.ln_x(x), ctx, valid)
        return x + self.ffn(self.ln2(x))


class DiT(nn.Module):
    """
    扩散 Transformer：输入（带噪声的潜变量或图像）切成 P×P 的块 → 线性嵌入 + 二维位置编码 → depth 个 DiT 块
    → 最后一层 LayerNorm（adaLN 版本同样由 c 调制）→ 线性层把每个 token 变回 C·P·P 个数 → 拼回图像形状（预测噪声 ε）。
    slides："DiT 的主干基本就是一个 ViT，加了一点改动"；"补丁越小，GFLOPs 越大，但参数量不变"。
    【参数】in_ch、size：输入的通道数和边长（潜空间 4×7×7；像素空间 1×28×28）；patch：块大小 P；d、depth、heads；
        cond：COND_KINDS 之一；d_txt、txt_layers：提示编码器
    【边界】DiT 原文还预测协方差 Σ（输出通道翻倍，用 Improved DDPM 的变分下界训练）；我们只预测 ε、用固定方差，
        和本课其余的扩散模型一致。
    """

    def __init__(self, in_ch=4, size=7, patch=1, d=64, depth=4, heads=4, cond="adaln_zero", d_txt=64, txt_layers=2):
        super().__init__()
        if size % patch:
            raise ValueError(f"边长 {size} 不能被块大小 {patch} 整除")
        self.in_ch, self.size, self.patch, self.d, self.cond = in_ch, size, patch, d, cond
        self.n_tok = (size // patch) ** 2
        self.text = PromptEncoder(d_txt, txt_layers)
        self.embed = nn.Linear(in_ch * patch * patch, d)
        self.register_buffer("pos", pos_2d(size // patch, d), persistent=False)
        self.t_mlp = nn.Sequential(nn.Linear(d, d), nn.SiLU(), nn.Linear(d, d))
        self.c_proj = nn.Linear(d_txt, d)           # 整句提示的向量 → d 维，和时间嵌入相加得到 c
        self.ctx_proj = nn.Linear(d_txt, d)         # 每个词的向量 → d 维（cross / in_context 用）
        self.blocks = nn.ModuleList(DiTBlock(d, heads, cond) for _ in range(depth))
        ada = cond in ("adaln", "adaln_zero")
        self.final_ln = nn.LayerNorm(d, elementwise_affine=not ada)
        self.final_ada = nn.Linear(d, 2 * d) if ada else None
        self.final = nn.Linear(d, in_ch * patch * patch)
        if cond == "adaln_zero":                    # DiT 也把最后的调制层和输出层初始化为 0：一开始预测的 ε 就是 0
            for lin in (self.final_ada, self.final):
                nn.init.zeros_(lin.weight)
                nn.init.zeros_(lin.bias)

    def forward(self, x, t, ids):
        B = len(x)
        tok = self.embed(patchify(x, self.patch)) + self.pos               # (B, n_tok, d)
        txt, valid = self.text(ids)
        c = self.t_mlp(time_embedding(t, self.d)) + self.c_proj(pool(txt, valid))
        if self.cond == "in_context":
            ctx = self.ctx_proj(txt)
            seq = torch.cat([tok, c[:, None], ctx], 1)                     # [图像 token…, 条件 token, 提示的词…]
            ok = torch.cat([torch.ones(B, self.n_tok + 1, dtype=torch.bool), valid], 1)
            mask = ((~ok).float() * -1e9)[:, None, None, :]
            for blk in self.blocks:
                blk.attn.extra_mask = mask
                seq = blk(seq, c)
                blk.attn.extra_mask = None
            tok = seq[:, :self.n_tok]                                       # 只保留图像 token 的输出
        elif self.cond == "cross":
            ctx = torch.cat([c[:, None], self.ctx_proj(txt)], 1)
            ok = torch.cat([torch.ones(B, 1, dtype=torch.bool), valid], 1)
            for blk in self.blocks:
                tok = blk(tok, c, ctx, ok)
        else:
            for blk in self.blocks:
                tok = blk(tok, c)
        h = self.final_ln(tok)
        if self.final_ada is not None:
            b, g = self.final_ada(F.silu(c))[:, None].chunk(2, -1)
            h = h * (1 + g) + b
        return unpatchify(self.final(h), self.patch, self.in_ch, self.size)


# =====================================================================
# 5. 训练与采样
# =====================================================================

def train_t2i(model, x0, attrs, sched, steps, batch=64, lr=1e-3, seed=0, p_drop=0.1, ema=0.995, warmup=200,
              templates=ALL_TEMPLATES, p_partial=0.5, snapshot_steps=(), log=None, log_every=0):
    """
    训练文本条件的扩散模型（预测噪声 ε）：
        每步抽 batch 张图，每张写成一句描述（以概率 p_partial 写成不完整的描述，否则随机挑一种完整的说法）；
        以概率 p_drop 把描述换成空提示（无分类器引导需要模型同时会
        "有条件"和"无条件"两种预测，用同一组参数——slides："shares parameters with ε_θ(z_t, t, c)"）；
        t ~ 均匀(1..T)，x_t = √ᾱ_t x_0 + √(1−ᾱ_t) ε，损失 ‖ε − ε_θ(x_t, t, 描述)‖²。
    【参数】x0：(N, C, H, W) 训练数据（潜变量 z/scale，或像素 2x − 1）；attrs (N, 4)；ema、warmup 同第 12 讲；
        snapshot_steps：在这些步数把当时的 EMA 参数存一份（04 号脚本画"训练计算量 vs FD"的曲线）
    【返回】{"loss": 最后 200 步的平均损失, "ms_per_step": ⏱, "snapshots": {步数: state_dict}}；训练完模型里是 EMA 参数
    """
    torch.manual_seed(seed)
    gen = torch.Generator().manual_seed(seed)
    rng = random.Random(seed)
    params = [p for p in model.parameters() if p.requires_grad]
    names = [n for n, p in model.named_parameters() if p.requires_grad]
    opt = torch.optim.Adam(params, lr=lr)
    shadow = [p.detach().clone() for p in params]
    hist, snaps, t_total = [], {}, 0.0
    model.train()
    for step in range(1, steps + 1):
        t0 = time.perf_counter()
        for gr in opt.param_groups:
            gr["lr"] = lr * min(1.0, step / warmup)
        idx = torch.randint(len(x0), (batch,), generator=gen)
        caps = sample_captions(attrs[idx], rng, templates, p_partial)
        drop = (torch.rand(batch, generator=gen) < p_drop).tolist()
        ids = caption_ids([[] if d else c for c, d in zip(caps, drop)])
        t = torch.randint(1, sched.T + 1, (batch,), generator=gen)
        eps = torch.randn(x0[idx].shape, generator=gen)
        ab = sched.alpha_bar[t].view(-1, *([1] * (x0.dim() - 1)))
        xt = ab.sqrt() * x0[idx] + (1 - ab).sqrt() * eps
        loss = ((model(xt, t, ids) - eps) ** 2).mean()
        opt.zero_grad()
        loss.backward()
        opt.step()
        with torch.no_grad():
            for sp, p in zip(shadow, params):
                sp.mul_(ema).add_(p.detach(), alpha=1 - ema)
        t_total += time.perf_counter() - t0
        hist.append(loss.item())
        if step in snapshot_steps:
            snaps[step] = {n: sp.clone() for n, sp in zip(names, shadow)}
        if log and log_every and step % log_every == 0:
            log(f"    step {step:5d}  最近 200 步的平均损失 {sum(hist[-200:]) / len(hist[-200:]):.4f}")
    with torch.no_grad():
        for sp, p in zip(shadow, params):
            p.copy_(sp)
    model.eval()
    return {"loss": sum(hist[-200:]) / len(hist[-200:]), "ms_per_step": 1000 * t_total / steps, "snapshots": snaps}


def guided_eps(eps_c, eps_null, w):
    """
    无分类器引导，slides 第一种写法：ε̂ = ε_θ(z_t, t, ∅) + w·(ε_θ(z_t, t, c) − ε_θ(z_t, t, ∅))。
        w = 0：只用无条件的预测；w = 1：普通的条件采样；w > 1：沿"条件减去无条件"的方向外推。
    slides 算法 1 的写法 ε̂ = (1 + w')ε_c − w'ε_∅ 是同一个东西，只是 w' = w − 1（w' = 0 是普通采样；slides 的 w' = 7.5
        对应这里的 8.5；而 Stable Diffusion / diffusers 的 guidance_scale 用的是本函数的写法）。02 号脚本数值上核对两种写法相同。
    """
    return eps_null + w * (eps_c - eps_null)


@torch.no_grad()
def ddim_t2i(model, ids, sched, steps, gen, shape, guidance=None, ctrl=None, x_T=None):
    """
    DDIM（η = 0，第 8 讲）按描述采样。
    【参数】ids (N, MAX_TXT)；shape：单个样本的形状，如 (4, 7, 7)；guidance：None 或无分类器引导的 w（guided_eps）；
        ctrl：注意力控制器（每一步开始前把 ctrl.step 设成第几步，0 = 噪声最大的那一步）；x_T：给定起点噪声（None 则用 gen 采）
    【返回】z_0 的估计 (N, *shape)
    【边界】guidance 不为 None 时，每步把"有提示"和"空提示"拼成一个 2N 的批一起算（代价是两倍的计算量）——
        批的排列是 [N 行有提示, N 行空提示]，P2PController 依赖这个顺序。
    """
    n = len(ids)
    ts = torch.linspace(sched.T, 1, steps).round().long().tolist()
    z = torch.randn(n, *shape, generator=gen) if x_T is None else x_T.clone()
    if guidance is not None:
        ids = torch.cat([ids, caption_ids([[]] * n)])
    for i, t in enumerate(ts):
        if ctrl is not None:
            ctrl.step = i
        ab = sched.alpha_bar[t]
        ab_prev = sched.alpha_bar[ts[i + 1]] if i + 1 < len(ts) else torch.tensor(1.0)
        if guidance is None:
            eps = model(z, torch.full((n,), t), ids)
        else:
            e = model(torch.cat([z, z]), torch.full((2 * n,), t), ids)
            eps = guided_eps(e[:n], e[n:], guidance)
        x0 = (z - (1 - ab).sqrt() * eps) / ab.sqrt()
        z = ab_prev.sqrt() * x0 + (1 - ab_prev).sqrt() * eps
    return z


# =====================================================================
# 6. 注意力控制：记录、对齐、替换（Prompt-to-Prompt）
# =====================================================================

class AttnRecorder:
    """
    记录交叉注意力图：对所有层、所有头取平均，按采样步分开累加，得到每张图 (n 个像素, m 个词) 的平均注意力。
    【用法】rec = AttnRecorder(model)；ddim_t2i(..., ctrl=rec)；rec.maps() → (N, n, m)；rec.remove()
    【边界】带无分类器引导采样时批是 2N 行，maps(n) 只返回前 n 行（有提示的那一半）。
    """

    def __init__(self, model):
        self.mods = cross_attention_modules(model)
        for m in self.mods:
            m.ctrl = self
        self.by_step, self.step = {}, 0

    def __call__(self, layer_id, A):
        a = A.mean(1)                                                       # 对头取平均：(B, n, m)
        s, c = self.by_step.get(self.step, (0, 0))
        self.by_step[self.step] = (s + a, c + 1)
        return A

    def maps(self, n=None, steps=None):
        """【参数】steps：只对这些采样步取平均（None = 全部步）"""
        keys = sorted(self.by_step) if steps is None else [k for k in sorted(self.by_step) if k in steps]
        out = sum(self.by_step[k][0] for k in keys) / sum(self.by_step[k][1] for k in keys)
        return out if n is None else out[:n]

    def remove(self):
        for m in self.mods:
            m.ctrl = None


def align_tokens(src, dst):
    """
    两句提示的词对齐（slides："Swap in just the appropriate parts!"）：对新提示的每个位置 j，给出原提示里对应的位置，
    没有对应（新插入的词）则为 −1。两句前面都有 <s>，所以位置 0 对 0。
        相同的片段：一一对应；
        被替换的片段（"orange" → "big tabby"）：新片段里的每个词都对应原片段的第一个词——slides 的做法：把 "orange"
            的注意力抄给 "big" 和 "tabby" 两个词；长度相同时一一对应；
        插入的片段：−1（这些词保留它们自己在新一轮采样里的注意力）。
    【参数】src、dst：词列表（不含 <s>）
    【返回】长度 MAX_TXT 的列表；<pad> 位置对应原提示的同一个位置（两边都被屏蔽，值是 0，抄不抄都一样）
    【Python 小知识】difflib.SequenceMatcher(None, a, b).get_opcodes() 给出把 a 变成 b 的一串操作：
        ('equal' / 'replace' / 'insert' / 'delete', i1, i2, j1, j2)，表示 a[i1:i2] 对应 b[j1:j2]。
    """
    a, b = ["<s>"] + list(src), ["<s>"] + list(dst)
    out = list(range(MAX_TXT))
    for op, i1, i2, j1, j2 in difflib.SequenceMatcher(None, a, b, autojunk=False).get_opcodes():
        for j in range(j1, j2):
            if op == "equal":
                out[j] = i1 + (j - j1)
            elif op == "replace":
                out[j] = i1 + (j - j1) if i2 - i1 == j2 - j1 else i1
            elif op == "insert":
                out[j] = -1
    return out


class P2PController:
    """
    Prompt-to-Prompt 的注意力替换（Hertz et al., 2022；slides 算法 1）：原提示和新提示**在同一个批里**、从**同一个噪声**
    出发一起去噪；前 n_inject 步，新提示那一行的交叉注意力图 A* 换成原提示那一行的 A（按 align_tokens 的对应关系
    逐列抄过来），之后用它自己的 A*。不论用哪张图，乘的都是新提示的 V（"regardless … you still attend to y*"）。
    【参数】n：每组（原提示 / 新提示）的行数；n_inject：前多少步替换；maps：每行一个 align_tokens 的结果；
        groups：1（不带引导）或 2（带无分类器引导时批是 [原, 新, 原·空, 新·空]，两组各自替换）；
        reweight：None 或 (位置列表, 系数 c)：把新提示这些位置的注意力乘以 c（P2P 的"重新加权"编辑，不重新归一化）
    【边界】替换后每一行不一定再和为 1（新插入的词保留了自己的注意力）——P2P 论文同样不重新归一化。
    """

    def __init__(self, model, n, n_inject, maps, groups=1, reweight=None):
        self.mods = cross_attention_modules(model)
        for m in self.mods:
            m.ctrl = self
        self.n, self.n_inject, self.groups, self.reweight = n, n_inject, groups, reweight
        self.src_idx = torch.tensor([[max(j, 0) for j in mp] for mp in maps])           # (n, MAX_TXT)
        self.use_src = torch.tensor([[j >= 0 for j in mp] for mp in maps])               # (n, MAX_TXT)
        self.step = 0

    def __call__(self, layer_id, A):
        A = A.clone()
        n = self.n
        for g in range(self.groups):
            src, dst = slice(2 * g * n, (2 * g + 1) * n), slice((2 * g + 1) * n, (2 * g + 2) * n)
            if self.step < self.n_inject:
                gathered = torch.gather(A[src], 3, self.src_idx[:, None, None, :].expand_as(A[src]))
                A[dst] = torch.where(self.use_src[:, None, None, :], gathered, A[dst])
            if self.reweight is not None:
                pos, c = self.reweight
                A[dst][..., pos] = A[dst][..., pos] * c
        return A

    def remove(self):
        for m in self.mods:
            m.ctrl = None


@torch.no_grad()
def p2p_edit(model, src_caps, dst_caps, sched, steps, gen, shape, inject_frac, guidance=None, reweight=None):
    """
    用 Prompt-to-Prompt 编辑：同一个噪声 z_T，原提示 y 和新提示 y* 一起去噪，前 inject_frac 比例的步数里把 y* 的
    交叉注意力换成 y 的。inject_frac = 0 就是"固定随机种子、只换提示"这个基线（slides："整张图的结构可能大变"）。
    【返回】(原图的 z_0, 编辑后的 z_0)
    """
    n = len(src_caps)
    ids = caption_ids(list(src_caps) + list(dst_caps))
    x_T = torch.randn(n, *shape, generator=gen)
    ctrl = P2PController(model, n, round(inject_frac * steps), [align_tokens(s, d) for s, d in zip(src_caps, dst_caps)],
                         groups=1 if guidance is None else 2, reweight=reweight)
    try:
        z = ddim_t2i(model, ids, sched, steps, gen, shape, guidance=guidance, ctrl=ctrl, x_T=torch.cat([x_T, x_T]))
    finally:
        ctrl.remove()
    return z[:n], z[n:]


@torch.no_grad()
def switch_edit(model, src_caps, dst_caps, sched, steps, gen, shape, switch_frac, guidance=None):
    """
    另一种只用推理的编辑（对照）：同一个噪声出发，前 switch_frac 比例的步数里两行都用**原提示**去噪（轨迹完全相同），
    之后第二行才换成新提示。它回答 slides 的问题"为什么先用一阵原来的注意力图、再换成新的"：
    如果构图是在前几步定下来的，那么共享前几步的轨迹就能保住构图。
    【返回】(原图的 z_0, 编辑后的 z_0)
    """
    n = len(src_caps)
    ts = torch.linspace(sched.T, 1, steps).round().long().tolist()
    k = round(switch_frac * steps)
    x_T = torch.randn(n, *shape, generator=gen)
    z = torch.cat([x_T, x_T])
    ids_s, ids_d = caption_ids(src_caps), caption_ids(dst_caps)
    for i, t in enumerate(ts):
        ids = torch.cat([ids_s, ids_s if i < k else ids_d])
        if guidance is None:
            eps = model(z, torch.full((2 * n,), t), ids)
        else:
            e = model(torch.cat([z, z]), torch.full((4 * n,), t), torch.cat([ids, caption_ids([[]] * (2 * n))]))
            eps = guided_eps(e[:2 * n], e[2 * n:], guidance)
        ab = sched.alpha_bar[t]
        ab_prev = sched.alpha_bar[ts[i + 1]] if i + 1 < len(ts) else torch.tensor(1.0)
        x0 = (z - (1 - ab).sqrt() * eps) / ab.sqrt()
        z = ab_prev.sqrt() * x0 + (1 - ab_prev).sqrt() * eps
    return z[:n], z[n:]


class ReweightController:
    """
    P2P 的"注意力重新加权"（slides 的 "down-weight existing descriptor"）：把提示里某些位置的交叉注意力乘以 c（不重新归一化），
    c < 1 减弱这个词的作用、c > 1 加强。只改"有提示"的那 n 行，带引导时后面的空提示行不动。
    """

    def __init__(self, model, n, positions, c):
        self.mods = cross_attention_modules(model)
        for m in self.mods:
            m.ctrl = self
        self.n, self.pos, self.c, self.step = n, list(positions), c, 0

    def __call__(self, layer_id, A):
        A = A.clone()
        A[:self.n, ..., self.pos] = A[:self.n, ..., self.pos] * self.c
        return A

    def remove(self):
        for m in self.mods:
            m.ctrl = None


@torch.no_grad()
def reweight_sample(model, caps, sched, steps, gen, shape, positions, c, guidance=None):
    """按提示采样，同时把 positions 这些位置的交叉注意力乘以 c。【返回】z_0"""
    ctrl = ReweightController(model, len(caps), positions, c)
    try:
        return ddim_t2i(model, caption_ids(caps), sched, steps, gen, shape, guidance=guidance, ctrl=ctrl)
    finally:
        ctrl.remove()


# =====================================================================
# 7. 实验配置、数据、缓存、评估、计算量
# =====================================================================

# 主模型（01、02、03、05 号脚本共用）：卷积 + 交叉注意力的潜空间扩散模型，只用完整的描述，10% 的空提示
MAIN_CONFIG = {"kind": "cross_ldm", "c": 64, "n": 4, "d_txt": 64, "txt_layers": 2, "steps": 8000, "batch": 64,
               "lr": 1e-3, "p_drop": 0.1, "p_partial": 0.0, "ema": 0.995, "warmup": 200, "snapshots": [4000]}
# 05 号脚本的对照：一半训练描述是不完整的（"a circle"）
PARTIAL_CONFIG = dict(MAIN_CONFIG, p_partial=0.5, snapshots=[])
# 潜空间 DiT（03、04 号脚本）：块大小 1（每个潜空间像素一个 token，共 49 个）
DIT_CONFIG = {"kind": "dit", "in_ch": 4, "size": 7, "patch": 1, "d": 64, "depth": 4, "heads": 4, "cond": "adaln_zero",
              "d_txt": 64, "txt_layers": 2, "steps": 4000, "batch": 64, "lr": 1e-3, "p_drop": 0.1, "p_partial": 0.0,
              "ema": 0.995, "warmup": 200, "snapshots": [1000, 2000, 3000]}


def build_model(cfg):
    """按配置字典构造模型（CrossLDM 或 DiT）。"""
    if cfg["kind"] == "cross_ldm":
        return CrossLDM(c=cfg["c"], n=cfg["n"], d_txt=cfg["d_txt"], txt_layers=cfg["txt_layers"])
    return DiT(cfg["in_ch"], cfg["size"], cfg["patch"], cfg["d"], cfg["depth"], cfg["heads"], cfg["cond"], cfg["d_txt"],
               cfg["txt_layers"])


def latent_data(log=print):
    """
    训练用的潜变量：第 12 讲的 8192 张训练图（seed = 1）去掉 HOLDOUT 的 4 种属性组合，编码成 z / scale。
    【返回】(ae, scale, z0 (N, 4, 7, 7), attrs (N, 4))
    """
    x_tr, a_tr, _ = make_captioned(8192, 1)
    ae, scale = load_autoencoder(x_tr, log=log)
    keep = torch.tensor([tuple(a.tolist()) not in HOLDOUT for a in a_tr])
    return ae, scale, (encode_all(ae, x_tr) / scale)[keep], a_tr[keep]


def get_model(name, cfg, x0, attrs, scale, quick=False, log=print):
    """
    训练（或读取缓存的）文本条件扩散模型。
    【参数】cfg：MAIN_CONFIG / PARTIAL_CONFIG / DIT_CONFIG 之类；scale：潜变量的尺度（写进缓存配置，自编码器变了就重新训练）；
        quick=True：冒烟测试——只训练 30 步、不读写缓存（数字没有意义）
    【返回】(model, info)：info 里有 loss、ms_per_step（⏱）、snapshots（EMA 参数的快照）、seconds
    """
    steps = 30 if quick else cfg["steps"]
    snaps = [] if quick else cfg.get("snapshots", [])

    def train_fn(m):
        return train_t2i(m, x0, attrs, Schedule(1000), steps, batch=cfg["batch"], lr=cfg["lr"], p_drop=cfg["p_drop"],
                         ema=cfg["ema"], warmup=cfg["warmup"], p_partial=cfg["p_partial"], snapshot_steps=snaps,
                         log=None if quick else log, log_every=max(steps // 4, 1))

    if quick:
        torch.manual_seed(0)
        m = build_model(cfg)
        return m, train_fn(m)
    return load_or_train(name, lambda: build_model(cfg), train_fn,
                         dict(cfg, scale=round(scale, 6), holdout=[list(h) for h in HOLDOUT]), log=log)


def load_or_train(name, build, train_fn, config, log=print):
    """
    有缓存就读取 checkpoints/<name>.pt（配置一致时），否则 build() → train_fn(model) → 保存。
    【返回】(model, info)
    【可复现】构造模型（随机初始化）之前先固定种子：这台机器上 PyTorch 每次启动时的默认种子是随机的（第 13 讲踩过的坑）。
    """
    path = os.path.join(CKPT_DIR, f"{name}.pt")
    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"])
            model.eval()
            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）")
    model.eval()
    return model, info


def evaluate(images, attrs, judge, ref_feats):
    """
    生成图的两把尺子：
        听不听话：用形状拟合读出 4 个属性，和描述比（逐属性准确率 + 4 个全对的比例）；
        质量 + 多样性：和真实参考图比 Fréchet 距离（第 6 讲的 FD，裁判网络的特征，越小越好）。
    【返回】{"acc": [4 个属性的准确率], "all": 4 个全对的比例, "fd": FD, "pred": 读出来的属性 (N, 4)}
    """
    pred = read_attributes(images)
    acc, all4 = attribute_accuracy(pred, attrs)
    return {"acc": acc, "all": all4, "fd": frechet_distance(judge.features(images), ref_feats), "pred": pred}


def per_shape(pred, target):
    """每种形状（圆 / 方 / 三角 / 十字）读对的比例。"""
    return [((pred[:, 0] == target[:, 0]) & (target[:, 0] == k)).sum().item() / max((target[:, 0] == k).sum().item(), 1)
            for k in range(4)]


def gflops(model, *inputs):
    """
    模型做一次前向传播的计算量（GFLOPs，每个样本）：用 PyTorch 的 FlopCounterMode 数矩阵乘法、卷积、注意力里的
    浮点运算（一次乘加算 2 次 FLOP）。和硬件无关——slides："GFLOPS 是一个与硬件无关的计算代价的度量"。
    【边界】不数逐元素的运算（激活函数、LayerNorm、加法），它们在 Transformer 里通常只占很小一部分。
    """
    from torch.utils.flop_counter import FlopCounterMode
    with torch.no_grad(), FlopCounterMode(display=False) as fc:
        model(*inputs)
    return fc.get_total_flops() / 1e9 / len(inputs[0])


__all__ = ["MAX_TXT", "ALL_TEMPLATES", "AE_CONFIG", "CKPT_DIR", "encode_all", "caption_ids", "partial_caption",
           "mentioned", "sample_captions", "load_autoencoder",
           "PromptEncoder", "pool", "CrossAttention", "cross_attention_modules", "CrossLDM", "DiT", "DiTBlock",
           "COND_KINDS", "unpatchify", "pos_2d", "train_t2i", "guided_eps", "ddim_t2i", "AttnRecorder", "align_tokens",
           "P2PController", "p2p_edit", "switch_edit", "ReweightController", "reweight_sample", "load_or_train", "evaluate", "per_shape", "gflops", "Schedule",
           "MAIN_CONFIG", "PARTIAL_CONFIG", "DIT_CONFIG", "build_model", "latent_data", "get_model"]