Stable Diffusion, but the transformer skips duplicate tokens

What it does

ToMe for SD is a drop-in patch that lightens the load on Stable Diffusion’s transformer blocks by merging redundant tokens during inference. Written in pure Python and PyTorch, it needs no fine-tuning or architectural swaps and claims to work across SD v1, v2, Latent Diffusion, and Diffusers pipelines. The trade-off is straightforward: fewer tokens mean less compute and memory, but the output image shifts slightly because the process is inherently lossy.

The interesting bit

This is not another optimized CUDA kernel or attention rewrite. It actually shrinks the total token count the transformer must process, so it stacks with existing speedups like xFormers or Flash Attention rather than replacing them. The authors report that even when merging 60% of tokens, FID scores stay within roughly 0.25 of the baseline.

Key highlights

• At a 60% merge ratio, the authors measured roughly 2× faster inference and nearly 6× lower memory usage on a 4090 generating 512×512 images, with only a marginal FID increase.
• Supports patching SD v1, SD v2, Latent Diffusion, and Hugging Face Diffusers pipelines via a single tomesd.apply_patch call.
• Implemented entirely in Python with no custom CUDA compilation.
• Composable with other efficiency layers such as xFormers, Flash Attention, and PyTorch 2.0 attention.

Caveats

• The process is lossy: images will change, and reproducibility requires fixing seeds because token merging introduces randomness.
• Memory savings largely vanish if you already use xFormers, since efficient attention already optimizes memory; the speedup, however, reportedly remains for large images.
• The first generated image may not reflect full speed gains while PyTorch builds the computation graph.

Verdict

Worth a look if you are running Stable Diffusion locally and want lower latency or VRAM usage without swapping models. Skip it if you need pixel-perfect consistency or are already satisfied with xFormers memory usage on small images.