Stable Diffusion + LoRA on WSL2 で自分の画風を学習させてイラスト生成する

# はじめに
最近、自作ゲーミングPCを新調しました。
イラスト生成AIをローカルで稼働させられるスペックです。せっかくなので試してみたくなりました。
調べていると、LoRAというAIモデルに対する追加学習の手法があるらしい。
自分は年賀状などで絵をたまに描くので、**自分の画風を学習させればお絵描きも自動化できるのでは？** と思いました。（なまけもの）

Windows環境を汚したくなかったので、環境構築はWSL上で行いました。
Linux環境での環境構築は地味に情報が少なく、結構な失敗があり、試行錯誤しました。

情報共有のため、本記事に導入過程をまとめておきます。

# 内容
## 前提
### PCスペック

まず自分のPCスペックを示します。
流用パーツもありますが、全部含めて13万円くらいでしょうか。

|   |スペック|だいたいの価格|
|---|---|--:|
|CPU|Intel Core i5 12400F|21,800|
|RAM|DDR5 32GB|7,980|
|GPU|nVIDIA RTX3060Ti GDDR6X VRAM8GB|49,800|
|Strage|M.2 NVMe 1TB|5,980|
|Power|850W GOLD認証|10,400|

[wiki](https://wikiwiki.jp/sd_toshiaki/pc_spec#ta9a079a)から抜粋すると、AIイラストにおける要求スペックは以下のようです。
ストレージが足りてないですが、本記事でやりたいことをやる分には問題ありませんでした。

* CPU: 影響は小さい
* RAM: 16GB以上
* GPU: nVIDIAのRTX系GPU。VRAM最低8GB以上、推奨12GB以上
* Strage: M.2 NVMe 2TB以上
* 電源容量十分かつGOLD認証

本PCはVRAM 8GBであり、LoRAが実行可能なスペックです。[参考](https://cdn.wikiwiki.jp/to/w/sd_toshiaki/%E3%83%AD%E3%83%BC%E3%82%AB%E3%83%AB%E7%89%88%E5%B0%8E%E5%85%A5/::attach/temp4_VRAM.png)

### 自分の絵

こんな絵を描きます。
<blockquote class="twitter-tweet"><p lang="ja" dir="ltr">あけおめ！ <a href="https://t.co/QvLOsx9Fi4">https://t.co/QvLOsx9Fi4</a></p>&mdash; すいばり@&#39;23年5戦3勝1敗1分⚾️ (@Suibari_cha) <a href="https://twitter.com/Suibari_cha/status/1609216444818149379?ref_src=twsrc%5Etfw">December 31, 2022</a></blockquote> <script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script>

以下が自分のpixivです。
https://www.pixiv.net/users/251033

特徴としては、画風が安定してない、アニメ塗りが好きというとこでしょうか。
自分のpixivから10枚の絵をチョイスしてAIに食わせます。

## 導入過程

目次です。
なお本記事におけるLoRAの学習方式は「DreamBooth+キャプション方式（正則化画像不使用）」です。

1. WSL2セットアップ
1. Stable Diffusionに必要なもののインストール
1. AUTOMATIC1111版Stable Diffusion web UIのダウンロード
1. モデルのダウンロード、配置
1. web UI起動スクリプトの編集
1. web UI起動、お試しイラスト生成
1. 追加学習用の画像収集、配置
1. WD-tagger1.5のダウンロード、実行
1. キャプション編集
1. sd-scriptsのダウンロード
1. LoRA生成
1. web UIでの生成モデル読み出し、自分風イラスト生成

### WSL2セットアップ

以下の記事などを参考にしてください。
systemdの問題でSSH接続ができなかったり色々と不便なので、Store版がおすすめです。

https://qiita.com/hiromasa-masuda/items/355fe0c882a1a0abe1e1

### Stable Diffusionに必要なもののインストール

```
sudo apt install wget git python3 python3-venv
```

CUDA11.7もインストールします。

```
wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin
sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda-repo-wsl-ubuntu-11-7-local_11.7.1-1_amd64.deb
sudo dpkg -i cuda-repo-wsl-ubuntu-11-7-local_11.7.1-1_amd64.deb
sudo cp /var/cuda-repo-wsl-ubuntu-11-7-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt update
sudo apt -y install cuda-11-7
```

パスを通すため、.bashrcに以下を追記しておきます。

```
# Setting CUDA PATH
export PATH="/usr/local/cuda-11.7/bin:$PATH"
export LD_LIBRARY_PATH="/usr/local/cuda-11.7/lib64:$LD_LIBRARY_PATH"
```

### AUTOMATIC1111版Stable Diffusion web UIのダウンロード

homeディレクトリにcloneします。

```
git clone --depth=1 --branch v1.2.1 https://github.com/AUTOMATIC1111/stable-diffusion-webui ~/stable-diffusion-webui
```

### モデルのダウンロード、配置

モデル、VAEをダウンロードします。
以下は例であり、モデル、VAEは自分で調べて好きなものをチョイスしてみてください。
モデルにはライセンス（商用利用禁止など）があるので事前に内容を理解しておきましょう。

```
wget https://huggingface.co/WarriorMama777/OrangeMixs/resolve/main/Models/AbyssOrangeMix3/AOM3_orangemixs.safetensors
wget https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors
```

適切な場所に格納します。

```
mv AOM3_orangemixs.safetensors ~/stable_diffusion_webui/models/Stable-diffusion
mv vae-ft-mse-840000-ema-pruned.safetensors ~/stable_diffusion_webui/models/VAE
```

### web UI起動スクリプトの編集

stable-diffusion-webui/webui-user.shを編集します。
対応する行を以下の通り書き換えてください。

```
export TORCH_COMMAND="pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117"
```

```
export COMMANDLINE_ARGS="--autolaunch --xformers"
```

スクリプトに実行権限を付与します。

```
chmod 755 webui-user.sh
bash ./webui-user.sh
```

### web UI起動、お試しイラスト生成

webui.shを実行します。

```
bash ./webui.sh
```

[http://127.0.0.1:7860](http://127.0.0.1:7860) にアクセスします。
Web UIが表示されれば成功です。
色々試してみてください。

### 追加学習用の画像収集、配置

追加学習させるイラストを10枚ほど集めてきてください。
[他記事](https://note.com/epxstudio/n/n3784d4b9ff21)より、10枚で最低限学習はできるようです。
以下の通りディレクトリを掘って配置します。
WindowsエクスプローラからWSLのディレクトリには"\\wsl.localhost\Ubuntu"でアクセスできます。

```
lora_train_data/
├── reg
└── train
    └── 10_suibari             # 10_<学習対象の名前>
        ├── 28236027_p0.jpg    # 画像10枚を配置。形式はなんでもいい
            ...
```

### WD-tagger1.5のダウンロード、実行

WebUI Extentionをcloneしてきます。
WebUIからもダウンロードできるらしいのですが、自分の環境ではTaggerタブが表示されずうまくいきませんでした。

```
cd ~/stable_diffusion_webui
git clone https://github.com/toriato/stable-diffusion-webui-wd14-tagger.git extensions/tagger
```

次にwikiの通りにタグを出力します。
画像を置いたディレクトリにタグ(*.txt)が出力されます。
https://wikiwiki.jp/sd_toshiaki/キャプション・タグ編集#f04af331

### キャプション編集

以下wikiのステップ3を参考に、「キャラの身体的特徴」と「キャラの衣服特徴」のタグを削除します。
最終的には"1girl, <キャラ名>"くらいになると思います。
https://wikiwiki.jp/sd_toshiaki/LoRA#z03bbbb9

ちなみにwikiのステップ4は不要です。昔は必要だったのかな？

### sd-scriptsのダウンロード

画像、キャプションが用意できたのでいよいよLoRA学習を行います。
まずはclone。

```
git clone https://github.com/kohya-ss/sd-scripts.git
cd sd-scripts
```

venvに入り、必要パッケージを導入。

```
source venv/bin/activate
pip install torch torchvision xformers
pip install --upgrade -r requirements.txt
```

PyTorchのconfigを設定。

```
accelerate config

# 出てくる質問は以下の通りに答える。
- This machine
- No distributed training
- NO
- NO
- NO
- all
- fp16
```

### LoRA生成

以下スクリプトを作っておく。

```config.toml
[general]
enable_bucket = true                                        # 画像サイズがバラバラの場合はtrueを指定

[[datasets]]
resolution = 512                                            # 学習解像度
batch_size = 4                                              # バッチサイズ

  [[datasets.subsets]]
  image_dir = '../lora_train_data/train/10_xxx'  # 学習用画像を入れたフォルダを指定
  caption_extension = '.txt'                                # キャプションファイルの拡張子　.txt を使う場合には書き換える
  num_repeats = 10                                          # 学習用画像の繰り返し回数
```

以下コマンドを実行。

```
accelerate launch --num_cpu_threads_per_process 1 train_network.py \
    --pretrained_model_name_or_path=../stable-diffusion-webui/models/Stable-diffusion/AOM3A1_orangemixs.safetensors \
    --dataset_config=./config.toml \
    --output_dir=../stable_diffusion_webui/models/Lora \
    --output_name=<学習したモデル出力時のファイル名> \
    --save_model_as=safetensors \
    --prior_loss_weight=1.0 \
    --max_train_steps=400 \
    --learning_rate=1e-4 \
    --optimizer_type="AdamW8bit" \
    --xformers \
    --mixed_precision="fp16" \
    --cache_latents \
    --gradient_checkpointing \
    --save_every_n_epochs=1 \
    --network_module=networks.lora
```

自分の環境では30分弱で実行完了し、../stable_diffusion_webui/models/Loraにsafetensorsファイルが生成されました。

### web UIでの生成モデル読み出し、自分風イラスト生成

以下wikiを参考にしてみてください。
https://wikiwiki.jp/sd_toshiaki/学習ファイルの使用方法#c27effe3

## 結果

こうして作った、自分風AIによる生成イラストが以下です。

<blockquote class="twitter-tweet"><p lang="ja" dir="ltr">左が生成AIモデルそのままで、右が自分の絵を学習させたもの。<br>なんだろう。自分の絵か？と言われると…ただ下手にしただけな気もするｗ<br>動きのなさ、立体感のなさ、アニメ塗りへの執着とか感じるあたりは自分っぽいのかもしれん。<a href="https://twitter.com/hashtag/%E7%94%9F%E6%88%90AI?src=hash&amp;ref_src=twsrc%5Etfw">#生成AI</a> <a href="https://twitter.com/hashtag/LoRA?src=hash&amp;ref_src=twsrc%5Etfw">#LoRA</a> <a href="https://t.co/6YJUnPFUFb">pic.twitter.com/6YJUnPFUFb</a></p>&mdash; すいばり@&#39;23年5戦3勝1敗1分⚾️ (@Suibari_cha) <a href="https://twitter.com/Suibari_cha/status/1687838497900756992?ref_src=twsrc%5Etfw">August 5, 2023</a></blockquote> <script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script>

そもそも自分の画風が安定していないということもあって、自分の絵に似ているかの評価は難しいですね。
元々上手い人が本手法をやらないと厳しいなと思いました。
もしくは本記事の手順で改善すべき点があるか…。（ぜひ指摘してください）
ただ、配色がアニメ塗りになるのは、おそらく学習結果だと思うので凄いですね。

# おわりに

環境構築は普通に大変でした。
かなりの試行錯誤が必要な点では、イラストを描くのと通ずる部分があるかもしれないです。
自分は構図を考えるのが苦手なので、AIに構図を考えてもらってそれを自分が描く、などの手法がいのかもしれません。
反目するのではなく、AIと協調していける世の中になったらいいなと思いました。