Scalable Diffusion Models with Transformers — Meta (Facebook Research)

一句话概括：用 Transformer 替代 U-Net 做扩散模型的骨干网络，在 ImageNet 上生成高质量的类别条件图像。

• 这个项目能做什么？
• 整体架构一览
• 项目文件结构说明
• VAE 在哪里？
• 环境安装
• 数据集要求
• 如何训练
• 如何推理（生成图像）
• 如何评估（FID 等指标）
• 可用的模型变体
• 关键参数速查表

典型用途：你想生成某个类别的图像（如金毛犬、火山、披萨等），只需指定 ImageNet 类别 ID 即可。

┌─────────────────────────────────────────────────────────┐
│                    完整生成流程                           │
│                                                         │
│  训练时：                                                │
│  原始图像 ──→ [VAE Encoder] ──→ 潜空间特征(latent)       │
│                                      │                  │
│                                      ▼                  │
│                              [DiT Transformer]          │
│                              (学习去噪过程)              │
│                                                         │
│  推理时：                                                │
│  随机噪声 ──→ [DiT Transformer] ──→ 去噪后的latent       │
│              (逐步去噪 250步)              │              │
│                                           ▼             │
│                                   [VAE Decoder]         │
│                                           │             │
│                                           ▼             │
│                                      生成的图像          │
└─────────────────────────────────────────────────────────┘

关键点：DiT 只负责"在潜空间中去噪"这一步，VAE 负责"图像 ↔ 潜空间"的转换。

DiT/
├── models.py              # ⭐ DiT 模型定义（Transformer 骨干网络）
├── train.py               # ⭐ 训练脚本（PyTorch DDP 分布式训练）
├── sample.py              # ⭐ 单机推理脚本（生成少量图像）
├── sample_ddp.py          # ⭐ 多 GPU 批量采样脚本（用于 FID 评估）
├── download.py            # 预训练模型自动下载工具
├── run_DiT.ipynb          # Colab/Jupyter 演示 notebook
├── environment.yml        # Conda 环境依赖
├── diffusion/             # 扩散过程实现（来自 OpenAI）
│   ├── __init__.py        #   create_diffusion() 入口函数
│   ├── gaussian_diffusion.py  # 高斯扩散核心（加噪、去噪、损失计算）
│   ├── respace.py         #   时间步重采样（加速采样用）
│   ├── timestep_sampler.py#   训练时的时间步采样策略
│   └── diffusion_utils.py #   KL散度、对数似然等数学工具
└── visuals/               # 示例生成图像

VAE 不在本项目代码中，而是通过 diffusers 库直接加载 Stability AI 的预训练 VAE：

from diffusers.models import AutoencoderKL
vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-ema")

具体来说：

• VAE 来源：Stable Diffusion 的 VAE（stabilityai/sd-vae-ft-ema 或 stabilityai/sd-vae-ft-mse）
• VAE 作用：将 256×256 的图像压缩为 32×32×4 的潜空间特征（缩小 8 倍）
• VAE 不参与训练：在训练和推理中 VAE 权重都是冻结的，只有 DiT 部分被训练
• 两个 VAE 变体：
  • ema：用 EMA 微调的 VAE（训练时默认用这个）
  • mse：用 MSE 损失微调的 VAE（推理时默认用这个，FID 更好）

首次运行时，VAE 权重会从 HuggingFace 自动下载（约 335MB）。

如果远程机器无法访问 HuggingFace/CDN，也可以先在联网机器上运行：

python download_vae.py --variant ema
python download_vae.py --variant mse

默认会生成：

• pretrained_models/vae/sd-vae-ft-ema
• pretrained_models/vae/sd-vae-ft-mse

然后把目录拷到远端，在训练或推理时通过 --vae-path /path/to/local/vae 指定本地 VAE 目录即可。

这里使用的是 Hugging Face 的 snapshot_download：下载的是某个 revision 的文件快照，不包含 git clone 的完整历史。

git clone https://github.com/facebookresearch/DiT.git
cd DiT
conda env create -f environment.yml
conda activate DiT

pip install torch torchvision   # 需要 >= 1.13，建议装 CUDA 版本
pip install timm                # Vision Transformer 组件
pip install diffusers           # 用于加载 Stable Diffusion 的 VAE
pip install accelerate          # diffusers 的依赖
pip install huggingface_hub     # 用于下载本地 VAE 快照

默认使用 ImageNet-1K 数据集（ILSVRC 2012），这是一个包含 1000 个类别、约 128 万张图像的分类数据集。

训练脚本使用 PyTorch 的 ImageFolder，所以数据必须按以下目录结构组织：

/path/to/imagenet/train/
├── n01440764/          # 类别文件夹（ImageNet synset ID）
│   ├── n01440764_10026.JPEG
│   ├── n01440764_10027.JPEG
│   └── ...
├── n01443537/
│   ├── ...
│   └── ...
├── n01484850/
│   └── ...
└── ... (共 1000 个类别文件夹)

