🧨 JAX / Flax での安定した拡散！

'🧨 Stable diffusion with JAX/Flax!'

🤗 Hugging Face Diffusersはバージョン0.5.1からFlaxをサポートしています！これにより、Colab、Kaggle、またはGoogle Cloud PlatformなどのGoogle TPU上での超高速な推論が可能になります。

この投稿では、JAX / Flaxを使用して推論を実行する方法を示します。Stable Diffusionの動作詳細やGPUでの実行方法について詳細を知りたい場合は、このColabノートブックを参照してください。

一緒に進める場合は、上のボタンをクリックしてこの投稿をColabノートブックとして開きます。

• ハギングフェイス推論エンドポイントの始め方
• MTEB 大規模テキスト埋め込みベンチマーク
• PyTorch DDPからAccelerateへ、そしてTrainerへ簡単に分散トレーニングをマスターしましょう

まず、TPUバックエンドを使用していることを確認してください。このノートブックをColabで実行している場合は、上のメニューでランタイムを選択し、「ランタイムのタイプを変更」オプションを選択し、ハードウェアアクセラレータの設定でTPUを選択します。

JAXはTPUに限定されているわけではありませんが、TPUサーバーごとに8つのTPUアクセラレータが並列に動作するため、そのハードウェア上で輝きます。

セットアップ

import jax
num_devices = jax.device_count()
device_type = jax.devices()[0].device_kind

print(f"Found {num_devices} JAX devices of type {device_type}.")
assert "TPU" in device_type, "Available device is not a TPU, please select TPU from Edit > Notebook settings > Hardware accelerator"

出力：

    Found 8 JAX devices of type TPU v2.

diffusersがインストールされていることを確認してください。

!pip install diffusers==0.5.1

次に、すべての依存関係をインポートします。

import numpy as np
import jax
import jax.numpy as jnp

from pathlib import Path
from jax import pmap
from flax.jax_utils import replicate
from flax.training.common_utils import shard
from PIL import Image

from huggingface_hub import notebook_login
from diffusers import FlaxStableDiffusionPipeline

モデルの読み込み

モデルを使用する前に、モデルのライセンスを承諾して重みをダウンロードし使用する必要があります。

ライセンスは、このような強力な機械学習システムの潜在的な有害な影響を軽減するために設計されています。ユーザーに対してライセンスの全文を注意深く読んでいただくようお願いします。以下は要約です：

1. モデルを意図的に違法または有害な出力やコンテンツを生成または共有するために使用することはできません。
2. 生成した出力に関して、私たちは権利を主張しません。それらを自由に使用することができ、使用に関してはライセンスで設定された規定に違反しないように責任を持つ必要があります。
3. 重みを再配布し、モデルを商業的におよび/またはサービスとして使用することができます。ただし、その場合、ライセンスの使用制限とCreativeML OpenRAIL-Mのコピーをすべてのユーザーに共有する必要があります。

Flaxの重みはStable Diffusionリポジトリの一部としてHugging Face Hubで利用できます。Stable DiffusionモデルはCreateML OpenRail-Mライセンスの下で配布されています。このオープンライセンスは、生成した出力に関して権利を主張せず、違法または有害なコンテンツを意図的に生成することを禁止しています。モデルカードには詳細が記載されているため、ライセンスを承諾するかどうかを慎重に検討し、読んでください。承諾する場合は、Hubの登録ユーザーであり、コードが機能するためのアクセストークンを使用する必要があります。アクセストークンを提供するには、次の2つのオプションがあります：

• ターミナルでhuggingface-cli loginコマンドラインツールを使用し、プロンプトにトークンを貼り付けます。トークンはコンピュータにファイルとして保存されます。
• または、ノートブックでnotebook_login()を使用します。これは同じことを行います。

次のセルは、このコンピュータで既に認証済みでない限り、ログインインターフェースを表示します。アクセストークンを貼り付ける必要があります。

if not (Path.home()/'.huggingface'/'token').exists(): notebook_login()

TPUデバイスはbfloat16、効率的なハーフフロートタイプをサポートしています。テストに使用しますが、代わりに完全な精度を持つfloat32を使用することもできます。

