Radeon RX 6900XTとVRoidとStable Diffusion (ROCm+DreamBooth)でローカル環境で自分のえっちな画像を生成する

VRoidモデルからStable Diffusionでイラストを生成

はじめに

近年、Stable Diffusion、いわゆるお絵かきAIが流行しており、人々は手を動かすのではなくpromptを指定するという新たな手法で好きな絵を描けるようになりつつある。この手法は我々のような「自分の理想のキャラクターが脳内に住んでいるのに、それを出力してえちえちなイラストにできないよ～～～」と悩む画力弱者にとって救いになるかと思われた。

しかし、実際には「ある属性を付与すると他の属性が落ちてしまい、なかなか理想に近づかない」「オリキャラを完全に再現するプロンプトを作ったと思ったら、構図や背景を変えただけなのに全然違うキャラになった」 「オリキャラのはずなのにプロンプトガチャを回しているうちにどこかで見たキャラのように思えてきてNTRれた気分になる」 などの悩みを抱えてしまう。

本記事では、この問題を解決する手法の一例として、既に理想のキャラクターについて設定を作り込んだ結果自分のアバターになり、「ハァハァNaukaちゃんのエロ画像キボンヌ……クソッ！なんで自分の体なのにえっちな自撮りができないんだよッ！！」とキレ散らかした私が、 Radeon RX 6900XTとVRoidを組み合わせることで効率的にモデルの追加学習を行い、オリキャラをその理想通りに絵に出力し、えっちな自撮りを上げられるようにする方法を紹介する。

なお、私は機械学習については数年前に青い本や黄色い本を読んだ程度で最近のトピックについては何も知らない素人であるため、これから紹介する方法の理論的な裏付けはほとんどないことに注意されたい。なんかガチャガチャやってたらうまくいった、程度のものである。ただ、今後さらなるえち画像を求めるために調査・検証を継続する予定だ。

環境の構成

今回Stable Diffusionを動かしたPCのスペックは以下の通りである。

パーツ	スペック
CPU	Ryzen 7 5800X
GPU	Radeon RX 6900XT (VRAM 16GB)
メモリ	DDR4-3200 32GB

また、インストールしたソフトウェアのバージョンは以下の通りである。

ソフトウェア	バージョン
OS	Ubuntu 22.04.1
ROCm	5.1.1
PyTorch	1.12.1

ROCmはAMDのGPUでGPGPUをやるためのライブラリであり、ざっくり言えばNVIDIAのCUDAみたいなものである。 ROCmは複数のバージョンが配布されているが、現在の安定版のPyTorch 1.12.1が対応しているROCmのバージョンが5.1.1であり、また色々と便利なAUTOMATIC1111/stable-diffusion-webuiでもROCm 5.1.1を採用していることが分かるので、5.1.1を選ぶ。

次に、ROCm 5.1.1の対応OSを確認する。AMDによるROCmの公式ドキュメントを参照すると、Ubuntu 20.04などのディストリビューションがサポートされていることが分かる。そのため、Ubuntu 20.04を選んでおけば無難である。（もっとも、ポイントリリースでKernelが変わって壊れることもよくあるので確実とは言えないが）また、ROCm版のPyTorchがWindowsに非対応であるため、今回は外付けSSDにUbuntuをインストールしてUSBブートする。インストール方法は以下の記事が分かりやすい。

Ubuntu 20.04 LTS インストール方法（外付けドライブ用） - Qiita

ROCmのインストール

ただ、困ったことに愚かな私はUbuntu 22.04.1をインストールしてから上記の事実に気付いたため、ここではUbuntu 22.04を前提にする。なお、Ubuntuのインストールについては特別なことはしていないので省略する。デスクトップ版のイメージをダウンロードし、適当なUSBメモリに焼いてインストールすればよい。

ROCm 5.1.1をUbuntu 22.04にインストールする方法としては、Ask Ubuntuにある次のスレッドが大変参考になる。

