CLIP 跨模态对齐与扩散模型原理

"""最小可运行示例：用 CLIP 给「一张图 + 多条候选文本」打分。

依赖：pip install torch transformers pillow
首次运行会自动下载 openai/clip-vit-base-patch32（约 600MB）。

整条链路只有五步，和讲义里监理的工作顺序完全一致：
    ① 读图  ② 读候选文本  ③ 双塔各自编码  ④ 算 cosine 相似度  ⑤ softmax 成概率
"""

import os

import torch
from PIL import Image
from transformers import CLIPModel, CLIPProcessor

# ① 模型名可以用环境变量覆盖，方便换成本地已下载好的路径
MODEL_NAME = os.environ.get("CLIP_MODEL", "openai/clip-vit-base-patch32")
IMAGE_PATH = os.environ.get("CLIP_IMAGE", "./demo.jpg")

# processor 负责「把图片缩放归一化 + 把文本切成 token」，model 负责两座塔的前向
processor = CLIPProcessor.from_pretrained(MODEL_NAME)
model = CLIPModel.from_pretrained(MODEL_NAME)
model.eval()  # 推理模式：CLIP 在这里参数冻结，不训练

# ② 候选文本。写成一句完整的话，比只写一个词效果好得多
candidates = [
    "a photo of a cat",
    "a photo of a dog",
    "a photo of a car",
    "a photo of a plate of noodles",
]

# ③ 图和文一起交给 processor，padding=True 把长短不一的句子补齐
image = Image.open(IMAGE_PATH).convert("RGB")
inputs = processor(text=candidates, images=image, return_tensors="pt", padding=True)

with torch.no_grad():
    outputs = model(**inputs)

# ④ logits_per_image 形状是 (图片数, 文本数)，值 = cosine 相似度 × 可学习的温度系数
logits = outputs.logits_per_image  # (1, 4)

# ⑤ 在「文本」这一维做 softmax，得到「这张图最像哪句话」的概率分布
probs = logits.softmax(dim=-1)[0]

print("图片：", IMAGE_PATH)
for text, p, raw in zip(candidates, probs.tolist(), logits[0].tolist()):
    print("  %-40s  概率 %6.2f%%   原始分 %7.3f" % (text, p * 100, raw))

best = int(torch.argmax(probs))
print("\n最匹配：%s（%.2f%%）" % (candidates[best], probs[best] * 100))

# 单独拿两座塔的输出向量看一眼：它们维度相同，因为住在同一个向量空间里
img_vec = model.get_image_features(pixel_values=inputs["pixel_values"])
txt_vec = model.get_text_features(input_ids=inputs["input_ids"],
                                  attention_mask=inputs["attention_mask"])
print("图像向量 shape:", tuple(img_vec.shape), " 文本向量 shape:", tuple(txt_vec.shape))

"""案例一：CLIP zero-shot 图像分类。

不训练任何分类头，只把「分类」改写成「哪句描述跟这张图最像」。
换一批类别标签 = 换一个分类器，这就是 zero-shot transfer 的全部秘密。

用法：
    python3 clip_zero_shot.py ./photos            # 批量跑一个目录
    python3 clip_zero_shot.py ./photos/cat.jpg    # 跑单张

依赖：pip install torch transformers pillow
"""

import os
import sys

import torch
from PIL import Image
from transformers import CLIPModel, CLIPProcessor

MODEL_NAME = os.environ.get("CLIP_MODEL", "openai/clip-vit-base-patch32")

# 类别列表：改这里就等于换一个分类器，不需要任何再训练
LABELS = ["猫", "狗", "汽车", "飞机", "蛋糕", "键盘", "盆栽", "自行车"]

# prompt 模板。CLIP 的训练语料是网页图文，句子形态的描述比裸词命中率高，
# 这个技巧在原论文里叫 prompt engineering / prompt ensembling。
TEMPLATES = [
    "a photo of a {}.",
    "a blurry photo of a {}.",
    "a close-up photo of a {}.",
]

# 中文类别转成英文提示词。base 版 CLIP 的文本塔主要吃英文，
# 直接用中文做提示词会明显掉点——想用中文要换 Chinese-CLIP 系列权重。
ZH2EN = {
    "猫": "cat", "狗": "dog", "汽车": "car", "飞机": "airplane",
    "蛋糕": "cake", "键盘": "keyboard", "盆栽": "potted plant", "自行车": "bicycle",
}