dtype = jnp.bfloat16

Flaxは関数型のフレームワークなので、モデルは状態を持たず、パラメータはそれらの外部に保存されます。事前学習されたFlaxパイプラインをロードすると、パイプライン自体とモデルの重み（またはパラメータ）の両方が返されます。私たちは重みのbf16バージョンを使用しており、これにより型の警告が発生しますが、安全に無視できます。

pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    revision="bf16",
    dtype=dtype,
)

推論

通常、TPUは8つのデバイスが並列に動作しているため、プロンプトをデバイスの数だけ複製します。その後、8つのデバイスで同時に推論を行い、各デバイスが1つの画像を生成する責任を持ちます。したがって、1つのチップが1つの画像を生成するのにかかる時間と同じ時間で、8つの画像を取得することができます。

プロンプトを複製した後、パイプラインのprepare_inputs関数を呼び出すことで、トークン化されたテキストのIDを取得します。トークン化されたテキストの長さは、基礎となるCLIPテキストモデルの設定によって77トークンに設定されています。

prompt = "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of field, close up, split lighting, cinematic"
prompt = [prompt] * jax.device_count()
prompt_ids = pipeline.prepare_inputs(prompt)
prompt_ids.shape

出力 :

    (8, 77)

複製と並列化

モデルのパラメータと入力は、8つの並列デバイスに複製する必要があります。パラメータ辞書はflax.jax_utils.replicateを使用して複製され、辞書をトラバースして重みの形状を8回繰り返すように変更します。配列はshardを使用して複製されます。

p_params = replicate(params)

prompt_ids = shard(prompt_ids)
prompt_ids.shape

出力 :

    (8, 1, 77)

その形状は、8つのデバイスのそれぞれが、形状が(1, 77)のjnp配列を入力として受け取ることを意味しています。したがって、1はデバイスごとのバッチサイズです。メモリが十分にあるTPUでは、1つのチップで複数の画像（チップごと）を生成したい場合、バッチサイズは1よりも大きくなる可能性があります。

画像を生成する準備がほぼ整いました！画像生成関数に渡すためのランダムな数値ジェネレータを作成する必要があります。これはFlaxの標準的な手続きであり、ランダムな数値に関連するすべての関数はジェネレータを受け取ることが期待されています。これにより、複数の分散デバイスでトレーニングしている場合でも再現性が確保されます。

以下のヘルパー関数は、シードを使用してランダムな数値ジェネレータを初期化します。同じシードを使用すれば、まったく同じ結果を得ることができます。後でノートブックで結果を調べる際には、異なるシードを使用しても構いません。

def create_key(seed=0):
    return jax.random.PRNGKey(seed)

ジェネレータを取得し、それを8回「分割」して各デバイスが異なるジェネレータを受け取るようにします。したがって、各デバイスは異なる画像を作成し、全体のプロセスは再現可能です。

rng = create_key(0)
rng = jax.random.split(rng, jax.device_count())

JAXのコードは、非常に高速に実行される効率的な表現にコンパイルできます。ただし、後続の呼び出しですべての入力が同じ形状であることを確認する必要があります。そうでない場合、JAXはコードを再コンパイルする必要があり、最適化された速度を活用することができません。

Flaxパイプラインは、引数としてjit = Trueを渡すと、コードをコンパイルしてくれます。また、モデルが8つの利用可能なデバイスで並列に実行されるようにもします。

次のセルを実行するのは最初の一回だけで、コンパイルには時間がかかりますが、それ以降の呼び出し（異なる入力でも）ははるかに速くなります。例えば、私がテストしたTPU v2-8では、コンパイルに1分以上かかりましたが、その後の推論実行には約7秒かかります。

images = pipeline(prompt_ids, p_params, rng, jit=True)[0]

出力 :

    CPU 時間: ユーザー 464 ms、システム: 105 ms、合計: 569 ms
    ウォールタイム: 7.07 s

返された配列の形状は (8, 1, 512, 512, 3) です。2番目の次元を取り除いて、512 × 512 × 3 の8つの画像を取得し、それらをPIL形式に変換します。

