Deep Learning

Trinart Stable Diffusionの実装

 

Stable Diffusionにはたくさんの派生モデルがあります。

今回はその中から、「Stable Diffusion」を二次元イラストのデータセットでファインチューニングした「Trinart Stable Diffusion」を実装して試します。

Google Colabで実装

はじめにランタイムの変更からGPUを有効化しておきます。

まずGPUが有効化されているか確認します。

!nvidia-smi

生成画像をgoogleドライブに保存するためにマウントします。
マウントの方法はGUIからでも次のコマンドからでもどちらからでも平気です。

from google.colab import drive
drive.mount('/content/drive')

任意のフォルダを作成して、そこに移動してgitのクローン先を作成したフォルダにします。

%mkdir drive/MyDrive/stable_diffusion/
%cd drive/MyDrive/stable_diffusion/

必要なライブラリをインストールする。

!pip install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
!pip install pytorch_lightning tensorboard==2.8 omegaconf einops taming-transformers==0.0.1 clip transformers kornia test-tube
!pip install diffusers invisible-watermark

Stable Diffusionのクローンとインストールに合わせて、生成画像の保存先を作成

!git clone https://github.com/CompVis/stable-diffusion.git
%cd stable-diffusion
!pip install -e .
%mkdir outputs

ステップが60000/95000/115000のチェックポイントを引っ張ってくる。

容量が大きいので、git lfsを使用する。

!git lfs install
!git clone https://huggingface.co/naclbit/trinart_stable_diffusion_v2

最後に画像生成に必要なオプション（重要なのは「–prompt」）を渡してあげれば完成。

n_samplesがバッチサイズでこれは１が安定する。大きくしすぎるとメモリが足らずエラーになることがある。

n_iterが生成する画像数なので好きな値を生成してください。

今回は「ピアノを弾く猫」を生成するようにしています。今回のモデルは英語のみなので日本語はNGです。

同じ画像を出力したい場合には、オプションにseedを設定してあげれば、テキストが同じ場合は同じ画像が毎回出力されます。

別のスタイルを試してみたい方は、trinart2_step115000.ckptをtrinart2_step60000.ckptに変更してみてください。

!python scripts/txt2img.py \
    --plms \
    --ckpt ./trinart_stable_diffusion_v2/trinart2_step115000.ckpt \
    --skip_grid \
    --n_samples 1  \
    --n_iter 10  \
    --outdir outputs \
    --ddim_steps 100 \
    --prompt "A cat playing the piano"

処理が完了すると、生成画像が「/content/drive/MyDrive/stable_diffusion/stable-diffusion/outputs/samples/」に出力され、表示してみると次のような画像が保存されているのが確認できます。

生成した画像の一覧表示は以下のコードで可能です。（フォルダ内の全ての画像を表示するので注意）

import cv2
from google.colab.patches import cv2_imshow

import glob

img_list = glob.glob("outputs/samples/*")

for i in img_list:
  img = cv2.imread(i)
  cv2_imshow(img)

以上がTrinart Stable Diffusionをcolab上で動かす流れになります。

次回はWaifu Diffusionを試します。

RELATED POST