Stable Diffusion 在 JAX/Flax 上运行

Hugging Face 博客宣布，Stable Diffusion 现已支持在 JAX/Flax 上运行，利用 TPU 加速推理。该实现基于 Diffusers 库，提供了更快的生成速度和更低的延迟，适用于大规模部署。用户可以通过简单的 API 调用在 TPU 上运行模型，并支持批处理。博客还提供了性能基准测试，显示在 TPU v4 上生成单张图像仅需约 0.8 秒，相比 GPU 有显著提升。该实现已开源，并提供了 Colab 笔记本供开发者快速上手。