Token merging

Token merging (ToMe) 在基于 Transformer 的网络的前向传递中逐步合并冗余 tokens/patches，这可以加速 StableDiffusionPipeline 的推理延迟。

从 pip 安装 ToMe

pip install tomesd

您可以将 ToMe 与 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 函数公开了许多参数，以帮助在 pipeline 推理速度和生成的 tokens 的质量之间取得平衡。最重要的参数是 ratio，它控制在前向传递期间合并的 tokens 数量。

正如 论文 中报告的那样，ToMe 可以在大大保持生成图像质量的同时提高推理速度。通过增加 ratio，您可以进一步加速推理，但会以牺牲一些图像质量为代价。

为了测试生成图像的质量，我们从 Parti Prompts 中采样了一些 prompts，并使用以下设置对 StableDiffusionPipeline 进行了推理

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

基准测试

我们还对 StableDiffusionPipeline 进行了基准测试，评估了在启用 xFormers 的情况下，tomesd 对不同图像分辨率的影响。结果来自以下开发环境中的 A100 和 V100 GPU

- `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 时相对于 vanilla pipeline 的加速百分比。

如上表所示，对于较大的图像分辨率，tomesd 的加速效果更加明显。有趣的是，请注意，使用 tomesd，可以在更高的分辨率（如 1024x1024）下运行 pipeline。您或许可以使用 torch.compile 进一步加速推理。