def build_text_bank(model, processor, device):
    """把每个类别的多条模板句编码后取平均，作为这一类的「文本原型」。"""
    protos = []
    for zh in LABELS:
        en = ZH2EN[zh]
        sentences = [t.format(en) for t in TEMPLATES]
        batch = processor(text=sentences, return_tensors="pt", padding=True).to(device)
        with torch.no_grad():
            feats = model.get_text_features(**batch)
        # 先做 L2 归一化再平均，避免长句因为向量模长大而被放大
        feats = feats / feats.norm(dim=-1, keepdim=True)
        proto = feats.mean(dim=0)
        proto = proto / proto.norm()
        protos.append(proto)
    return torch.stack(protos)  # (类别数, 向量维度)


def collect_images(path):
    exts = (".jpg", ".jpeg", ".png", ".webp", ".bmp")
    if os.path.isfile(path):
        return [path]
    files = []
    for name in sorted(os.listdir(path)):
        if name.lower().endswith(exts):
            files.append(os.path.join(path, name))
    return files


def main():
    target = sys.argv[1] if len(sys.argv) > 1 else "./photos"
    files = collect_images(target)
    if not files:
        print("没有找到图片：", target)
        return

    device = "cuda" if torch.cuda.is_available() else "cpu"
    processor = CLIPProcessor.from_pretrained(MODEL_NAME)
    model = CLIPModel.from_pretrained(MODEL_NAME).to(device).eval()

    text_bank = build_text_bank(model, processor, device)  # 只算一次，可复用

    for path in files:
        image = Image.open(path).convert("RGB")
        batch = processor(images=image, return_tensors="pt").to(device)
        with torch.no_grad():
            img_feat = model.get_image_features(**batch)
        img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True)

        # 归一化之后的点积就是 cosine 相似度；乘 100 只是把差距拉开便于 softmax
        sims = (img_feat @ text_bank.T)[0]
        probs = (sims * 100).softmax(dim=-1)

        top = torch.topk(probs, k=3)
        head = "  ".join("%s %.1f%%" % (LABELS[i], p * 100)
                         for p, i in zip(top.values.tolist(), top.indices.tolist()))
        print("%-28s -> %s" % (os.path.basename(path), head))


if __name__ == "__main__":
    main()

"""案例二：CLIP 跨模态检索 —— 一批图 + 一句话，取 top-k。

思路和 zero-shot 分类是同一套，只是把矩阵转了个方向：
    分类 = 一张图 对 N 句话   取最大
    检索 = 一句话 对 N 张图   取前 k

工程上的关键点是「图库向量只算一次」：把 N 张图编码成一个矩阵存下来，
之后每来一句查询只需要编码 1 条文本 + 做一次矩阵乘法，毫秒级返回。

用法：
    python3 clip_retrieval.py ./photos "一只趴在键盘上的猫"
依赖：pip install torch transformers pillow
"""

import json
import os
import sys

import torch
from PIL import Image
from transformers import CLIPModel, CLIPProcessor

MODEL_NAME = os.environ.get("CLIP_MODEL", "openai/clip-vit-base-patch32")
INDEX_FILE = os.environ.get("CLIP_INDEX", "./clip_index.pt")
TOP_K = int(os.environ.get("TOP_K", "5"))


def list_images(folder):
    exts = (".jpg", ".jpeg", ".png", ".webp", ".bmp")
    return [os.path.join(folder, n) for n in sorted(os.listdir(folder))
            if n.lower().endswith(exts)]


def build_index(folder, model, processor, device, batch_size=16):
    """把整个图库编码成 (N, D) 的归一化矩阵，落盘复用。"""
    paths = list_images(folder)
    if not paths:
        raise SystemExit("图库为空：" + folder)

    vectors = []
    for i in range(0, len(paths), batch_size):
        chunk = paths[i:i + batch_size]
        images = [Image.open(p).convert("RGB") for p in chunk]
        batch = processor(images=images, return_tensors="pt").to(device)
        with torch.no_grad():
            feats = model.get_image_features(**batch)
        # 归一化：之后点积 == cosine 相似度，省掉每次算模长
        feats = feats / feats.norm(dim=-1, keepdim=True)
        vectors.append(feats.cpu())
        print("已编码 %d / %d" % (min(i + batch_size, len(paths)), len(paths)))

    matrix = torch.cat(vectors, dim=0)
    torch.save({"paths": paths, "matrix": matrix}, INDEX_FILE)
    print("索引已写入", INDEX_FILE, "shape =", tuple(matrix.shape))
    return paths, matrix


