Stable Diffusionをクラウド環境(Modal)で動かす

|
featured image

手元の PC で stable-diffusion-webui を動かしましたが、画像の生成速度が遅いため、より高いスペックで実行できるクラウド環境を使ってみることにしました。 また、今回は webui を使わず、Modal のチュートリアルを参考に Pythonでコーディングしました。CLIから手軽に実行でき、ControlNet 1.1 TileとUpscalerによる高解像度な画像を生成できるよう工夫しました。モデルの取り替えもある程度自由に行うことができます。

Modal

Modal はモデル推論、バッチ処理、タスクキュー、ウェブアプリケーションなどをクラウド上で実行できるサービスです。 利用してみた所管ですが、AWS Lambda のような、serverless アーキテクチャに近いサービスという印象。Dockerで推論に必要な環境をModal上に構築し、実装したPythonのクラスや関数を登録、実行できます。

料金は月 30$まで無料(2023/07/16時点)。 無料枠を超える場合は従量課金です。デフォルトで 30$の上限が設定されていて、支払い方法を登録していない限りはクラウド破産することはないです。2ヶ月利用してみましたが、趣味で使う分には十分な容量です。しかし非常に多くの画像を生成したいヘビーユーザーは、Google Colab や Paperspace など他の定額制のサービスと比較して選ぶのが良いと思います。

事前準備

  1. Modal のアカウント申請
  2. API トークンの発行

Modalでアカウント申請。waitlist に登録されて、半日ほどでアカウントが払い出されます。

API トークンの発行

アカウントが払い出されたあと、Modal のホーム画面に Getting Started があり、API トークンの発行から動作確認までの手順があります。

Stable Diffusion CLI を Modal 上で実行するスクリプト

txt2imgの推論、ControlNet 1.1 TileとReal-ESRGANを利用したアップスケール機能を搭載したアプリケーションを実装しました。GitHub のリポジトリ。

下記の手順で画像が生成され、outputs ディレクトリに出力されます。

  1. リポジトリをgit clone
  2. .envファイルを作成し、.env.example を参考に huggingface の API トークンとモデルを設定
  3. ./setup_files/config.sample.yml を ./setup_files/config.ymlにコピー
  4. Makefile を開いてプロンプトを設定
  5. make deployをコマンドラインで実行(Modal上にアプリケーションが構築されます)
  6. make run(スクリプトが起動します)

スクリプトの説明

ディレクトリ構成

 1.
 2├── .env                    # huggingfaceのAPIトークンを設定するファイル
 3├── Makefile                # タスクランナー
 4├── sdcli/                  # アプリケーションを起動するためのエントリーポイントがあるディレクトリ
 5│   ├── outputs/            # 画像の出力先
 6│   ├── txt2img.py          # txt2imgを実効するエントリーポイント
 7│   └── util.py             # utilファイル
 8└── setup_files/            # アプリケーションを構築するためのコードがあるディレクトリ
 9    ├── __main__.py         #
10    ├── Dockerfile          # アプリケーションの基盤部分
11    ├── config.yml          # 推論に使うモデル、VAE、ControlNetなどを設定するファイル
12    ├── requirements.txt    # 必要なライブラリのリスト
13    ├── setup.py            # Dockerfileとrequirements.txt、config.yamlを元にビルドするためのスクリプト
14    └── txt2img.py          # 推論を実行するためのクラスとメソッドが実装されているファイル

setup_files/のファイルについて

アプリケーションの構築部分であるsetup_files/の中身についてまとめます。

config.yml

推論に使うモデル、VAE、ControlNetなどを設定します。Textual InversionやLoRAも設定できます（詳細はリポジトリのファイルを参照のこと）。

ControlNetがarrayですが、複数指定できるコードにはなっていないです(あとで修正予定)。

詳細
1model:
2  name: stable-diffusion-2-1
3  repo_id: stabilityai/stable-diffusion-2-1
4vae:
5  name: sd-vae-ft-mse
6  repo_id: stabilityai/sd-vae-ft-mse
7controlnets:
8  - name: control_v11f1e_sd15_tile
9    repo_id: lllyasviel/control_v11f1e_sd15_tile

Dockerfile

python:3.11が乗っかっているdebian-slimのイメージをベースに、requirements.txtに記載されているライブラリのインストールとReal-ESRGANのダウンロードを行います。

詳細
 1FROM python:3.11.3-slim-bullseye
 2COPY ./requirements.txt /
 3RUN apt update \
 4    && apt install -y wget git libgl1-mesa-glx libglib2.0-0 \
 5    && apt autoremove -y \
 6    && apt clean -y \
 7    && pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu117 \
 8    && mkdir -p /vol/cache/esrgan \
 9    && wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P /vol/cache/esrgan \
