Accelerating Stable Diffusion XL Inference with Jax on Cloud TPU v5e

huggingface.co

23 points·rayshan··6 comments

TPUs are amazing for Stable Diffusion.

We've been doing training (Dreambooth) and inference on TPUs since the beginning of the year at https://dreamlook.ai.

We basically get 2.5x the training speed for Stable Diffusion 1.5 compared to A100, a very nice "unfair advantage"!

That's very fast!

I just tried it on my RTX 3090, with a riced linux environment + pytorch/xformers nightly, and 4 images take 36.7 seconds on the ComfyUI backend (used by Fooocus-MRE).

...But the issue is, right now, you can either pick high quality tooling (The ComfyUI/automatic1111 backend UIs) or speed (diffusers-based UIs), not both. InvokeAI and VoltaML do not support SDXL as well as Fooocus at the moment, and all the other UIs use the slow Stability backend with no compilation support.

Double check your facts there I think InvokeAI support SDXL

See https://github.com/invoke-ai/InvokeAI/releases

I believe invoke supports caching and a number of interesting features this release.

Yeah, I just tried the release candidate two days ago, but my results with SDXL were very poor compared to Fooocus, and its missing other niceties like the known style presets, FreeU, the better Fooocus prompt expansion...

I really wanted to use Invoke for the better performance and torch.compile support, but the switch to the refiner seems slow, and torch.compile doesn't seem to work either. I need to investigate torch.compile more, but maybe they did something to the vanilla diffusers pipeline?

Thanks for your perspective here. I'll definitely look into some of these other projects.

With some CPU offloading I'm able to run SDXL at 1.5it/s on an RTX2070S w/ 8GB VRAM.

When used with ControlNet it still runs but more layers offloaded at 1.4s/it

Reply on news.ycombinator.com