Stable Diffusionを試す - TadaoYamaokaの開発日記

拡散モデルを使用した画像生成AIであるStable Diffusionがオープンソースで公開された。
ソースだけでなくモデルの重みも自由に利用できる。商用利用も可能だ。

Stable Diffusion Public Release — Stability.Ai

早速、ローカルのPCで試してみた。
丁寧な使ったみた記事は他にあるのでほぼ個人メモである。

手順は、公式ブログのリンクにあるノートブックがすべてなので、この記事を見る必要はない。

WSL2上にconda環境作成

WSL2のUbuntu 20.04上に環境構築する。

wget https://repo.anaconda.com/archive/Anaconda3-2022.05-Linux-x86_64.sh
sh Anaconda3-2022.05-Linux-x86_64.sh

conda環境作成

conda環境を作成する。

conda create -n stable-diffusion python=3.9
conda activate stable-diffusion

PyTorchインストール

PyTorchをインストールする。

conda install pytorch torchvision torchaudio cudatoolkit=11.6 -c pytorch -c conda-forge

Jupyterをインストールする(Stable Diffutionには不要だが、ついでにpandasとmatplotlibも)。バージョン不整合を避けるためにconda-forgeからインストールする。

conda install jupyter pandas matplotlib -c conda-forge

Hugging Faceセットアップ

公式ブログにあるHugging Faceのリンクから、stable-diffusion-v1-4をクリックし、Stable Diffusion v1-4 Model Cardのページの説明の通りセットアップする。
CompVis/stable-diffusion · Hugging Face

Hugging Faceアカウントが必要になるので、アカウントを作成する。

ライブラリインストール

必要なライブラリをインストールする。モデルロード時に警告がでるのでftfyもインストールする。

pip install diffusers transformers scipy ftfy

ログインする。

huggingface-cli login

トークンの入力が必要になるので、説明文に従いURLにアクセスしてトークンを作成する。

git config --global credential.helper store

を実行するように提案されるので実行しておく。

多重ログインはできない仕様のようなので、一旦ログオフする。

huggingface-cli logout

ノートブックから実行

Jupyterノートブックから実行する。
VS CodeでRemote-WSLに接続して、適当な名前でノートブック(work.ipynbなど)を作成する。
VS Codeで開く。
カーネルにcondaのstable-diffusion環境を選択する。

モデルロード

モデルをダウンロードしてロードする。
use_auth_tokenには、上記で作成したトークン文字列を指定する。
ソース中にトークンを記述したくない場合は、

from huggingface_hub import notebook_login
notebook_login()

で入力する方法もある。

import torch
from torch import autocast
from diffusers import StableDiffusionPipeline

model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda"

pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=r'xxxxxxxx')
pipe = pipe.to(device)

初回ダウンロードには12分くらいかかった。

2回目からはキャッシュされるため、10秒くらいでロードできる。

GPUメモリが足りない場合は、float16モデルも用意されている。
詳細はノートブックを参照。

画像生成

「a photograph of an astronaut riding a horse(馬に乗った宇宙飛行士の写真)」というプロンプト(サンプルのまま)で画像を生成する。

from torch import autocast

prompt = "a photograph of an astronaut riding a horse"
with autocast("cuda"):
  image = pipe(prompt)["sample"][0]  # image here is in [PIL format](https://pillow.readthedocs.io/en/stable/)

# Now to display an image you can do either save it such as:
image.save(f"astronaut_rides_horse.png")

# or if you're in a google colab you can directly display it with 
image

ノートブックには複数枚生成してタイル表示する方法や、画像解像度を変える方法も記載されている。

まとめ

Stable DiffusionをWindowsのWSL2上でJupyterノートブックで動かす手順について記載した。
Stable Diffusionは重みとソースが公開されているため、ファインチューニングや、日本語化対応とかも工夫すればできそうである。
しばらくは、いろんなプロンプトで画像生成して遊んでみたい。