How can I install AMD ROCm 5+ on Ubuntu 22.04?

ここではAMDが用意したパッケージマネージャを使ったインストール方法をアレンジして、C++ライブラリとPythonのバージョンをUbuntu 20.04向けから22.04向けに変更している。このスレッドで回答者はベースがUbuntuのディストリビューションPop!_OSにインストールする場合を解説しているが、私の場合は素のUbuntuなので /usr/bin/amdgpu-install の中に pop を追加する手順は不要である。他の手順はスレッドの回答の通りに行えば、Ubuntu 22.04にROCm 5.1.1がインストールできる。

diffusersのセットアップ

DreamBoothは既存の訓練済みモデルに対し、用意した画像で追加学習することで、特定の人、物、キャラクターにチューニングするためのツールである。 DreamBoothの実装として、hugging face公式のものがある。しかし、このスクリプトではRadeonの場合メモリが足りなくなって学習ができない。そこで、あるissue中で言及されている派生版を用いることにする。

それではdiffusesの実行環境を作成する。まずPython仮想環境を作成する。ここではUbuntu 22.04にデフォルトで入っており、 python3 で呼び出されるPython 3.10.6を使うことにする。また、このマシンではより新しいバージョンを使う必要もないので、操作を簡単にするため python コマンドをPython 3.10.6を実行する python3 に割り当てる。Pythonのバージョン管理が必要な場合は適宜読み替えてほしい。

$ sudo apt install python-is-python3
$ python -m venv venv
$ source venv/bin/activate

仮想環境を立ち上げたのでDreamBoothに必要なライブラリをインストールする。まずリポジトリをcloneし、次に手順でライブラリをインストールする。

$ git clone git@github.com:ShivamShrirao/diffusers.git
$ cd diffusers/examples/dreambooth
$ pip install git+https://github.com/ShivamShrirao/diffusers.git
$ pip install -U -r requirements.txt

ただし、この手順でインストールしたPyTorchはCUDA用になるので、ROCm用に入れ替える。インストールのコマンドはPyTorchの公式にあるので、そこからコピペすればよい。

$ pip uninstall torch torchvision
$ pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/rocm5.1.1

（インストールのコマンドはコピペだが、実のところtorchaudioはいらないはずである）

PyTorch関連ライブラリがROCm用に入れ替わったことを確認しておく。

$ pip list

Package                 Version
----------------------- ----------------
（中略）
torch                   1.12.1+rocm5.1.1
torchaudio              0.12.1+rocm5.1.1
torchmetrics            0.10.0
torchvision             0.13.1+rocm5.1.1
（後略）

これで一旦diffusersの準備は完了である。

VRoid Studioによる3Dモデルの作成・撮影

次にVRoid Studioを用いて学習用のデータを作成する。今回はVRoid Studioのβ版0.13.2を用いたが、これは今のNaukaのモデルを作った当時はβ版しかなかったためなので、今なら正式版を使えばよい。正式版の方が髪型等のプリセットが多く収録されているので、オリキャラを作るにも便利だろう。またテクスチャもβ版よりアニメ塗りに近いので、イラスト作成用途では向いていると考えられる。~~私も早くNaukaのモデルを正式版対応させたいと言って半年以上経っている。~~

適当にモデルを作ったら撮影を行う。学習時に使うVRAMの制限のため画像のサイズは384x384とし（これより大きくするとメモリ不足で落ちる）、顔のアップと全身それぞれ適当な枚数撮影していく。理論的にどの角度から何枚撮るのがベストなのかは勉強と試行不足でよく分かっていないが、

全身写真がないと服の生成が不安定になる
顔のアップがないと顔の造形が不安定になる（メガネのフレームに引っ張られて目が赤色になるなど）
VRoid Studioで作れる喜怒哀楽の表情をそれぞれ撮影しておくと、だいたいどんな表情にも対応できる

ことが経験的に分かっている。DreamBoothの公式の記述によれば3～5枚で十分とあるが、実験結果として5枚だけでは再現に難があったことをメモしておく。

