Token 合併

Token 合併 (ToMe) 在基於 Transformer 的網路的正向傳播中逐步合併冗餘的 token/補丁，這可以加速 StableDiffusionPipeline 的推理延遲。

從 pip 安裝 ToMe

pip install tomesd

你可以使用 tomesd 庫的 ToMe，透過 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 函式公開了一些引數，以幫助在管道推理速度和生成 token 的質量之間取得平衡。最重要的引數是 ratio，它控制在正向傳播過程中合併的 token 數量。

如論文所述，ToMe 可以極大地保留生成影像的質量，同時提高推理速度。透過增加 ratio，你可以進一步加速推理，但這會犧牲一些影像質量。

為了測試生成影像的質量，我們從 Parti 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 時相對於原始管道的加速百分比。

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)	記憶體不足 (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 進一步加速推理。

< > 在 GitHub 上更新

Diffusers

Token 合併

基準測試