Hugging Face Diffusers 支持 JAX/Flax 实现 Stable Diffusion TPU 并行推理

Hugging Face 在 Diffusers 0.5.1 中加入对 JAX/Flax 的支持，允许用户在 Google TPU 上快速运行 Stable Diffusion 推理。利用 TPU 包含 8 个加速核心的特性，结合 jax.pmap 和 JIT 编译，系统可将参数与提示词分片后并行生成 8 张图片；在 TPU v2-8 上单次推理测试耗时约 7 秒。