Diffusers documentation

Token Merging

Model Database's logo
Join the Model Database community

and get access to the augmented documentation experience

to get started

Token Merging

Token Merging (introduced in Token Merging: Your ViT But Faster) works by merging the redundant tokens / patches progressively in the forward pass of a Transformer-based network. It can speed up the inference latency of the underlying network.

After Token Merging (ToMe) was released, the authors released Token Merging for Fast Stable Diffusion, which introduced a version of ToMe which is more compatible with Stable Diffusion. We can use ToMe to gracefully speed up the inference latency of a DiffusionPipeline. This doc discusses how to apply ToMe to the StableDiffusionPipeline, the expected speedups, and the qualitative aspects of using ToMe on the StableDiffusionPipeline.

Using ToMe

The authors of ToMe released a convenient Python library called tomesd that lets us apply ToMe to a DiffusionPipeline like so:

from diffusers import StableDiffusionPipeline
import tomesd

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

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

And that’s it!

tomesd.apply_patch() exposes a number of arguments to let us strike a balance between the pipeline inference speed and the quality of the generated tokens. Amongst those arguments, the most important one is ratio. ratio controls the number of tokens that will be merged during the forward pass. For more details on tomesd, please refer to the original repository https://github.com/dbolya/tomesd and the paper.

Benchmarking tomesd with StableDiffusionPipeline

We benchmarked the impact of using tomesd on StableDiffusionPipeline along with xformers across different image resolutions. We used A100 and V100 as our test GPU devices with the following development environment (with Python 3.8.5):

- `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

We used this script for benchmarking: https://gist.github.com/sayakpaul/27aec6bca7eb7b0e0aa4112205850335. Following are our findings:

A100

Resolution Batch size Vanilla ToMe ToMe + xFormers ToMe speedup (%) ToMe + xFormers speedup (%)
512 10 6.88 5.26 4.69 23.54651163 31.83139535
768 10 OOM 14.71 11
8 OOM 11.56 8.84
4 OOM 5.98 4.66
2 4.99 3.24 3.1 35.07014028 37.8757515
1 3.29 2.24 2.03 31.91489362 38.29787234
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 2.81 43.59375 56.09375

The timings reported here are in seconds. Speedups are calculated over the Vanilla timings.

V100

Resolution Batch size Vanilla ToMe ToMe + xFormers ToMe speedup (%) ToMe + xFormers speedup (%)
512 10 OOM 10.03 9.29
8 OOM 8.05 7.47
4 5.7 4.3 3.98 24.56140351 30.1754386
2 3.14 2.43 2.27 22.61146497 27.70700637
1 1.88 1.57 1.57 16.4893617 16.4893617
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 2.82 37.75322284 48.06629834
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

As seen in the tables above, the speedup with tomesd becomes more pronounced for larger image resolutions. It is also interesting to note that with tomesd, it becomes possible to run the pipeline on a higher resolution, like 1024x1024.

It might be possible to speed up inference even further with torch.compile().

Quality

As reported in the paper, ToMe can preserve the quality of the generated images to a great extent while speeding up inference. By increasing the ratio, it is possible to further speed up inference, but that might come at the cost of a deterioration in the image quality.

To test the quality of the generated samples using our setup, we sampled a few prompts from the “Parti Prompts” (introduced in Parti) and performed inference with the StableDiffusionPipeline in the following settings:

We didn’t notice any significant decrease in the quality of the generated samples. Here are samples:

tome-samples

You can check out the generated samples here. We used this script for conducting this experiment.