tomesd

tomesd is the official implementation of Token Merging for Fast Stable Diffusion (ToMe for SD), a technique that accelerates diffusion models by merging redundant tokens inside their transformer blocks so the network does less work per step. Using nothing but pure Python and PyTorch, it delivers meaningful speed-ups and memory savings with minimal impact on image quality.

It works as a drop-in patch: a single apply_patch(model) call wraps the attention blocks of an existing Stable Diffusion model - including diffusers pipelines - and the effect can be tuned or fully removed at any time. Because it is training-free, it works out of the box on any Stable Diffusion checkpoint.

What You Get

• A single apply_patch(model, ratio=...) call that adds token merging to an existing Stable Diffusion model
• Support for both raw model architectures and Hugging Face diffusers pipelines
• A tunable merge ratio to balance speed and memory savings against image fidelity
• A remove_patch call to cleanly restore the original model
• A training-free approach that works on any Stable Diffusion checkpoint out of the box

Common Use Cases

• Reducing per-image generation time in a Stable Diffusion inference service
• Lowering GPU memory usage to generate at higher resolutions on the same hardware
• Speeding up batch or high-resolution image generation with minimal quality loss
• Adding an optional performance mode to a diffusers-based application

Under The Hood

Architecture - The public surface is apply_patch/remove_patch in tomesd/patch.py. apply_patch walks a Stable Diffusion model, and for each transformer block dynamically subclasses it (make_tome_block / make_diffusers_tome_block) so that compute_merge runs bipartite soft matching before attention/MLP and unmerges afterward, storing configuration on a _tome_info dict attached via hook_tome_model. The actual token-matching and merge/unmerge math lives in tomesd/merge.py, keeping the merging algorithm cleanly separated from the model-patching glue.

Tech Stack - Pure Python and PyTorch with no custom CUDA extensions; it targets Stable Diffusion models both in their native form and as Hugging Face diffusers pipelines. Packaged with setuptools.

Code Quality - The code is compact and focused, with clear separation between the merging algorithm (merge.py) and the runtime patching (patch.py), plus an examples/ directory demonstrating usage. It is a reference implementation backed by a CVPR 2023 workshop paper rather than a heavily unit-tested product.

API Design - Extremely low friction: import the package, call apply_patch(model, ratio=0.5), and generation is faster; call remove_patch(model) to revert. The single merge-ratio knob makes the speed/quality trade-off obvious, and support for diffusers means most users need no architectural changes.