算力平台:
令牌合并
令牌合并 (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 时相对于原始管道的加速百分比。
GPU | Resolution | Batch size | Vanilla | ToMe | ToMe + xFormers |
---|---|---|---|---|---|
A100 | 512 | 10 | 6.88 | 5.26 (+23.55%) | 4.69 (+31.83%) |
768 | 10 | OOM | 14.71 | 11 | |
8 | OOM | 11.56 | 8.84 | ||
4 | OOM | 5.98 | 4.66 | ||
2 | 4.99 | 3.24 (+35.07%) | 2.1 (+37.88%) | ||
1 | 3.29 | 2.24 (+31.91%) | 2.03 (+38.3%) | ||
1024 | 10 | OOM | OOM | OOM | |
8 | OOM | OOM | OOM | ||
4 | OOM | 12.51 | 9.09 | ||
2 | OOM | 6.52 | 4.96 | ||
1 | 6.4 | 3.61 (+43.59%) | 2.81 (+56.09%) | ||
V100 | 512 | 10 | OOM | 10.03 | 9.29 |
8 | OOM | 8.05 | 7.47 | ||
4 | 5.7 | 4.3 (+24.56%) | 3.98 (+30.18%) | ||
2 | 3.14 | 2.43 (+22.61%) | 2.27 (+27.71%) | ||
1 | 1.88 | 1.57 (+16.49%) | 1.57 (+16.49%) | ||
768 | 10 | OOM | OOM | 23.67 | |
8 | OOM | OOM | 18.81 | ||
4 | OOM | 11.81 | 9.7 | ||
2 | OOM | 6.27 | 5.2 | ||
1 | 5.43 | 3.38 (+37.75%) | 2.82 (+48.07%) | ||
1024 | 10 | OOM | OOM | OOM | |
8 | OOM | OOM | OOM | ||
4 | OOM | OOM | 19.35 | ||
2 | OOM | 13 | 10.78 | ||
1 | OOM | 6.66 | 5.54 |
如上表所示,tomesd
在处理更大图像分辨率时,速度提升更加明显。另外,值得注意的是,使用 tomesd
,可以运行 1024x1024 等更高分辨率的管道。你可能可以通过 torch.compile
进一步加速推理。