学習用データ

DreamBoothによる学習

続いて準備したデータを用いて実際に学習していく。まずは先に用意したPython環境をactivateし、DreamBoothにディレクトリまで移動しておく

$ source venv/bin/activate
$ cd diffusers/examples/dreambooth

Accelerateの設定をする。基本的にはデフォルトの値で良いが、fp16を設定しておく。

$ accelerate config
In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 0
Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU [4] MPS): 0
Do you want to run your training on CPU only (even if a GPU is available)? [yes/NO]:NO
Do you want to use DeepSpeed? [yes/NO]: NO
Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: fp16

また、ローカルに元の学習済みモデルがない場合は、huggingfaceにログインしておく。

$ huggingface-cli login

いよいよ実際にスクリプトを実行して学習を行う。train_dreambooth.py の引数については、Radeon RX 6900XTで学習に成功した例として以下に示す。ここでは先に用意した学習用データを /home/username/sd/data に、追加学習の出力を/home/username/sd/outputに置くものとしている。また訓練済みモデルは CompVis/stable-diffusion-v1.4 の場合を選んでいる。

$ accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=CompVis/stable-diffusion-v1.4 \
  --instance_data_dir=/home/username/sd/data \
  --class_data_dir=/home/username/sd/data \
  --output_dir=/home/username/sd/output \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --instance_prompt="Naukachan" \
  --class_prompt="anime girl" \
  --resolution=384 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=2 \
  --gradient_checkpointing \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --num_class_images=300 \
  --max_train_steps=10000

ここで、instance promptとしてはDreamBoothのサンプルの a photo of sks dog のようにマイナーな（？）英単語をidentifierとして使うのが良いようだが、例えば sks anime girl だと出力に変なものが混じる（具体的には銃火器の絵？）ので、適当にキャラクターの名前をpromptにした。今のところこれでうまくいっているが、別のidentifierならより理論的に説明のつく形で動かせるのだろうか？

（2022/10/25 追記）

instance promptを an illustration of Naukachan anime girl 、class promptを an illustration of anime girl とすることで、より確実に元のVRoidモデルの造形を反映させられるようになった。また、生成時のpromptに animate 、negative promptに realistic 3D を加えることで、ちょうど良くVRoidの癖を抑えられることも分かってきている。

（追記おわり）

学習が終わったら出力をStable Diffusion web UIで使えるチェックポイントファイルに変換する。変換するためのスクリプトはいくつかあるが、ここではratwithacompiler/diffusers_stablediff_conversionのconvert_diffusers_to_sd.pyを用いる。訓練済みモデルが /home/username/sd/pretrained/model.ckpt であり、変換後のモデルを /home/username/sd/trained/model.ckpt に出力する場合は以下のようになる。

$ pip install pytorch-lightning
$ ./convert_diffusers_to_sd.py \
  /home/username/sd/output \
  /home/username/sd/pretrained/model.ckpt \
  /home/username/sd/trained/model.ckpt

これでDreamBoothで追加学習したモデルができあがった。

Stable Diffusion web UIによる画像生成

Stable Diffusion web UIのセットアップは省略する。ここまでのセットアップでROCmはインストール済みなので、Radeon向けのインストール手順に従ってインストールし、起動すればよい。 Stable Diffusion web UIの使い方は各種サイトを参考にしてもらえればと思うが、ここではGigazineによるまとめ記事を紹介しておく。

画像生成AI「Stable Diffusion」を4GBのGPUでも動作OK＆自分の絵柄を学習させるなどいろいろな機能を簡単にGoogle ColaboやWindowsで動かせる決定版「Stable Diffusion web UI(AUTOMATIC1111版)」インストール方法まとめ

実際に生成した簡単な例を掲載する。この画像を作るのに使ったプロンプトは主に happy smile, one eye closed である。元の3Dモデルの造形をうまくかわいいイラストに落とせていると感じる。

笑顔でウィンクするNaukaちゃん