10    ...省略

requirements.txt

Stable Diffusion を動かすための必須ライブラリのリスト。

詳細
 1accelerate
 2diffusers[torch]==0.17.1
 3onnxruntime==1.15.1
 4safetensors==0.3.1
 5torch==2.0.1+cu117
 6transformers==4.30.2
 7xformers==0.0.20
 8
 9realesrgan==0.3.0
10basicsr>=1.4.2
11facexlib>=0.3.0
12gfpgan>=1.3.8
13numpy
14opencv-python
15Pillow
16torchvision
17tqdm
18
19controlnet_aux
20pyyaml

setup.py

Dockerfileとrequirements.txt、およびconfig.yamlを元にビルドを実行し、dockerイメージをModal上に構築します。コードが長いので、詳細はリポジトリのファイルを参照のこと。

このファイルで工夫した点は下記の箇所です。

 1stub = Stub("stable-diffusion-cli")
 2base_stub = Image.from_dockerfile(
 3    path="Dockerfile",
 4    context_mount=Mount.from_local_file("requirements.txt"),
 5)
 6stub.image = base_stub.extend(
 7    dockerfile_commands=[
 8        "FROM base",
 9        "COPY config.yml /",
10    ],
11    context_mount=Mount.from_local_file("config.yml"),
12).run_function(
13    build_image,
14    secrets=[Secret.from_dotenv(__file__)],
15)

Dockerfileとrequirements.txtでベースのイメージを作った後、config.ymlに記載したモデル類をダウンロードするマルチステージビルドを行っています。このように実装することで、モデルだけ変更した際に0からdockerイメージを構築することなく、config.ymlの差分だけ再ビルドしてくれます。

txt2img.py

推論を実行するためのクラスとメソッドを実装したファイルです。こちらも長いので、詳細はリポジトリのファイルを参照のこと。

のちのちコードが肥大化していきそうなので、interfaceを実装しています。これでテストコードも書きやすくなると思います（書くとは言っていない）。

単純なtxt2imgではなく、ControlNet 1.1TileとReal-ESRGANを用いて、高解像度な画像を生成できるように工夫しました。高解像度化の詳細についてはこちらの記事にまとめました。

sdcli/のファイルについて

txt2img.py

txt2imgの推論を行うためのエントリーポイントです。Modal上に構築したアプリケーション（というかクラスそのもの）を実行する入り口部分のファイルです。CLIからプロンプトやseedなどの引数を受け取りModal上に展開したクラスへ渡す、ハンドラー的な役割を果たします。

util.py

生成した画像をローカルに保存したり、seedを指定しない場合に自動で生成したり、ちょっとした処理を実装しています。

その他のファイルについて

.env

hugging faceのAPIトークンを設定するファイルです。今はトークンのみですが、今後秘匿したい情報を追加する場合はこのファイルに追加していきます。

1HUGGINGFACE_TOKEN=""

Makefile

タスクランナーです。本当はビルドに使うべき役割のファイルですが、コマンドライン引数が多いのでMakefile で実行するようにしました。プロンプトや画像のサイズなどをここに書いて置きます。

今のところ下記のコマンドを用意しています。

  • deploy: make deployでModal上にアプリケーションをデプロイする。
  • run: make runで画像生成を実行する。

--upscalerという引数で、高解像度化するためのアップスケーラを指定するのですが、下記のReal-ESRGANのみ対応しています。今後diffusersのアップスケーラも使えるようにする予定です。

  • RealESRGAN_x4plus
  • RealESRNet_x4plus
  • RealESRGAN_x4plus_anime_6B
  • RealESRGAN_x2plus
 1deploy:
 2	cd ./setup_files && modal deploy main.py
 3
 4run:
 5	cd ./sdcli && modal run txt2img.py \
 6	--prompt "a photograph of an astronaut riding a horse" \
 7	--n-prompt "" \
 8	--height 512 \
 9	--width 512 \
10	--samples 1 \
11	--steps 50 \
12	--upscaler "" \
13	--use-face-enhancer "False" \
14	--fix-by-controlnet-tile "False"

まとめ

DiffusersとModalを使い、CLIから画像生成可能なアプリケーションを作ることができました。このアプリケーションはControlNet 1.1 TileとUpscalerを利用した高解像度化に対応しており、モデルの取り替えを自由に行えます（再ビルドは必要ですが、マルチステージビルドを採用して効率よく取り替えできるよう工夫しました）。 そのうち img2imgを実行できるよう、少しずつ改修していきます。

living room

living room

living room