TadaoYamaokaの開発日記

個人開発しているスマホアプリや将棋AIの開発ネタを中心に書いていきます。

Stable Diffusionにdanbooruデータセットを追加学習する

Stable Diffusionに独自データセットで追加学習できるようになるために、まずは既存のデータセットを使用した学習方法を試した。

追加学習の方法には、画像3~5枚を用いてスタイルを学習させるTextual Inversionという方法があるが、ここでは追加学習にデータセットを用いて画像をテキスト条件付きで追加学習する方法を試す。

GitHubStable Diffusionには追加学習の方法についてドキュメントが用意されていないため、Waifu Diffusionの方法を参考にした。

Waifu Diffusionは、Stable Diffusionをdanbooruデータセットでファイチューニングしたものである。
Waifu Diffusionがどのように追加学習を行ったか手順は書かれていないが、ファイチューニング用のconfigファイルが用意されているため、それを使って、トライ&エラーでなんとか追加学習できるようになった。

以下、追加学習の手順について記述する。

環境準備

NVIDIAのPyTorchのdockerイメージを使用する。

docker run --gpus all --shm-size=32g --network host -v /work:/work -w /work -it nvcr.io/nvidia/pytorch:21.10-py3
ライブラリインストール

requirements.txtに書かれているライブラリをインストールする。

pip install albumentations==0.4.3 opencv-python==4.5.5.64 opencv-python-headless==4.5.5.64 pudb==2019.2 imageio==2.9.0 imageio-ffmpeg==0.4.2 pytorch-lightning==1.4.2 omegaconf==2.1.1 test-tube streamlit einops==0.3.0 torch-fidelity==0.3.0 transformers==4.19.2 torchmetrics==0.6.0 kornia==0.6 gradio Pillow==9.2.0

opencv-pythonとPillowは、使用しているdockerイメージで動作するバージョンに変更している。

以下のライブラリはGitHubリポジトリからインストールする。

pip install git+https://github.com/openai/CLIP.git@main#egg=clip
pip install git+https://github.com/hlky/k-diffusion-sd#egg=k_diffusion

git clone https://github.com/CompVis/taming-transformers.git
cd taming-transformers
pip install -e .

Waifu Diffusionのリポジトリclone

cd /work
git clone https://github.com/harubaru/waifu-diffusion.git

danbooruデータセットダウンロード

あらかじめdanbooruで、アカウント作成を行い、APIキーを作成しておく。

スクレイピング
cd waifu-diffusion
python danbooru_data/scrape.py -user username -key apikey -t "solo -rating:general order:score age:<1month"

usernameとapikeyの部分は書き換える。

-tオプションで指定するタグの条件は、デフォルトのままだと複数タグ指定によるエラーになったため、soloのみにしている。

カレントディレクトリにlinks.jsonができる。

画像とキャプションのダウンロード
python danbooru_data/download.py -f links.json -o danbooru-aesthetic/

danbooru-aestheticディレクトリ配下のimgとtextディレクトリにそれぞれ画像とキャプションがダウンロードされる。
※画像は、職場での閲覧に不適切な画像を含むので注意

画像変換

画像サイズを512×512に変換する。

danbooru_data/local/convert.pyを以下のように書き換える。

from PIL import Image, ImageOps

import os
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('in_dir')
parser.add_argument('out_dir')
args = parser.parse_args()

directory = args.in_dir

for filename in os.listdir(directory):
    var1 = os.path.join(directory, filename)
    if os.path.isfile(var1):
        print(var1)
        try:
            im = Image.open(var1)
            im = ImageOps.pad(im, (512, 512), color='black')
            im.save(os.path.join(args.out_dir, filename))
        except:
            print('skip')

元のディレクトリ名を変更してから、変換を行う。

cd danbooru-aesthetic
mv img img_org
python ../danbooru_data/local/convert.py img_org img
textディレクトリの名前変更

キャプションのディレクトリ名は、txtにする必要があるので変更する。

mv text txt

Stable Diffusionのモデル準備

apt update
apt install git-lfs
git lfs install
cd /work
git clone https://huggingface.co/CompVis/stable-diffusion-v-1-4-original

cloneにはHuggingfaceのアカウントが必要になる。

チェックポイントに空のoptimizer_statesを追加

Stable Diffusionのチェックポイントファイルには、追加学習時のresumeに必要となるoptimizer_statesが含まれていないため、以下のようなスクリプトで空のoptimizer_statesを追加して保存する。

import torch
ckpt = torch.load('/work/stable-diffusion-v-1-4-original/sd-v1-4-full-ema.ckpt')
ckpt['optimizer_states'] = []
torch.save(ckpt, '/work/waifu-diffusion/logs/original/model.ckpt')

スクリプトの実行前に、waifu-diffusionのディレクトリいlogs/originalディレクトリを作成しておく。

DejaVuSans.ttfのダウンロード

学習時にDejaVuSans.ttfがないとエラーになるため、ダウンロードする。

cd /work
wget http://sourceforge.net/projects/dejavu/files/dejavu/2.37/dejavu-fonts-ttf-2.37.zip
unzip dejavu-fonts-ttf-2.37.zip
cp dejavu-fonts-ttf-2.37/ttf/DejaVuSans.ttf waifu-diffusion/data/

追加学習

準備が整ったので、追加学習を実行する。
GPUはA100×4を使用した。V100だとCUDA out of memoryで学習できなかった。

cd waifu-diffusion
python main.py --resume logs/original/model.ckpt --base ./configs/stable-diffusion/v1-finetune-4gpu.yaml -t --no-test --seed 25 --scale_lr False --gpus 0,1,2,3

なお、使用しているdockerイメージだとPillowのワーニングが大量に表示されるため、抑止したい場合は、main.pyの先頭に以下のコードを追加する。

import warnings
warnings.simplefilter('ignore')

学習結果

logs/images/trainに学習に生成された画像が出力される。
元のStable Diffusionが5エポック学習されているため、6エポックから開始して13エポック時点で生成された画像は以下の通り。

プロンプトは訓練データに含まれるものである。

※職場での閲覧に適さない画像が生成される場合があるので要注意

まとめ

Waifu Diffusionのリポジトリを使用して、Stable Diffusionにdanbooruデータセットを追加学習することができた。
画像とキャプションを1行記載したテキストファイルを用意できれば、独自のデータセットでも同じ手順で追加学習することができる。

別途、独自データセットでの学習を試したい。