Token Merging (ToMe) for SD - possible 2x speedup

[MrSaad]

"Token Merging (ToMe) speeds up transformers by merging redundant tokens, which means the transformer has to do less 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 SD doesn't require training and should work out of the box for any Stable Diffusion model."

paper:
https://arxiv.org/abs/2303.17604
github:
https://github.com/dbolya/tomesd

[catboxanon]

Implementation here: https://git.mmaker.moe/mmaker/sd-webui-tome

[jifeng123guo]

it's lost

[wywywywy]

This is probably important enough that it should be a built-in feature rather than an external extension.

But of course it should only happen after @dbolya has published it as a module on pip.

[papuSpartan]

This is probably important enough that it should be a built-in feature rather than an external extension.

But of course it should only happen after @dbolya has published it as a module on pip.

I'm working on merging this with #9256.

[dbolya]

Hi, original author here.

Some quick suggestions:

• With xformers, tome for small images doesn't really increase speed much (see the benchmark some else did here: Token Merging (ToMe) for Stable Diffusion huggingface/diffusers#2940). So, I'd suggest having an option to only apply ToMe during the highres fix.
• Most options other than the ratio don't really need to be shown to the user, since there's not much point in changing them (both in terms of quality and speed). Idk what the best solution for this is though, since some people might still want to tinker with them.

Note that the tome patch can be applied right before generation just fine, it doesn't need to be on model load. Applying the patch is free (it just sets some class variables) and it can be applied any number of times to the same model without adverse effects (e.g., to change the parameters). The way I implemented it for testing was in the txt2img processing function, right before sampling.

I think the best way to implement it would be to add one of those boxes under the generation parameters, like the controlnet plugin.

[recoilme]

Not ready for Mac on Apple Silicone?
TypeError: Operation 'neg_out_mps()' does not support input type 'int64' in MPS backend