Stable Diffusion XL inference on Cloud TPU v5e via Diffusers JAX integration

AI Impact Summary

Diffusers Diffusion JAX integration now supports SDXL on Cloud TPU v5e, using JAX/JIT and pmap to parallelize across TPU chips. In tests, four TPU v5e-4 chips generated four 1024×1024 images in about 4 seconds, with actual generation around 2.3 seconds, by leveraging the FlaxStableDiffusionXLPipeline and 16‑bit parameter storage (bfloat16) while keeping the scheduler state in float32. However, the approach relies on static input/output shapes to avoid recompilation, so deployments must use fixed prompts and image sizes or accept compilation overhead and latency spikes when shapes change.