Elixir Bumblebee のStable Diffusionを Docker 上の Livebook で動かして画像生成してみた

はじめに

@the_haigo さんの記事を参考に Docker 上で動かした記事です

Bumblebee　を使うと、 Stable Diffusion を使った画像生成が Elixir で超簡単に実装できます

他にも Bumblebee でできることを一通りやってみました

このシリーズの記事

• 画像分類: ResNet50
• 画像生成: Stable Diffusion (ここ)
• 文章の穴埋め: BERT
• 文章の判別: BERTweet
• 文章の生成: GPT2
• 質疑応答: RoBERTa
• 固有名詞の抽出: bert-base-NER

実装の全文はこちら

実行環境

• MacBook Pro 13 inchi
  • 2.4 GHz クアッドコアIntel Core i5
  • 16 GB 2133 MHz LPDDR3
• macOS Ventura 13.0.1
• Rancher Desktop 1.6.2
  • メモリ割り当て 12 GB
  • CPU 割り当て 6 コア

Livebook 0.8.0 の Docker　イメージを元にしたコンテナで動かしました

コンテナ定義はこちらを参照

セットアップ

必要なモジュールをインストールし、 Nx のバックエンドを EXLA.Backend に指定します

Mix.install(
  [
    {:bumblebee, "~> 0.1"},
    {:nx, "~> 0.4"},
    {:exla, "~> 0.4"},
    {:kino, "~> 0.8"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)

設定

Nx バックエンドを確認します

Nx.default_backend()

ちゃんと EXLA.Backend になっています

Stable Diffusion モデルのダウンロード元リポジトリーを指定します

repository_id = "CompVis/stable-diffusion-v1-4"

このように指定すると、以下の Hugging Face のリポジトリーからダウンロードしてきます

モデルファイルのキャッシュ用ディレクトリーを指定します

ここは Docker ならではの設定です（詳細は以下の注釈参照）

cache_dir = "/tmp/bumblebee_cache"

私の用意した実行環境だと docker-compose.yml で以下のように設定しています

...
    volumes:
      - ./tmp:/tmp

つまりローカルの ./tmp をコンテナ上の /tmp にマウントしています

この状態で Bumblebee のモデルロードを動かすと、まずコンテナ上の /tmp にモデルファイルをダウンロードします

キャッシュディレクトリーを指定していない場合、ダウンロード完了後、 ~/.cache 配下にダウンロードしたファイルを移動しようとします

このとき、 /tmp は macOS のファイルシステムで、 ~/.cache は Linux のファイルシステムになっている関係で、以下のようなエラーが発生します

could not rename ~ cross-domain link ~

これは Elixir に限らず Docker 上だと度々発生する現象です

ファイルシステムが同じ /tmp 内であればコピー可能なので、キャッシュディレクトリーを /tmp/bumblebee_cache に指定しておきます

また、こうすることでコンテナをビルドし直してもキャッシュがローカルに残るため、ダウンロードし直さなくて良くなります

モデルのダウンロード

以下のコードを逐次実行して Hugging Face からモデルファイルをダウンロードし、読み込みます

{:ok, tokenizer} =
  Bumblebee.load_tokenizer({
    :hf,
    "openai/clip-vit-large-patch14",
    cache_dir: cache_dir
  })

{:ok, clip} =
  Bumblebee.load_model({
    :hf,
    repository_id,
    subdir: "text_encoder", cache_dir: cache_dir
  })

{:ok, unet} =
  Bumblebee.load_model(
    {
      :hf,
      repository_id,
      subdir: "unet", cache_dir: cache_dir
    },
    params_filename: "diffusion_pytorch_model.bin"
  )

{:ok, vae} =
  Bumblebee.load_model(
    {
      :hf,
      repository_id,
      subdir: "vae", cache_dir: cache_dir
    },
    architecture: :decoder,
    params_filename: "diffusion_pytorch_model.bin"
  )

{:ok, scheduler} =
  Bumblebee.load_scheduler({
    :hf,
    repository_id,
    subdir: "scheduler", cache_dir: cache_dir
  })

{:ok, featurizer} =
  Bumblebee.load_featurizer({
    :hf,
    repository_id,
    subdir: "feature_extractor", cache_dir: cache_dir
  })

{:ok, safety_checker} =
  Bumblebee.load_model({
    :hf,
    repository_id,
    subdir: "safety_checker", cache_dir: cache_dir
  })

初回は以下のような進行状況が表示されて 100% になればダウンロード終了です

|=============================================================| 100% (334.70 MB)

2回目以降はキャッシュが読まれるので、モデルファイルのロードだけが実行されます

ちなみにモデルファイルは合計 5.48 GB あるので、ストレージの空き容量には気をつけましょう

画像生成の実行

Nx 0.4.1 で追加された Nx.Serving を利用して、画像生成サービスを提供します

serving =
  Bumblebee.Diffusion.StableDiffusion.text_to_image(
    clip,
    unet,
    vae,
    tokenizer,
    scheduler,
    num_steps: 20,
    num_images_per_prompt: 2,
    safety_checker: safety_checker,
    safety_checker_featurizer: featurizer,
    compile: [batch_size: 1, sequence_length: 60],
    defn_options: [compiler: EXLA]
  )

画像生成の元となるプロンプトの入力エリアを作ります

prompt_input = Kino.Input.text("PROMPT")

今回は以下のような文言を入れてみました

Cartoon of a blue raccoon-like robot pulling futuristic tools out of its pockets.

画像生成サービスにプロンプトを渡して実行します

※私の M1 ではない MacBook では10分以上かかります

output = Nx.Serving.run(serving, Kino.Input.read(prompt_input))

実行結果を表示します

output.results
|> Enum.map(fn result ->
  Kino.Image.new(result.image)
end)
|> Kino.Layout.grid(columns: 2)

まとめ

コンテナ特有の気をつけるべき点はありますが、 Docker 上で実行できました