With Flax, it’s possible to train a Stable Diffusion model faster on TPUs and GPUs thanks to @duongna211. This is very efficient on TPU hardware but works great on GPUs too. The Flax training script doesn’t support features like gradient checkpointing or gradient accumulation yet, so you’ll need a GPU with at least 30GB of memory or a TPU v3.