Stable Diffusionに独自データセットで追加学習できるようになるために、まずは既存のデータセットを使用した学習方法を試した。
追加学習の方法には、画像3～5枚を用いてスタイルを学習させるTextual Inversionという方法があるが、ここでは追加学習にデータセットを用いて画像をテキスト条件付きで追加学習する方法を試す。
GitHubのStable 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のみにしている。
画像とキャプションのダウンロード
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
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エポック時点で生成された画像は以下の通り。
プロンプトは訓練データに含まれるものである。
※職場での閲覧に適さない画像が生成される場合があるので要注意
