令牌合并

令牌合并（ToMe）在基于 Transformer 的网络的前向传递中逐步合并冗余令牌/补丁，这可以加速 StableDiffusionPipeline 的推理延迟。

从 pip 安装 ToMe：

pip install tomesd

您可以使用 tomesd 库中的 apply_patch 函数：

  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 进行了推理，设置如下：

我们没有注意到生成样本的质量有任何显著下降，您可以在此 WandB 报告中查看生成的样本。如果您有兴趣重现此实验，请使用此脚本。

基准测试

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

- `diffusers` 版本：0.15.1
- Python 版本：3.8.16
- PyTorch 版本（GPU？）：1.13.1+cu116 (True)
- Huggingface_hub 版本：0.13.2
- Transformers 版本：4.27.2
- Accelerate 版本：0.18.0
- xFormers 版本：0.0.16
- tomesd 版本：0.1.2

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

GPU	分辨率	批处理大小	原始	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

如上表所示，tomesd 带来的加速效果在更大的图像分辨率下变得更加明显。有趣的是，使用 tomesd 可以在更高分辨率如 1024x1024 上运行管道。您可能还可以通过 torch.compile 进一步加速推理。

< > Update on GitHub

Diffusers

令牌合并

基准测试