Stable diffusionの追加学習によって、お気に入りの絵師にお気に入りのキャラを書かせる。

2022年10月3日 22:28

　タイトル通り、追加学習によりAIに、お気に入りの絵師の画風でお気に入りのキャラを書かせることを目標とします。元々はmimicの炎上を見て、技術的に難しいことはしておらず、簡単に再現できてしまうのではないかと思っていたので、それを証明するためにやっていました。倫理的にも法的にも問題ありそうなので、このような投稿をするかは迷っていたのですが、NovelAIがとんでもないサービスを公開して弾除けになってくれそうなので、投稿します。

参考サイト

Waifu diffusionのREADME

「dlshogi」開発者としても有名な山岡さんのブログ

環境を構築

　追加学習には30GB程度のVRAMをもつGPUが必要です。私はGCPのスポットでA100を１台借りました。（１時間1.1ドル）。
追記：VRAM13.6GBでできました

　Waifu diffusionの追加学習用dockerイメージがあるので入ります。（NUM_GPUはGPU数を指定）

docker run --gpus all -it -e NUM_GPU=1 ghcr.io/derfred/waifu-diffusion

GCPの場合？ --gpus allをつけないとnvidiaのドライバーを読み込んでくれませんした。

学習データの用意

　~~次にDanbooruの画像をスクレイピングします。やり方は山岡さんのブログそのままです。~~

waifu/danbooru-aesthetic
├── img
│ └── hogepiyo.jpg 512×512のjpg,png,bmp
└── txt
└── hogepiyo.txt 画像の説明文

モデルのダウンロード

　追加学習するモデルを選びます。私はWaifu diffusion v1.3のepoch 5を選びました。fullがつくモデルを選びましょう。

Waifu diffusion v1.3のcheckpoint

コンフィグファイルの書き換え

　デフォルトの設定では、学習データのかさ増しのために、画像を左右反転してしまいます。左右対称のキャラクターを学ばせるには良いのですが、非対称のキャラクターを学ばせたい場合、変更しましょう。（私が発見したわけではなくツイッターで知りました。）
　変更するファイルはconfigs/stable-diffusion/v1-finetune-4gpu.yamlです。

    train:
      target: ldm.data.local.LocalBase
      params:
        size: 512
        mode: "train"
        flip_p: 0.0 #これを追加！

学習開始！

　私が何か忘れていなければ、これで学習できるはずです。

sh train.sh -t -n "aesthetic" --resume_from_checkpoint <モデルのパス> --base ./configs/stable-diffusion/v1-finetune-4gpu.yaml --no-test --seed 25 --scale_lr False --data_root "./danbooru-aesthetic"

　学習中は損失ベスト3までのcheckpointと最終的なcheckpointを保存します。1ファイル14GB越えですので、容量は多めにしておきましょう。放っておくといつ止まるのかよくわかりません。ctrl+Cで現時点でのcheckpointを保存してプログラムが止まります。

必要ないデータの除去

　学習済みモデルには主に以下のデータが入っています。

最終的な重み
重みの指数移動平均(ema)
- 推論用の重み（過学習を防ぐため、推論では過去の重みとの指数移動平均をとったものを使う）
optimizer_states：
- Adam用の勾配指数移動平均と勾配二乗指数移動平均

　推論は１，２のどちらかがあればできるので、他は除去します。そのためのスクリプトがscripts/prune.pyです。ただしdockerイメージに入っているものだと、対象のモデルが指定できず使いづらいので、githubにある最新のものに置き換えたほうが便利です。除去した後は4GBほどになります。

学習結果(10エポック)

　完成したckptファイルを使う方法は既に紹介記事が大量にあると思いますので省略します。
　diffusersで使いたい場合は、scripts/convert-diffusers.pyで変換できます。ただしbertをtext_encoderに名前変更、vqvaeをvaeに名前変更、feature_extractorとsafety_cheker、model_index.jsonをWaifu diffusionのものからコピーしてやっと動いた気がします。
　でき上がった画像です。promptはタグ1個のみで、厳選はしていません。私のお気に入りのキャラがバレたら私の目標は達成です。