Stable Diffusion JAX と Flax の統合
Hugging Face はバージョン 0.5.1 から diffusers ライブラリに Flax のサポートを統合し、Google TPU 上で Stable Diffusion を高効率で実行できるようにしました。この統合により、ユーザーは通常 8 台のアクセラレータを備える TPU サーバーの並列処理能力を活用し、1 回の生成にかかる時間で複数の画像を同時に生成できます。
JAX と Flax を用いた高速 TPU 推論
TPU 上での Stable Diffusion 推論は JAX と Flax を使用して最適化され、標準的な GPU 実装と比較して大幅な高速化が実現します。TPU v2-8 では、初回のコンパイル後の推論は約 7 秒 かかります。
Key technical optimizations include:
- bfloat16 Precision: TPU デバイスは
bfloat16をサポートしており、メモリオーバーヘッドを削減しつつ性能を維持する効率的な半精度浮動小数点型です。 - JIT Compilation: Flax パイプラインに
jit=Trueを渡すことで、JAX はモデルを効率的な表現にコンパイルします。最初の実行はコンパイルに時間がかかります（TPU v2-8 では 1 分以上）、その後の呼び出しは大幅に高速です。 - Stateless Models: Flax は関数型フレームワークであるため、モデルはステートレスで、パラメータはモデル本体の外部に保存されます。
SPMD による並列化
diffusers の Flax パイプラインは Single-Program, Multiple-Data (SPMD) 並列化を利用して TPU ハードウェアの使用率を最大化します。これは主に jax.pmap 関数によって実現されます。
並列化の実装方法
jax.pmap は 2 つの重要な機能を実行します：コードをコンパイルする（jax.jit() に似ています）ことと、コンパイルされたコードが利用可能なすべてのデバイスで並列に実行されることを保証することです。
この並列実行のために、パイプラインは以下の手順を踏みます：
- Replication: モデルパラメータは
flax.jax_utils.replicateを使用してすべてのデバイスに複製されます。 - Sharding: トークン化されたプロンプト ID などの入力データは
shardを使用して分割されます。例えば、8 台のデバイスがある場合、プロンプト配列は分割され、各デバイスが入力の特定の部分を受け取ります。 - PRNG Handling: 生成画像の再現性と多様性を確保するために、乱数生成器 (RNG) を作成し、各デバイス用に複数の生成器に分割します。
このアーキテクチャにより、各デバイスがバッチ項目を独立して処理するため、パイプラインは 8 枚の異なる画像（または同一画像の 8 コピー）を同時に生成できます。
モデルへのアクセスとライセンス
Flax 用の Stable Diffusion 重みは CompVis/stable-diffusion-v1-4 リポジトリの Hugging Face Hub で入手可能です。アクセスには CreativeML OpenRAIL-M ライセンスへの同意が必要で、以下の条件が含まれます：
- ユーザーは違法または有害なコンテンツを意図的に生成・共有するためにモデルを使用してはなりません。
- ユーザーは生成した出力に対する権利を保持し、その使用について責任を負います。
- 商用利用および重みの再配布は許可されますが、同じ使用制限と CreativeML OpenRAIL-M ライセンスのコピーをすべてのユーザーと共有する必要があります。