关键要求：

• 每个子文件夹 = 一个类别
• 文件夹名就是类别标签（ImageFolder 会自动按文件夹名排序分配 0~999 的标签）
• 图像格式：JPEG/PNG 均可
• 图像尺寸：任意（训练时会自动 center crop 到 256×256 或 512×512）

1. 官方渠道：前往 image-net.org 注册账号，申请下载 ILSVRC2012
2. 学术机构：很多大学/实验室有内部镜像
3. HuggingFace：huggingface-cli download imagenet-1k（需要申请访问权限）
4. Kaggle：搜索 "ImageNet Object Localization Challenge"

可以！ 只要满足 ImageFolder 格式即可。但需要注意：

• 修改 --num-classes 参数为你的类别数
• 如果类别数不是 1000，推理时也要相应修改
• 预训练模型是在 ImageNet-1K 上训练的，自定义数据需要从头训练

训练时自动执行以下预处理（在 train.py 中）：

1. Center Crop：将图像裁剪为正方形（256×256 或 512×512）
2. 随机水平翻转：数据增强
3. 归一化：像素值归一化到 [-1, 1]
4. VAE 编码：在训练循环中，图像会被 VAE 编码为 32×32×4 的潜空间特征

# 单节点 N 卡训练 DiT-XL/2（256×256）
torchrun --nnodes=1 --nproc_per_node=N train.py \
    --model DiT-XL/2 \
    --data-path /path/to/imagenet/train \
    --image-size 256

results/
└── 000-DiT-XL-2/
    ├── checkpoints/
    │   ├── 0050000.pt    # 每 50K 步保存一次
    │   ├── 0100000.pt
    │   └── ...
    └── log.txt           # 训练日志

每个 checkpoint 包含：model（模型权重）、ema（EMA 权重）、opt（优化器状态）、args（训练参数）

• 必须使用 torchrun 启动（DDP 分布式训练），即使单卡也需要
• A100 用户：脚本已默认开启 TF32 加速，训练速度大幅提升
• 不支持断点续训（原版代码未实现，可参考 fast-DiT）

# 生成 256×256 图像（自动下载预训练权重）
python sample.py --image-size 256 --seed 1

# 生成 512×512 图像
python sample.py --image-size 512 --seed 1

运行后会在当前目录生成 sample.png，包含 8 张图像（默认类别：金毛犬、狼、雪豹等）。

python sample.py --model DiT-XL/2 --image-size 256 --ckpt /path/to/your/checkpoint.pt

python sample.py --image-size 256 --vae-path /path/to/sd-vae-ft-mse

打开 run_DiT.ipynb，可以在 Colab 或本地 Jupyter 中交互式运行。

在 sample.py 第 38 行修改类别列表：

class_labels = [207, 360, 387, 974, 88, 979, 417, 279]

这些数字是 ImageNet 类别 ID。常用类别示例：

完整 ImageNet 类别列表：查看

首次运行 sample.py 时会自动下载，也可以手动下载：

或使用下载脚本：

python download.py   # 下载所有预训练模型到 pretrained_models/
python download_vae.py --variant both   # 下载 VAE 到 pretrained_models/vae/

使用 sample_ddp.py 批量生成图像，然后用 ADM 评估工具计算指标：

# 多 GPU 采样 50K 张图
torchrun --nnodes=1 --nproc_per_node=N sample_ddp.py \
    --model DiT-XL/2 \
    --num-fid-samples 50000 \
    --image-size 256 \
    --cfg-scale 1.5

# 生成的 .npz 文件可用 ADM 评估套件计算 FID/IS
# https://github.com/openai/guided-diffusion/tree/main/evaluations

模型命名规则：DiT-{规模}/{patch大小}

• Patch 大小越小（如 /2），输入 token 越多，计算量越大，效果越好
• 预训练权重只提供 DiT-XL/2（256×256 和 512×512 两个版本）

# 最小可运行命令（单节点 1 卡）
torchrun --nnodes=1 --nproc_per_node=1 train.py \
    --model DiT-B/4 \
    --data-path /path/to/imagenet/train \
    --image-size 256 \
    --global-batch-size 32 \
    --epochs 100

# 最简单的推理（自动下载模型）
python sample.py --image-size 256 --seed 42

# 快速采样（减少步数，质量略降）
python sample.py --image-size 256 --num-sampling-steps 50 --cfg-scale 4.0

# 高质量采样
python sample.py --image-size 256 --num-sampling-steps 250 --cfg-scale 4.0

Q: 没有 GPU 能跑吗？ A: 推理可以在 CPU 上跑（很慢），训练必须有 GPU。

Q: 能生成任意内容的图像吗？ A: 不能。这是类别条件模型，只能生成 ImageNet 1000 个类别中的图像，不支持文本描述。

Q: 和 Stable Diffusion 什么关系？ A: DiT 使用了 Stable Diffusion 的 VAE 来做图像编解码，但去噪网络用的是 Transformer 而非 U-Net。两者共享"潜空间扩散"的思路。

