Stable Diffusion on Tpu using Colab

Flax/JAX Projects
1

I’ve seen a few implementations of this but haven’t been able to get it to run properly using FlaxStableDiffusionPipeline… it froze and took forever.
Any help getting this installed is greatly appreciated.