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

## はじめに

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

https://qiita.com/the_haigo/items/cbec41b5d1fd0db76ba4

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

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

このシリーズの記事

- [画像分類: ResNet50](https://qiita.com/RyoWakabayashi/items/b9b4799683af5aaa29dd)
- 画像生成: Stable Diffusion (ここ)
- [文章の穴埋め: BERT](https://qiita.com/RyoWakabayashi/items/dcc03a8d19d7320dda8f)
- [文章の判別: BERTweet](https://qiita.com/RyoWakabayashi/items/1f8624d717fb2b1dbdc0)
- [文章の生成: GPT2](https://qiita.com/RyoWakabayashi/items/6a780c2858ec952fa1e1)
- [質疑応答: RoBERTa](https://qiita.com/RyoWakabayashi/items/316d0c24020468d07f18)
- [固有名詞の抽出: bert-base-NER](https://qiita.com/RyoWakabayashi/items/82ac1dddfdddc8a2f79d)

実装の全文はこちら

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/bumblebee/stable_diffusion.livemd

## 実行環境

- 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　イメージを元にしたコンテナで動かしました

https://hub.docker.com/layers/livebook/livebook/0.8.0/images/sha256-6ef3fdd439fd72af9ae01d58e947664a46a8493f031e179cc0e91c08ddac02e5?context=explore

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

https://github.com/RyoWakabayashi/elixir-learning

## セットアップ

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

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

## 設定

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

```elixir
Nx.default_backend()
```

![スクリーンショット 2022-12-12 16.19.17.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/b77d4afc-7d1a-592b-5635-53825e013d36.png)

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

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

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

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

https://huggingface.co/CompVis/stable-diffusion-v1-4

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

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

```elixir
cache_dir = "/tmp/bumblebee_cache"
```

::: note info

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

```yml
...
    volumes:
      - ./tmp:/tmp
```

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

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

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

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

```text
could not rename ~ cross-domain link ~
```

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

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

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

:::

## モデルのダウンロード

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

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

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

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

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

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

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

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

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

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

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

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

## 画像生成の実行

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

```elixir
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]
  )
```

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

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

![スクリーンショット 2022-12-12 16.50.43.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/cc63a757-2c61-e0d4-fd7e-7a619f44b0ca.png)

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

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

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

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

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

実行結果を表示します

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

![スクリーンショット 2022-12-12 17.18.33.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/bd61a3e3-105a-f745-9fcb-82e25821b40f.png)

## まとめ

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