Settings

Theme

Accelerating Stable Diffusion XL Inference with Jax on Cloud TPU v5e

huggingface.co

23 points by rayshan 2 years ago · 6 comments

Reader

MasterScrat 2 years ago

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"!

brucethemoose2 2 years ago

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.

  • FloatArtifact 2 years ago

    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.

    • brucethemoose2 2 years ago

      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?

iAkashPaul 2 years ago

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

Keyboard Shortcuts

j
Next item
k
Previous item
o / Enter
Open selected item
?
Show this help
Esc
Close modal / clear selection