Q: 显存需要多少？ A: 推理 DiT-XL/2 约需 8GB 显存；训练 DiT-XL/2 建议 8×A100（80GB）。小模型如 DiT-B/4 训练资源需求小很多。

DiT 是一个类别条件生成模型，它的架构中 y（类别标签）是必须的输入：

# models.py - DiT.forward()
c = t + y   # 时间步嵌入 + 类别嵌入，两者相加作为条件

类别标签 y 通过 LabelEmbedder（一个 nn.Embedding 查找表）转换为向量，和时间步嵌入相加后，控制每个 Transformer Block 的 AdaLN 调制。

而 CelebA-HQ 是一个无类别标签的人脸数据集（只有人脸图像，没有分类标签），所以你需要选择一种策略来处理。

这是最简单也最常用的做法。把所有图像当作同一个类别。

/path/to/celeba_hq/
└── train/
    └── 0/                    # 只有一个子文件夹，文件夹名为 "0"
        ├── 00000.jpg
        ├── 00001.jpg
        ├── 00002.jpg
        └── ... (30000 张人脸图像)

把所有图像放进一个名为 0 的子文件夹即可，ImageFolder 会自动给所有图像分配标签 0。

torchrun --nnodes=1 --nproc_per_node=1 train.py \
    --model DiT-XL/2 \
    --data-path /path/to/celeba_hq/train \
    --image-size 256 \
    --num-classes 1

关键：--num-classes 1

sample.py 需要做两处修改： [读取文件中] 好的，现在我来完整说明。

DiT 的 forward 中，类别标签 y 是必须的输入：

# models.py 第 236-237 行
y = self.y_embedder(y, self.training)    # (N, D)
c = t + y                                # 时间步 + 类别 → 条件向量

y 经过 LabelEmbedder（一个 nn.Embedding(num_classes+1, hidden_size) 查找表）变成向量，和时间步嵌入相加后控制整个去噪过程。CelebA-HQ 没有类别标签，所以需要处理。

思路：所有图像都是"人脸"这一个类别，设 num_classes=1，所有标签都是 0。

1. 数据目录结构

/path/to/celeba_hq/train/
└── 0/                          # 唯一的子文件夹
    ├── 00000.jpg
    ├── 00001.jpg
    └── ... (所有 30K 张图)

只需把所有图片放进一个名为 0 的文件夹，ImageFolder 自动分配标签 0。

2. 训练命令

torchrun --nnodes=1 --nproc_per_node=1 train.py \
    --model DiT-XL/2 \
    --data-path /path/to/celeba_hq/train \
    --image-size 256 \
    --num-classes 1

3. 推理命令

python sample.py \
    --model DiT-XL/2 \
    --image-size 256 \
    --num-classes 1 \
    --ckpt /path/to/your/checkpoint.pt

但 sample.py 中有两处硬编码需要修改：

• 第 50 行 class_labels = [207, 360, ...] → 改为 [0, 0, 0, ...]
• 第 59 行 y_null = torch.tensor([1000] * n, ...) → 空类别 ID 应该是 num_classes 的值，即 1

修改如下：

# 第 50 行：所有标签都是 0（唯一类别）
class_labels = [0] * 8   # 生成 8 张人脸

# 第 59 行：null class 的 ID = num_classes = 1
y_null = torch.tensor([args.num_classes] * n, device=device)

思路：CelebA-HQ 继承了 CelebA 的 40 个二值属性（如 Male、Eyeglasses、Smiling 等）。你可以选几个关键属性组合成类别。

举例：选 3 个属性（Male、Smiling、Eyeglasses），组合出 2³ = 8 个类别：

做法：

1. 读取 CelebA 的 list_attr_celeba.txt 属性文件
2. 写一个预处理脚本，根据属性组合把图片分到对应的子文件夹
3. 训练时设 --num-classes 8

# preprocess_celeba.py - 示例预处理脚本
import os, shutil, pandas as pd

attr = pd.read_csv("list_attr_celeba.txt", sep=r"\s+", skiprows=1)
# CelebA 属性值为 -1/1，转为 0/1
attr = (attr + 1) // 2

src_dir = "/path/to/celeba_hq/images"
dst_dir = "/path/to/celeba_hq/train"

for idx, row in attr.iterrows():
    # 3 个属性组合成类别 ID (0~7)
    class_id = row["Male"] * 4 + row["Smiling"] * 2 + row["Eyeglasses"]
    class_dir = os.path.join(dst_dir, str(class_id))
    os.makedirs(class_dir, exist_ok=True)
    shutil.copy(os.path.join(src_dir, idx), class_dir)

训练：--num-classes 8，推理时指定想要的类别 ID 即可控制生成属性。

对于大多数人脸生成任务，方案 A（无条件生成）就够了。CelebA-HQ 只有 3 万张图，数据量不大，类别越少模型越容易学好。