Package: tomesd
Token Merging for Stable DiffusionUsing nothing but pure python and pytorch, ToMe for SD speeds up diffusion by mergingredundanttokens.This is the official implementation ofToMe for SDfrom our short paper:Token Merging for Fast Stable DiffusionDaniel Bolya,Judy HoffmanGitHub|arXiv|BibTeXToMe for SD is an extension of the originalToMe:Token Merging: Your ViT but FasterDaniel Bolya,Cheng-Yang Fu,Xiaoliang Dai,Peizhao Zhang,Christoph Feichtenhofer,Judy HoffmanICLR '23 Oral (Top 5%)|GitHub|arXiv|Blog|BibTeXNote:this extension of ToMe is not affiliated in any way with Meta.What is ToMe for SD?Token Merging (ToMe) speeds up transformers bymerging redundant tokens, which means the transformer has to doless work. We apply this to the underlying transformer blocks in Stable Diffusion in a clever way that minimizes quality loss while keeping most of the speed-up and memory benefits. ToMe for SDdoesn'trequire training and should work out of the box for any Stable Diffusion model.Note:this is a lossy process, so the imagewillchange, ideally not by much. Here are results withFIDscores vs. time and memory usage (lower is better) when using Stable Diffusion v1.5 to generate 512x512 images of ImageNet-1k classes on a 4090 GPU with 50 PLMS steps using fp16:Methodr%FID ↓Time (s/im) ↓Memory (GB/im) ↓Baseline(Original Model)033.123.093.41w/ToMe for SD1032.862.56 (1.21xfaster)2.99 (1.14xless)2032.862.29 (1.35xfaster)2.17 (1.57xless)3032.802.06 (1.50xfaster)1.71 (1.99xless)4032.871.85 (1.67xfaster)1.26 (2.71xless)5033.021.65 (1.87xfaster)0.89 (3.83xless)6033.371.52 (2.03xfaster)0.60 (5.68xless)Even with more than half of the tokens merged (60%!), ToMe for SD still produces images close to the originals, while being2xfasterand using~5.7xless memory. Moreover, ToMe is not another efficient reimplementation of transformer modules. Instead, it actuallyreducesthe total work necessary to generate an image, so it can functionin conjunctionwith efficient implementations (seeUsage).News[2023.04.02]ToMe for SD is now available via pip astomesd. Thanks @mkshing![2023.03.31]ToMe for SD now supportsDiffusers. Thanks @JunnYu and @ExponentialML![2023.03.30]Initial release.See thechangelogfor more details.Supported EnvironmentsThis repo includes code to patch an existing Stable Diffusion environment. Currently, we support the following implementations:Stable Diffusion v2Stable Diffusion v1Latent DiffusionDiffusersAnd potentially othersNote:This also supports most downstream UIs that use these repositories.InstallationToMe for SD requirespytorch >= 1.12.1(forscatter_reduce), which you can get fromhere. Then after installing your choice of stable diffusion environment (supported environments), use the corresponding python environment to install ToMe for SD:pipinstalltomesdInstalling from sourceIf you'd like to install from source to get the latest updates, clone the repository:gitclonehttps://github.com/dbolya/tomesdcdtomesdThen set up the tomesd package with:pythonsetup.pybuilddevelopThat's it! ToMe for SD is implemented in pure python, no CUDA compilation required. 🙂UsageApply ToMe for SD to any Stable Diffusion model withimporttomesd# Patch a Stable Diffusion model with ToMe for SD using a 50% merging ratio.# Using the default options are recommended for the highest quality, tune ratio to suit your needs.tomesd.apply_patch(model,ratio=0.5)# However, if you want to tinker around with the settings, we expose several options.# See docstring and paper for details. Note: you can patch the same model multiple times.tomesd.apply_patch(model,ratio=0.9,sx=4,sy=4,max_downsample=2)# Extreme merging, expect diminishing returnsSee above for what speeds and memory savings you can expect with different ratios.
If you want to remove the patch later, simply usetomesd.remove_patch(model).ExampleTo apply ToMe to the txt2img script of SDv2 or SDv1 for instance, add the following tothis line(SDv2) orthis line(SDv1):importtomesdtomesd.apply_patch(model,ratio=0.5)That's it! More examples and demos coming soon (hopefully).Note:You may not see the full speed-up for the first image generated (as pytorch sets up the graph). Since ToMe for SD uses random processes, you might need to set the seed every batch if you want consistent results.DiffusersToMe can also be used to patch a 🤗 Diffusers Stable Diffusion pipeline:importtorch,tomesdfromdiffusersimportStableDiffusionPipelinepipe=StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",torch_dtype=torch.float16).to("cuda")# Apply ToMe with a 50% merging ratiotomesd.apply_patch(pipe,ratio=0.5)# Can also use pipe.unet in place of pipe hereimage=pipe("a photo of an astronaut riding a horse on mars").images[0]image.save("astronaut.png")You can remove the patch withtomesd.remove_patch(pipe).ToMe + xformers / flash attn / torch 2.0Since ToMe only affects the forward function of the block, it should support most efficient transformer implementations out of the box. Just apply the patch as normal!Note:when testing with xFormers, I observed the most speed-up with ToMe when usingbigimages (i.e., 2048x2048 in the parrot example above). You can get even more speed-up with more aggressive merging configs, but quality obviously suffers. For the result above, I had each method img2img from the same 512x512 res image (i.e., I only applied ToMe during the second pass of "high res fix") and used the default config with 60% merging. Also, the memory benefits didn't stack with xFormers (efficient attention already takes care of memory concerns).CitationIf you use ToMe for SD or this codebase in your work, please cite:@article{bolya2023tomesd,
  title={Token Merging for Fast Stable Diffusion},
  author={Bolya, Daniel and Hoffman, Judy},
  journal={arXiv},
  year={2023}
}If you use ToMe in general please cite the original work:@inproceedings{bolya2023tome,
  title={Token Merging: Your {ViT} but Faster},
  author={Bolya, Daniel and Fu, Cheng-Yang and Dai, Xiaoliang and Zhang, Peizhao and Feichtenhofer, Christoph and Hoffman, Judy},
  booktitle={International Conference on Learning Representations},
  year={2023}
}