Stable Diffusion を NVIDIA Container Toolkit で動かす
目次
経緯
最近よく Twitter で Midjourney や Stable Diffusion といった画像生成AIを見かけるようになり、ふーんと思っていたのですが、 Stable Diffusion がオープンソース化されたということで、せっかくGPUも積んでいることだしローカルで構築してみようと思いました。
また、調べてみるとホストマシンに Python をインストールしたり、Google Colab を使って構築している人はいました。しかし、宗教上の理由であまりホストマシンに Python 開発環境をそのまま入れたくなかったため、前職でも使ったことのある NVIDIA Container Toolkit を使ってみようと思いました。
環境
| 項目 | 値 |
|---|---|
| OS | Windows 10 Pro(WSL2) |
| CPU | Intel Core i9-9900K |
| RAM | DDR4-3600 64GB |
| GPU | Geforce RTX 3080 Ti 12GB |
| Docker | 20.10.11 |
| NVIDIA Driver | 512.15 |
| CUDA | 11.6 |
| ベースイメージ | nvcr.io/nvidia/pytorch:22.07-py3(※) |
※このイメージは CUDA11.7 なので、互換モード(？)で動くみたいです。 本当は自分の環境だと PyTorch:22.02 を使用したほうがよさそうです。
環境構築
NVIDIA Container Toolkit のインストール
// TODO: 別な記事にまとめるかも……
PyTorch コンテナの立ち上げ
下記コマンドでコンテナに GPU を割り当てて起動します。
$ docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:22.07-py3
必要パッケージのインストール
下記コマンドでインストールします。
pip install diffusers==0.2.4 transformers scipy ftfy
Hugging Face の登録・アクセストークン生成
学習済みモデルが公開されている Hugging Face のアカウントを作成し、アクセストークンを生成します。
Hugging Face にアクセスし、右上の Sign Up からアカウントを作成し、メールアドレス認証を済ませます。
次にアクセストークンを生成します。右上のアカウントアイコンをクリックし、 Settings を開きます。
左メニューの Access Tokens を開き、 New Token をクリックし、トークンを生成します。トークンの種類は Read で大丈夫です。
Stable Diffusion の学習済みモデルへアクセス
Stable Diffusion のページへアクセスし、ライセンスをよく読んで Access Repository をクリックします。
コーディング
ここからは公式の記事をなぞっていく感じになります。
from diffusers import StableDiffusionPipeline
YOUR_TOKEN={先ほど生成したアクセストークン}
# 学習済みモデルのロード
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=YOUR_TOKEN)
pipe.to("cuda")
# 生成に使用する文
prompt = "a photograph of an astronaut riding a horse"
# 生成
image = pipe(prompt)["sample"][0]
# 保存
image.save(f"hoge.png")
学習済みモデルをロードする際に、VRAM が 10GB 未満の場合は、下記のように半精度浮動小数を指定したほうが良いらしいです。
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, use_auth_token=YOUR_TOKEN)
躓いたところ
恐らくどっちも自分の環境の問題だと思います。
Docker が WSL2 を巻き込んで落ちた
- モデルのロード時や
docker commitを実行した際に、Docker と WSL2 が落ちました。 - 原因は恐らくディスクとメモリがいっぱいだったからっぽいです……たぶん……。
- 複数ブラウザ併用してタブを数百個とか開いてる人が悪い。
ホストのグラフィックドライバが落ちた
- 画像生成した後放置して他の作業をしていたら突然グラフィックドライバが落ちました。
- 原因は VRAM が足りていなかったみたいです。
- これも恐らくブラウザとかでVRAMが使用されていたみたいです。
参考
- Stable Diffusion with 🧨 Diffusers huggingface.co
