深色模式
L5 扩散模型推理与加速
目标
- 用 Diffusers 运行一个 DiT 结构的文生图模型，测出每一步去噪的耗时；
- 做一个“动机实验”：测量相邻去噪步之间中间特征的相似度，理解为什么缓存能加速；
- 实现一个最朴素的缓存，量化它带来的加速和质量损失；
- 掌握生成质量的评测方法；
- （选做）对照组内论文 BWCache，画出各模块特征随去噪步变化的热力图，并运行组内论文的官方代码。
做完这个实验，你就能读懂方向三的大部分论文。
环境
- 一张 24 GB 显存的 GPU（RTX 4090、A100 等）
- Python 包：
uv pip install torch diffusers transformers accelerate sentencepiece numpy pillow
本实验使用 PixArt-Σ：约 0.6B 参数的 DiT 文生图模型，显存需求小，也是不少缓存加速论文的评测模型。它的文本编码器是 T5，提示词用英文。
为什么不用 WebUI
Stable Diffusion WebUI、ComfyUI 适合出图，但做加速研究需要改模型内部的计算过程，用 Python 直接调用 Diffusers 才方便插入计时、缓存和统计代码。
第 1 步：生成一张图并计时
python
# step1_timing.py
import time
import torch
from diffusers import PixArtSigmaPipeline
pipe = PixArtSigmaPipeline.from_pretrained(
"PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
).to("cuda")
prompt = "a shiba inu running in the snow, cinematic lighting, highly detailed"
# 用前后钩子给每次 transformer 调用（即每个去噪步）计时
step_times = []
def before(module, args, kwargs):
torch.cuda.synchronize()
module._t0 = time.perf_counter()
def after(module, args, kwargs, output):
torch.cuda.synchronize()
step_times.append(time.perf_counter() - module._t0)
pipe.transformer.register_forward_pre_hook(before, with_kwargs=True)
pipe.transformer.register_forward_hook(after, with_kwargs=True)
pipe(prompt, num_inference_steps=5) # 预热：第一次运行包含初始化开销
step_times.clear()
torch.cuda.synchronize()
start = time.perf_counter()
image = pipe(prompt, num_inference_steps=20, generator=torch.Generator("cuda").manual_seed(0)).images[0]
torch.cuda.synchronize()
total = time.perf_counter() - start
image.save("baseline.png")
print(f"总耗时 {total:.2f} s，其中 transformer {sum(step_times):.2f} s，共 {len(step_times)} 步")
print(f"平均每步 {sum(step_times) / len(step_times) * 1000:.0f} ms")预期结果：transformer 的耗时占总耗时的绝大部分，这就是加速研究的重点。剩下的时间花在文本编码和 VAE 解码上（视频模型的 VAE 解码占比会明显更高）。
第 2 步：动机实验——相邻步的特征有多像
python
# step2_similarity.py（接着上面的 pipe 和 prompt）
import torch.nn.functional as F
features = []
block = pipe.transformer.transformer_blocks[14] # 取中间一层
handle = block.register_forward_hook(
lambda module, args, output: features.append(output.detach().float().flatten().cpu())
)
pipe(prompt, num_inference_steps=20, generator=torch.Generator("cuda").manual_seed(0))
handle.remove()
sims = [F.cosine_similarity(features[i], features[i - 1], dim=0).item() for i in range(1, len(features))]
for i, s in enumerate(sims, start=1):
print(f"第 {i} 步 vs 第 {i - 1} 步：{s:.4f}")把相似度画成折线图。你会看到大部分相邻步的相似度非常高，而开头几步和个别位置变化较大。相似度高说明可以复用上一步的结果，变化大的地方必须重新计算，几乎所有缓存方法都建立在这个观察上。换不同的层、不同的提示词再测一次，看规律是否一致。
第 3 步：最朴素的缓存
每隔一步才真正计算一次 transformer，其余步直接复用上一次的输出：
python
# step3_naive_cache.py（接着上面的 pipe 和 prompt）
class NaiveCache:
def __init__(self, transformer, interval=2):
self.forward = transformer.forward
self.interval = interval
self.reset()
def reset(self):
self.step, self.cached = 0, None
def __call__(self, *args, **kwargs):
if self.cached is None or self.step % self.interval == 0:
self.cached = self.forward(*args, **kwargs)
self.step += 1
return self.cached
cache = NaiveCache(pipe.transformer, interval=2)
pipe.transformer.forward = cache
cache.reset()
image = pipe(prompt, num_inference_steps=20, generator=torch.Generator("cuda").manual_seed(0)).images[0]
image.save("naive_cache.png")
# 做完后恢复原来的 forward（删掉实例上的属性，就会用回类上原来的方法），
# 否则后面的步骤都会在缓存开着的情况下运行
del pipe.transformer.forward用第 1 步的方法计时，再和 baseline.png 对比画质。把 interval 改成 3、4，记录“加速比—质量”的变化。你会发现这种固定间隔的做法质量下降很快。研究中的方法都在回答：什么时候复用、复用哪一部分、复用时怎么修正误差：
| 思路 | 代表工作 |
|---|---|
| 复用整步，按输入变化自适应决定是否跳过 | TeaCache |
| 复用部分层或部分模块 | DeepCache、FORA、BWCache |
| 按 token 区分，只重算重要的 token | ToCa |
| 不直接复用，而是预测下一步的特征 | TaylorSeer、Predict to Skip |
交互演示：特征缓存
下面用示意数据演示缓存方法的基本逻辑：累计变化低于阈值时复用上一次的结果，否则重新计算。拖动阈值，看看加速比和误差怎么变化，再和第 3 步“每隔一步复用”的做法对比。
第 4 步：评测生成质量
加速方法要同时报告速度和质量。
和原始输出比（衡量加速带来的损失，加速研究最常用）：
python
import numpy as np
from PIL import Image
def psnr(path_a, path_b):
a = np.asarray(Image.open(path_a), dtype=np.float64)
b = np.asarray(Image.open(path_b), dtype=np.float64)
mse = np.mean((a - b) ** 2)
return 10 * np.log10(255 ** 2 / mse)
print(f"PSNR: {psnr('baseline.png', 'naive_cache.png'):.2f} dB")SSIM 和 LPIPS 可以用 torchmetrics 计算。PSNR、SSIM 越高越好，LPIPS 越低越好。
常用指标汇总
| 类别 | 指标 | 说明 |
|---|---|---|
| 与原始输出的差异 | PSNR、SSIM、LPIPS | 同一提示词和随机种子下，与不加速版本逐张对比 |
| 图文一致性 | CLIP Score | 图像与提示词的向量相似度 |
| 人类偏好 | PickScore、ImageReward、HPSv2 | 用人类偏好数据训练的打分模型，可从多个候选中选最佳图像 |
| 分布质量 | FID、IS | FID 衡量生成图像与真实图像分布的差异，越低越好；IS 衡量清晰度和多样性，越高越好，现在用得较少 |
| 视频 | VBench | 从画质、运动、时间一致性等多个维度评测视频生成 |
第 5 步（选做）：对照组内论文
画出每个模块的变化
BWCache（ICLR 2026）的出发点是一张热力图：横轴是去噪步，纵轴是 DiT 的各个模块，颜色表示该模块的输出和上一步相比变化了多少（相对 L1 距离）。在 PixArt-Σ 上画一张同样的图：
先确认第 3 步之后已经恢复了原来的 forward（del pipe.transformer.forward）。如果缓存还开着，有一半的去噪步不会真正执行 DiT 模块，画出的热力图是错的。
python
# step5_block_heatmap.py（接着上面的 pipe 和 prompt）
import matplotlib.pyplot as plt
import numpy as np
blocks = pipe.transformer.transformer_blocks
prev, diff = {}, {i: [] for i in range(len(blocks))}
def make_hook(i):
def hook(module, args, output):
out = output.detach().float()
if i in prev: # 相对 L1 距离：|这一步 − 上一步| 的均值 ÷ |上一步| 的均值
diff[i].append(((out - prev[i]).abs().mean() / prev[i].abs().mean()).item())
prev[i] = out # 只保留上一步的输出，避免占用太多内存
return hook
handles = [b.register_forward_hook(make_hook(i)) for i, b in enumerate(blocks)]
pipe(prompt, num_inference_steps=20, generator=torch.Generator("cuda").manual_seed(0))
for h in handles:
h.remove()
plt.imshow(np.array([diff[i] for i in range(len(blocks))]), aspect="auto", cmap="viridis")
plt.xlabel("Timestep")
plt.ylabel("Block")
plt.colorbar(label="Relative L1 distance")
plt.savefig("block_heatmap.png", dpi=150, bbox_inches="tight")对照 BWCache 论文中的图：各模块的变化随去噪步怎样变化？是否也是开头和结尾变化大、中间变化小？不同模块之间有没有差别？换几个提示词再画一次。
运行官方代码
做完前 4 步，就可以运行组内论文的官方代码了：
| 论文 | 代码 | 支持的模型 |
|---|---|---|
| BWCache（ICLR 2026） | hsc113/BWCache | Open-Sora 等视频模型 |
| PrediT（ACM MM 2026） | hsc113/PrediT | FLUX.1-dev、DiT-XL/2、HunyuanVideo |
| SCOPE（ACM MM 2026） | hsc113/SCOPE | MAGI-1、SkyReels-V2 |
按各仓库的 README 安装。PrediT 和 SCOPE 可以先在没有 GPU 的电脑上运行单元测试，确认安装正确后再上服务器跑模型。复现的顺序和注意事项见复现组内论文。
现成的加速工具
- Diffusers 内置了几种缓存方法（如 First Block Cache、Pyramid Attention Broadcast、FasterCache），用法见官方文档的缓存章节。
- cache-dit：统一的 DiT 缓存加速库，支持较多模型。
- SGLang Diffusion 集成了 TeaCache、TaylorSeer 等方法。
- xDiT：多卡并行推理 DiT。
这些工具适合做对比实验的 baseline。
视频生成模型
| 模型 | 说明 |
|---|---|
| Wan2.2 | 阿里通义万相，Apache-2.0 许可。TI2V-5B 版本可在单张 4090 上生成 720p 视频 |
| HunyuanVideo 1.5 | 腾讯混元，约 8.3B 参数，可在消费级显卡上运行；许可证有使用限制 |
| LTX-2 系列 | Lightricks 开源，支持音视频同步生成；注意许可证条款 |
| CogVideoX | 智谱开源，较早但资料多 |
| Open-Sora | 全流程开源（含训练），适合研究训练与推理全链路 |
| Stable Video Diffusion | 2023 年的 UNet 结构图生视频模型，已较旧，作为了解历史 |
视频模型更新很快，选用前先查看各项目的最新版本、显存需求和许可证。
其他生成任务（按需了解）：口型同步 LatentSync、Wav2Lip；图像驱动对口型 Sonic；人脸动画 SkyReels-A1（组内实测单张 4090 显存不够）。
常见问题
- 显存不足：使用
pipe.enable_model_cpu_offload()代替.to("cuda")，或换更小的模型；计时实验要在不 offload 的条件下做，否则数据没有意义。 - 下载模型慢：设置
HF_ENDPOINT，见网络加速。 - 结果每次不一样：固定
generator的随机种子，对比实验必须用相同的种子和提示词。