def load_or_build(folder, model, processor, device):
    if os.path.exists(INDEX_FILE):
        data = torch.load(INDEX_FILE, map_location="cpu")
        # 图库增删过就重建，避免路径与向量对不上号
        if data["paths"] == list_images(folder):
            print("复用已有索引", INDEX_FILE)
            return data["paths"], data["matrix"]
        print("图库有变化，重建索引")
    return build_index(folder, model, processor, device)


def search(query, paths, matrix, model, processor, device, k=TOP_K):
    batch = processor(text=[query], return_tensors="pt", padding=True).to(device)
    with torch.no_grad():
        txt = model.get_text_features(**batch)
    txt = (txt / txt.norm(dim=-1, keepdim=True)).cpu()

    sims = (matrix @ txt.T)[:, 0]          # (N,) 每张图与这句话的 cosine
    k = min(k, sims.shape[0])
    top = torch.topk(sims, k=k)
    return [(paths[i], float(s)) for s, i in zip(top.values, top.indices)]


def main():
    folder = sys.argv[1] if len(sys.argv) > 1 else "./photos"
    query = sys.argv[2] if len(sys.argv) > 2 else "a cat sitting on a keyboard"

    device = "cuda" if torch.cuda.is_available() else "cpu"
    processor = CLIPProcessor.from_pretrained(MODEL_NAME)
    model = CLIPModel.from_pretrained(MODEL_NAME).to(device).eval()

    paths, matrix = load_or_build(folder, model, processor, device)
    hits = search(query, paths, matrix, model, processor, device)

    print("\n查询：", query)
    for rank, (path, score) in enumerate(hits, 1):
        print("  #%d  %.4f  %s" % (rank, score, path))

    # cosine 相似度是「相对排序」用的，绝对值没有统一阈值含义：
    # 同一句话在不同图库里的最高分可能差很多，别拿 0.3 之类的数字当硬门槛。
    print("\n结果 JSON：")
    print(json.dumps([{"path": p, "score": round(s, 4)} for p, s in hits],
                     ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()

"""案例三：从零实现一个极简扩散模型（DDPM 的最小骨架）。

规模很小：数据是 2D 平面上的一个双月牙点云，网络只有三层 MLP，CPU 上一两分钟能跑完。
它不生成图片，只生成「二维坐标」——但前向加噪、噪声预测、反向采样三段逻辑
与真正的图像扩散模型逐行同构，把维度从 2 换成 (3, 512, 512)、把 MLP 换成 U-Net 即可。
目的只有一个：让加噪 / 去噪循环真正跑在眼前，而不是停留在公式上。

依赖：pip install torch numpy（画图可选 matplotlib）
用法：python3 tiny_diffusion.py
"""

import math
import os

import numpy as np
import torch
import torch.nn as nn

torch.manual_seed(0)
np.random.seed(0)

T = 200            # 扩散总步数：步数越多每步要预测的噪声越少，学起来越容易
BATCH = 256
STEPS = 3000       # 训练迭代次数
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# ---------------------------------------------------------------- 数据
def make_moons(n):
    """手写一个双月牙点云，省掉 sklearn 依赖。"""
    t = np.random.rand(n // 2) * math.pi
    outer = np.stack([np.cos(t), np.sin(t)], axis=1)
    inner = np.stack([1 - np.cos(t), 0.5 - np.sin(t)], axis=1)
    pts = np.concatenate([outer, inner], axis=0)
    pts += np.random.randn(*pts.shape) * 0.06        # 一点点抖动
    pts = (pts - pts.mean(0)) / pts.std(0)           # 标准化到 0 均值 1 方差
    return torch.tensor(pts, dtype=torch.float32)


DATA = make_moons(8192)


# ------------------------------------------------------- 噪声调度表
# beta 从小到大线性增长：早期每步只掺一点噪声，后期掺得多，
# 这样 x_T 才能稳稳落到标准高斯上。
betas = torch.linspace(1e-4, 0.02, T)
alphas = 1.0 - betas
abar = torch.cumprod(alphas, dim=0)                  # \bar{alpha}_t，累乘
sqrt_abar = torch.sqrt(abar)
sqrt_one_minus_abar = torch.sqrt(1.0 - abar)


def q_sample(x0, t, noise):
    """前向过程一步到位：x_t = sqrt(abar_t) * x0 + sqrt(1-abar_t) * eps

    马尔可夫链本来要一步步加，但连续加高斯噪声仍是高斯，
    所以可以用闭式公式从 x0 直接跳到任意第 t 步 —— 训练因此不必真的循环 T 次。
    """
    a = sqrt_abar[t].unsqueeze(-1)
    b = sqrt_one_minus_abar[t].unsqueeze(-1)
    return a * x0 + b * noise


# ---------------------------------------------------------------- 网络
class TimeEmbedding(nn.Module):
    """把整数时间步 t 编码成向量，网络才知道「现在噪声有多重」。"""

    def __init__(self, dim=64):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        half = self.dim // 2
        freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
        args = t.float().unsqueeze(-1) * freqs.unsqueeze(0)
        return torch.cat([torch.sin(args), torch.cos(args)], dim=-1)


class EpsNet(nn.Module):
    """噪声预测网络：输入 (x_t, t)，输出「这一步里混进去的噪声」。

    注意它的输出维度和输入 x 完全一样 —— 图像扩散里 U-Net 也是这个形状约定。
    """

    def __init__(self, dim=2, hidden=128, tdim=64):
        super().__init__()
        self.temb = TimeEmbedding(tdim)
        self.net = nn.Sequential(
            nn.Linear(dim + tdim, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, dim),
        )

    def forward(self, x, t):
        return self.net(torch.cat([x, self.temb(t)], dim=-1))


# ---------------------------------------------------------------- 训练
def train():
    model = EpsNet().to(DEVICE)
    opt = torch.optim.Adam(model.parameters(), lr=2e-3)
    data = DATA.to(DEVICE)

    for step in range(1, STEPS + 1):
        idx = torch.randint(0, data.shape[0], (BATCH,), device=DEVICE)
        x0 = data[idx]

        # 每个样本随机抽一个时间步，等价于「所有步数一起练」
        t = torch.randint(0, T, (BATCH,), device=DEVICE)
        noise = torch.randn_like(x0)
        xt = q_sample(x0, t.cpu(), noise.cpu()).to(DEVICE)

        pred = model(xt, t)
        # 训练目标就是一个普通的回归 loss：预测的噪声要贴近真实撒进去的噪声
        loss = ((pred - noise) ** 2).mean()

        opt.zero_grad()
        loss.backward()
        opt.step()

        if step % 500 == 0:
            print("step %4d   loss %.4f" % (step, loss.item()))
    return model


# ---------------------------------------------------------------- 采样
@torch.no_grad()
def sample(model, n=1000):
    """反向过程：从纯噪声出发，循环 T 次，每次减掉网络预测的那一点噪声。"""
    x = torch.randn(n, 2, device=DEVICE)          # x_T ~ N(0, I)
    traj = []
    for i in reversed(range(T)):
        t = torch.full((n,), i, device=DEVICE, dtype=torch.long)
        eps = model(x, t)

        a_t = alphas[i]
        abar_t = abar[i]
        # DDPM 的后验均值：把预测出来的噪声按系数扣掉
        mean = (x - (1 - a_t) / torch.sqrt(1 - abar_t) * eps) / torch.sqrt(a_t)

        if i > 0:
            # 除最后一步外都要再注入一点随机性，否则样本会塌到同一个点上
            x = mean + torch.sqrt(betas[i]) * torch.randn_like(x)
        else:
            x = mean
        if i % 50 == 0:
            traj.append((i, x.cpu().clone()))
    return x.cpu(), traj


def describe(name, pts):
    print("%-10s 均值 (%6.3f, %6.3f)  标准差 (%5.3f, %5.3f)"
          % (name, pts[:, 0].mean(), pts[:, 1].mean(), pts[:, 0].std(), pts[:, 1].std()))


def main():
    model = train()
    out, traj = sample(model)

    print("\n反向采样过程中点云的统计量变化：")
    for i, pts in traj:
        describe("t=%03d" % i, pts)
    describe("真实数据", DATA)
    describe("生成结果", out)

    np.savetxt("tiny_diffusion_samples.csv", out.numpy(), delimiter=",", fmt="%.5f")
    print("\n已写出 tiny_diffusion_samples.csv，可用任意绘图工具画散点看双月牙形状")

    if os.environ.get("PLOT") == "1":
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(1, 2, figsize=(9, 4))
        ax[0].scatter(DATA[:, 0], DATA[:, 1], s=3)
        ax[0].set_title("real")
        ax[1].scatter(out[:, 0], out[:, 1], s=3, color="crimson")
        ax[1].set_title("sampled")
        plt.tight_layout()
        plt.savefig("tiny_diffusion.png", dpi=120)
        print("已保存 tiny_diffusion.png")


if __name__ == "__main__":
    main()

"""CLIP 打分骨架模板 —— 复制后只改 5 处 TODO。

用途：任何「给一张图和若干候选描述，问哪句最像」的任务都能套。
    图片审核：候选文本写成合规/违规的若干种描述
    素材标注：候选文本写成你的标签体系
    出图质检：候选文本写成 prompt 本身，用得分判断生成结果贴不贴题

依赖：pip install torch transformers pillow
"""

import os

import torch
from PIL import Image
from transformers import CLIPModel, CLIPProcessor

# ---------------------------------------------------------------- TODO 1
# 换模型权重。英文场景用 openai/clip-vit-base-patch32；
# 中文场景换成 OFA-Sys/chinese-clip-vit-base-patch16 一类的中文权重。
MODEL_NAME = os.environ.get("CLIP_MODEL", "openai/clip-vit-base-patch32")

# ---------------------------------------------------------------- TODO 2
# 换候选文本。写成完整句子，不要只写单个名词。
CANDIDATES = [
    "a photo of a cat",
    "a photo of a dog",
]

# ---------------------------------------------------------------- TODO 3
# 换判定阈值。softmax 概率是「候选之间的相对排序」，
# 想做「都不像就拒绝」要额外设一个最低概率或最低 cosine 门槛。
MIN_PROB = 0.5

_device = "cuda" if torch.cuda.is_available() else "cpu"
_processor = CLIPProcessor.from_pretrained(MODEL_NAME)
_model = CLIPModel.from_pretrained(MODEL_NAME).to(_device).eval()


def score(image_path, candidates=None):
    """返回 [(候选文本, 概率), ...]，按概率从高到低排序。"""
    texts = candidates or CANDIDATES
    image = Image.open(image_path).convert("RGB")
    batch = _processor(text=texts, images=image,
                       return_tensors="pt", padding=True).to(_device)
    with torch.no_grad():
        logits = _model(**batch).logits_per_image[0]
    probs = logits.softmax(dim=-1).tolist()
    pairs = list(zip(texts, probs))
    pairs.sort(key=lambda kv: kv[1], reverse=True)
    return pairs


def decide(image_path):
    """在打分之上加一层业务判定。"""
    pairs = score(image_path)
    best_text, best_prob = pairs[0]
    if best_prob < MIN_PROB:
        # ------------------------------------------------------ TODO 4
        # 低置信度时的兜底策略：转人工、落到「其他」类、或者换一套候选重试。
        return {"label": "uncertain", "prob": best_prob, "detail": pairs}
    return {"label": best_text, "prob": best_prob, "detail": pairs}


if __name__ == "__main__":
    # ---------------------------------------------------------- TODO 5
    # 换成你的图片路径或目录遍历。
    result = decide(os.environ.get("CLIP_IMAGE", "./demo.jpg"))
    print(result["label"], "%.2f%%" % (result["prob"] * 100))
    for text, prob in result["detail"]:
        print("   %-40s %6.2f%%" % (text, prob * 100))

"""扩散模型训练循环骨架模板 —— 复制后只改 5 处 TODO。

这份骨架把 DDPM 的训练拆成四个可替换的零件：
    数据集 / 噪声调度表 / 噪声预测网络 / 训练循环
换数据换网络都不影响另外两块，训练循环那段几乎不用动。

依赖：pip install torch
"""

import math

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# ---------------------------------------------------------------- TODO 1
# 扩散步数与训练超参。步数越大单步越好学、采样越慢；
# 玩具数据 200 步够用，图像任务常见 1000 步。
T = 1000
EPOCHS = 20
BATCH = 64
LR = 2e-4


# ---------------------------------------------------------------- TODO 2
def build_dataset():
    """换成你的数据。要求：已归一化到大致 [-1, 1]，形状 (N, ...) 任意。

    图像任务典型做法：
        transforms.Compose([Resize(64), CenterCrop(64), ToTensor(),
                            Normalize([0.5]*3, [0.5]*3)])
    """
    x = torch.randn(4096, 2)          # 占位：换成真实张量
    return TensorDataset(x)


# ------------------------------------------------------- 噪声调度表
def make_schedule(kind="linear"):
    """beta 调度。linear 是原始 DDPM 的做法，cosine 在小分辨率上更稳。"""
    if kind == "cosine":
        s = 0.008
        steps = torch.arange(T + 1, dtype=torch.float32) / T
        f = torch.cos((steps + s) / (1 + s) * math.pi / 2) ** 2
        abar = f / f[0]
        betas = torch.clip(1 - abar[1:] / abar[:-1], 0, 0.999)
    else:
        betas = torch.linspace(1e-4, 0.02, T)
    alphas = 1.0 - betas
    abar = torch.cumprod(alphas, dim=0)
    return betas, alphas, abar


BETAS, ALPHAS, ABAR = make_schedule("linear")
SQRT_ABAR = torch.sqrt(ABAR)
SQRT_1M_ABAR = torch.sqrt(1.0 - ABAR)


def q_sample(x0, t, noise):
    """前向加噪闭式解，一步跳到第 t 步。形状广播按 x0 的维度数自动补齐。"""
    shape = (-1,) + (1,) * (x0.dim() - 1)
    a = SQRT_ABAR.to(x0.device)[t].view(shape)
    b = SQRT_1M_ABAR.to(x0.device)[t].view(shape)
    return a * x0 + b * noise


# ---------------------------------------------------------------- TODO 3
class EpsNet(nn.Module):
    """噪声预测网络。输出形状必须与输入 x 完全一致。

    2D 玩具数据用 MLP 即可；图像任务把这里整个换成 U-Net
    （例如 diffusers 的 UNet2DModel），时间步走它自带的 timestep 入参。
    """

    def __init__(self, dim=2, hidden=128, tdim=64):
        super().__init__()
        self.tdim = tdim
        self.net = nn.Sequential(
            nn.Linear(dim + tdim, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, dim),
        )

    def time_embed(self, t):
        half = self.tdim // 2
        freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
        args = t.float().unsqueeze(-1) * freqs.unsqueeze(0)
        return torch.cat([torch.sin(args), torch.cos(args)], dim=-1)

    def forward(self, x, t):
        return self.net(torch.cat([x, self.time_embed(t)], dim=-1))


# ---------------------------------------------------------------- 训练循环
def train():
    loader = DataLoader(build_dataset(), batch_size=BATCH, shuffle=True, drop_last=True)
    model = EpsNet().to(DEVICE)
    opt = torch.optim.AdamW(model.parameters(), lr=LR)

    for epoch in range(1, EPOCHS + 1):
        total, seen = 0.0, 0
        for (x0,) in loader:
            x0 = x0.to(DEVICE)
            t = torch.randint(0, T, (x0.shape[0],), device=DEVICE)
            noise = torch.randn_like(x0)
            xt = q_sample(x0, t, noise)

            # ------------------------------------------------ TODO 4
            # 想做条件生成（文生图）就在这里把文本 embedding 传进网络：
            #     pred = model(xt, t, text_emb)
            # 并按 10% 概率把 text_emb 置空，训练出 classifier-free guidance 需要的无条件分支。
            pred = model(xt, t)
            loss = ((pred - noise) ** 2).mean()

            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()

            total += loss.item() * x0.shape[0]
            seen += x0.shape[0]
        print("epoch %3d   loss %.4f" % (epoch, total / max(seen, 1)))

    # ------------------------------------------------------------ TODO 5
    # 换成你的保存路径；真实训练还要存 optimizer 状态和 EMA 权重以便断点续训。
    torch.save(model.state_dict(), "eps_net.pt")
    return model


@torch.no_grad()
def sample(model, n=256, shape=(2,), guidance=None):
    """反向采样。guidance 传入 (cond, uncond, scale) 时启用 classifier-free guidance。"""
    x = torch.randn((n,) + shape, device=DEVICE)
    for i in reversed(range(T)):
        t = torch.full((n,), i, device=DEVICE, dtype=torch.long)
        if guidance is None:
            eps = model(x, t)
        else:
            cond, uncond, scale = guidance
            e_c, e_u = model(x, t, cond), model(x, t, uncond)
            # 把「有文本」相对「无文本」多出来的那部分放大 scale 倍
            eps = e_u + scale * (e_c - e_u)

        a_t, abar_t = ALPHAS[i].to(DEVICE), ABAR[i].to(DEVICE)
        mean = (x - (1 - a_t) / torch.sqrt(1 - abar_t) * eps) / torch.sqrt(a_t)
        x = mean if i == 0 else mean + torch.sqrt(BETAS[i].to(DEVICE)) * torch.randn_like(x)
    return x.cpu()


if __name__ == "__main__":
    net = train()
    out = sample(net)
    print("采样结果 shape:", tuple(out.shape))