images = images.reshape((images.shape[0],) + images.shape[-3:])
images = pipeline.numpy_to_pil(images)

可視化

画像をグリッド状に表示するためのヘルパー関数を作成しましょう。

def image_grid(imgs, rows, cols):
    w,h = imgs[0].size
    grid = Image.new('RGB', size=(cols*w, rows*h))
    for i, img in enumerate(imgs): grid.paste(img, box=(i%cols*w, i//cols*h))
    return grid

image_grid(images, 2, 4)

異なるプロンプトの使用

すべてのデバイスで同じプロンプトを複製する必要はありません。どんなことでもできます：2つのプロンプトを4回生成する、または一度に8つの異なるプロンプトを生成することさえできます。それをやってみましょう！

まず、入力の準備コードを便利な関数にリファクタリングします：

prompts = [
    "Labrador in the style of Hokusai",
    "Painting of a squirrel skating in New York",
    "HAL-9000 in the style of Van Gogh",
    "Times Square under water, with fish and a dolphin swimming around",
    "Ancient Roman fresco showing a man working on his laptop",
    "Close-up photograph of young black woman against urban background, high quality, bokeh",
    "Armchair in the shape of an avocado",
    "Clown astronaut in space, with Earth in the background",
]

prompt_ids = pipeline.prepare_inputs(prompts)
prompt_ids = shard(prompt_ids)
images = pipeline(prompt_ids, p_params, rng, jit=True).images
images = images.reshape((images.shape[0], ) + images.shape[-3:])
images = pipeline.numpy_to_pil(images)
image_grid(images, 2, 4)

---

並列化はどのように機能しますか？

以前に述べたように、diffusers Flaxパイプラインはモデルを自動的にコンパイルし、利用可能なすべてのデバイスで並列に実行します。ここではそのプロセスの内部を簡単に見て、それがどのように機能するかを示します。

JAXの並列化は複数の方法で行うことができます。もっとも簡単な方法は、jax.pmap 関数を使用して単一プログラム、複数データ（SPMD）並列化を実現することです。これは、同じコードの複数のコピーを異なるデータ入力で実行することを意味します。より高度なアプローチも可能ですが、興味がある場合はJAXのドキュメントとpjit のページを参照してこのトピックを探索することをお勧めします。

jax.pmap は次の2つのことを行います：

• コードをコンパイル（またはjit する）こと。これはpmap を呼び出したときには行われませんが、最初にpmapped関数が呼び出されるときに行われます。
• コンパイルされたコードがすべての利用可能なデバイスで並列に実行されるようにします。

それがどのように機能するかを示すために、パイプラインの_generate メソッドをpmap します。これは、画像を生成するプライベートメソッドです。注意してください、このメソッドは将来のdiffusers のリリースで名前が変更されるか削除される可能性があります。

p_generate = pmap(pipeline._generate)

pmap を使用した後、準備された関数p_generate は概念的に次のことを行います：

• 各デバイスで基礎となる関数pipeline._generate のコピーを呼び出します。
• 各デバイスに異なる部分の入力引数を送信します。これにはシャーディングが使用されます。この例では、prompt_ids の形状は(8, 1, 77, 768) です。この配列は8に分割され、各_generate のコピーは形状(1, 77, 768)の入力を受け取ります。

私たちは、_generate を並列で呼び出されることを無視して完全にコード化することができます。この例では、バッチサイズ（1）とコードに意味がある次元に関心を持ち、並列で動作させるために何も変更する必要はありません。

パイプライン呼び出しを使用した場合と同様に、最初に以下のセルを実行すると時間がかかりますが、その後ははるかに高速になります。

images = p_generate(prompt_ids, p_params, rng)
images = images.block_until_ready()
images.shape

出力：

    CPU times: user 118 ms, sys: 83.9 ms, total: 202 ms
    Wall time: 6.82 s

    (8, 1, 512, 512, 3)

JAXは非同期ディスパッチを使用し、できるだけ早くPythonループに制御を返すため、推論時間を正しく測定するためにblock_until_ready()を使用しています。コードでそれを使用する必要はありません。まだ具現化されていない計算の結果を使用する場合には、自動的にブロッキングが発生します。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!