Stable Diffusion JAX and Flax Integration

Hugging Face has integrated Flax support into the diffusers library starting with version 0.5.1, enabling Stable Diffusion to run with high efficiency on Google TPUs. This integration allows users to leverage the parallel processing power of TPU servers—which typically feature eight accelerators—to generate multiple images simultaneously in the time it takes to generate one.

High-Speed TPU Inference with JAX and Flax

Stable Diffusion inference on TPUs is optimized through the use of JAX and Flax, allowing for significant speedups compared to standard GPU implementations. On a TPU v2-8, subsequent inference runs after the initial compilation take approximately 7 seconds.

Key technical optimizations include:

  • bfloat16 Precision: TPU devices support bfloat16, an efficient half-float type that reduces memory overhead while maintaining performance.
  • JIT Compilation: By passing jit=True to the Flax pipeline, JAX compiles the model into an efficient representation. While the first run requires a compilation period (exceeding one minute on TPU v2-8), all subsequent calls are significantly faster.
  • Stateless Models: Because Flax is a functional framework, models are stateless, meaning parameters are stored outside the model itself.

Parallelization via SPMD

The diffusers Flax pipeline utilizes Single-Program, Multiple-Data (SPMD) parallelization to maximize TPU hardware utilization. This is primarily achieved through the jax.pmap function.

How Parallelization is Implemented

jax.pmap performs two critical functions: it compiles the code (similar to jax.jit()) and ensures the compiled code runs in parallel across all available devices.

To execute this in parallel, the pipeline follows these steps:

  1. Replication: Model parameters are replicated across all devices using flax.jax_utils.replicate.
  2. Sharding: Input data, such as tokenized prompt IDs, are sharded using shard. For example, if there are 8 devices, a prompt array is split so each device receives a specific portion of the input.
  3. PRNG Handling: To ensure reproducibility and diversity in generated images, a random number generator (RNG) is created and split into multiple generators—one for each device.

This architecture allows the pipeline to generate eight different images (or eight copies of the same image) simultaneously, as each device processes one batch item independently.

Model Access and Licensing

Stable Diffusion weights for Flax are available on the Hugging Face Hub under the CompVis/stable-diffusion-v1-4 repository. Access requires accepting the CreativeML OpenRAIL-M license, which includes the following stipulations:

  • Users may not use the model to deliberately produce or share illegal or harmful content.
  • Users retain rights to the outputs they generate and are accountable for their use.
  • Commercial use and redistribution of weights are permitted, provided the same use restrictions and a copy of the CreativeML OpenRAIL-M license are shared with all users.

Sources