running on tpu

[entrpn]

This PR adds the ability to run on TPUs, although its very slow right now. I'm unsure how to this hooks into uv as I used pip. Please review.

• Loads encoders on CPU when device_type == 'tpu'.
• Can run on 32GB of HBM, i.e., TPUv4.

[SauravMaheshkar]

Thanks a lot for the PR @entrpn. I've had a look at the diff and it makes sense to me. I'll just confirm by running this on TPU later today.

[ariG23498]

LGTM!

I have asked for some TPU support. Once I have them, I can check this PR.

On another note I have the following queries:

1. We have FLAX implementation of T5 and CLIP in Hugging Face, would it be easier for us to load them instead of PyTorch versions?
2. Do you see any place we can apply parallelization to make this implementation faster? I understand that this is a more dense question which requires one to look into the code, if that is too much to ask at this time I completely understand if you don't want to answer it now.

[ariG23498]

I had checked the implementation on my side on tpu v4.

Thanks for the great contribution. ❤️

[cataluna84]

LGTM 🤩

This conversation to use TPUs looks fantastic

[entrpn]

LGTM!

I have asked for some TPU support. Once I have them, I can check this PR.

On another note I have the following queries:

1. We have FLAX implementation of T5 and CLIP in Hugging Face, would it be easier for us to load them instead of PyTorch versions?
2. Do you see any place we can apply parallelization to make this implementation faster? I understand that this is a more dense question which requires one to look into the code, if that is too much to ask at this time I completely understand if you don't want to answer it now.

I will take some time to understand the code better. This should be able to run faster and be parallelizable across all TPU devices.

For reference, I created a pytorch xla version that runs flux on TPUs and generates 4 images in parallel in under 10 seconds. https://github.com/entrpn/diffusers/tree/flux_ptxla