🦒 Colab TPU & GPU 🦆 Kaggle TPU & GPU with 🤗 HuggingFace Diffusers

[camenduru]

https://github.com/camenduru/stable-diffusion-diffusers-colab

I also converted popular models for tpu

🤗 HuggingFace Path	Model Page
flax/anything-v3.0	https://huggingface.co/Linaqruf/anything-v3.0
flax/samsamart	https://huggingface.co/Danzelus/samsamart
flax/Stable_Diffusion_PaperCut_Model	https://huggingface.co/Fictiverse/Stable_Diffusion_PaperCut_Model
flax/redshift-diffusion	https://huggingface.co/nitrosocke/redshift-diffusion
flax/sd15	https://huggingface.co/runwayml/stable-diffusion-v1-5
flax/sd15-non-ema	https://huggingface.co/runwayml/stable-diffusion-v1-5
flax/midjourney-v4-diffusion	https://huggingface.co/prompthero/midjourney-v4-diffusion
flax/StudioGhibli	https://huggingface.co/IShallRiseAgain/StudioGhibli
flax/Tron-Legacy-diffusion	https://huggingface.co/dallinmackay/Tron-Legacy-diffusion
flax/Cyberpunk-Anime-Diffusion	https://huggingface.co/DGSpitzer/Cyberpunk-Anime-Diffusion
flax/waifu-diffusion	https://huggingface.co/hakurei/waifu-diffusion-v1-3
flax/Arcane-Diffusion	https://huggingface.co/nitrosocke/Arcane-Diffusion
flax/classic-anim-diffusion	https://huggingface.co/nitrosocke/classic-anim-diffusion
flax/mo-di-diffusion	https://huggingface.co/nitrosocke/mo-di-diffusion
camenduru/plushies	https://huggingface.co/camenduru/plushies

[ClashSAN]

I remember someone made it really fast for TPUs. It is like 1 second for 1 image. Were you able to get to such speeds? https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/stable_diffusion_fast_jax.ipynb

[ClashSAN]

incredible.

[camenduru]

https://www.youtube.com/watch?v=X6AousudHOk 🚀

[ClashSAN]

@camenduru Are there other sampler options available? Or there needs to be an update to diffusers pipeline first?

[camenduru]

Hi @ClashSAN Flax Pipe has DDIM, DPMSolverMultistep, LMSDiscrete, PNDM
PT Pipe has DDIM, DPMSolverMultistep, EulerAncestralDiscrete, EulerDiscrete, LMSDiscrete, PNDM
🐢 soon I will add samplers and embeddings functionality in colab code

[B0rner]

This colab-example isn't workin anymore.
"As of JAX 0.4.0, JAX only supports TPU VMs, not the older Colab TPUs."

Currently I haven't found a single way to get sd to work with a TPU.
Can anyone help and maybe have a script ready to use to generate images with TPU hardware acceleration?

[donlinglok1]

Great Work! I have tested in the colab get the below result:

anything-v3 model
512x512
28 steps
11 guidance_scale

= 5 sec for 2x4 images gird

I wonder if I can change it to 1 or 16 images but it seems the flax is using jax.local_devices (in colab is 8).
I will keep studying for it...

Anyway, thank you for you sharing again!

[camenduru]

@donlinglok yes colab free TPU v2-8 kaggle v3-8 and you can generate 8192 images under 5 seconds with v4-8192 little expensive https://cloud.google.com/tpu/pricing 🤣 but the good news is we can add PCIe card in our pc IBM did this https://www.youtube.com/watch?v=hCxvS1dVufs maybe 2 years maybe 1 🙄

[donlinglok1]

@donlinglok yes colab free TPU v2-8 kaggle v3-8 and you can generate 8192 images under 5 seconds with v4-8192 little expensive https://cloud.google.com/tpu/pricing 🤣 but the good news is we can add PCIe card in our pc IBM did this https://www.youtube.com/watch?v=hCxvS1dVufs maybe 2 years maybe 1 🙄

holy shit🤯 8000+ images under 5 sec!

[donlinglok1]

...
I wonder if I can change it to 1 or 16 images but it seems the flax is using jax.local_devices (in colab is 8). I will keep studying for it...

Now I can see I was in the wrong direction, I can't let 8 cores draw me 1 image, but I may can generate 8 dynamic prompts to 8 cores for drawing.🤩

    # with fixed seed
    prompt_n = num_samples * [prompt]
    tempPromptList = prompt.split(",")
    for index in range(len(tempPromptList)):
      if(index>=num_samples-1):
        break;
      if(index>0):
        prompt_n[index] = prompt_n[index - 1] +","+ tempPromptList[index];
      if(index==0):
        prompt_n[index] = tempPromptList[index];
    print(prompt_n)
    prompt_ids = pipe.prepare_inputs(prompt_n)
    prompt_ids = shard(prompt_ids)

OR a queue that runs every 5 sec with 8 seats.
(and in the future the 8 seats may scale to 8192🤯)

[mudapanda2]

Shouldn't be too hard to swap this into the automatic1111 webui no?
I'll give it a shot

[ClashSAN]

what do you mean? you use @TheLastBen 's ?
straight up implementing it here ?
Extension?

[SKYwebui]

@mudapanda2 Did you succeed in running automatic1111's webui with tpu?

[JokingLemon]

@camenduru can you share how you converted the models to flax type? I'm trying to convert some of the models I like to use