Skip to content

令牌合并

令牌合并 (ToMe) 在基于 Transformer 的网络的前向传递中逐步合并冗余令牌/补丁,这可以加快 [StableDiffusionPipeline] 的推理延迟。

pip 安装 ToMe:

bash
pip install tomesd

你可以使用 tomesd 库中的 apply_patch 函数使用 ToMe:

diff
  from diffusers import StableDiffusionPipeline
  import torch
  import tomesd

  pipeline = StableDiffusionPipeline.from_pretrained(
        "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True,
  ).to("cuda")
+ tomesd.apply_patch(pipeline, ratio=0.5)

  image = pipeline("a photo of an astronaut riding a horse on mars").images[0]

apply_patch 函数提供了一些 参数,帮助你在管道推理速度和生成令牌的质量之间取得平衡。最重要的参数是 ratio,它控制着在正向传递过程中合并的令牌数量。

正如 论文 中所述,ToMe 可以极大地保留生成图像的质量,同时提高推理速度。通过增加 ratio,你可以进一步加快推理速度,但代价是图像质量会下降。

为了测试生成图像的质量,我们从 Parti Prompts 中采样了一些提示,并使用以下设置执行了 [StableDiffusionPipeline] 推理:

基准测试

我们还对 tomesd 在启用 xFormers 的 [StableDiffusionPipeline] 上的影响进行了基准测试,涵盖了多种图像分辨率。结果是在以下开发环境中的 A100 和 V100 GPU 上获得的:

bash
- `diffusers` version: 0.15.1
- Python version: 3.8.16
- PyTorch version (GPU?): 1.13.1+cu116 (True)
- Huggingface_hub version: 0.13.2
- Transformers version: 4.27.2
- Accelerate version: 0.18.0
- xFormers version: 0.0.16
- tomesd version: 0.1.2

要重现此基准测试,请随时使用此 脚本。结果以秒为单位报告,在适用情况下,我们报告了使用 ToMe 和 ToMe + xFormers 时相对于原始管道的加速百分比。

GPUResolutionBatch sizeVanillaToMeToMe + xFormers
A100512106.885.26 (+23.55%)4.69 (+31.83%)
76810OOM14.7111
8OOM11.568.84
4OOM5.984.66
24.993.24 (+35.07%)2.1 (+37.88%)
13.292.24 (+31.91%)2.03 (+38.3%)
102410OOMOOMOOM
8OOMOOMOOM
4OOM12.519.09
2OOM6.524.96
16.43.61 (+43.59%)2.81 (+56.09%)
V10051210OOM10.039.29
8OOM8.057.47
45.74.3 (+24.56%)3.98 (+30.18%)
23.142.43 (+22.61%)2.27 (+27.71%)
11.881.57 (+16.49%)1.57 (+16.49%)
76810OOMOOM23.67
8OOMOOM18.81
4OOM11.819.7
2OOM6.275.2
15.433.38 (+37.75%)2.82 (+48.07%)
102410OOMOOMOOM
8OOMOOMOOM
4OOMOOM19.35
2OOM1310.78
1OOM6.665.54

如上表所示,tomesd 在处理更大图像分辨率时,速度提升更加明显。另外,值得注意的是,使用 tomesd,可以运行 1024x1024 等更高分辨率的管道。你可能可以通过 torch.compile 进一步加速推理。