WSL2でJoePenna版のDreamboothを動かす
Windows11のローカルでDreamboothを動かす環境構築を行います。
今回紹介しているJoePenna版のDreamboothは、成功事例が多く精度が高いと評価を受けています。また、ckptファイルを出力してくれるので、変換することなくAUTOMATIC1111などで使用できます。
ただし最低VRAM24GB必要なので注意してください。
筆者のPCの環境は以下の通りです。
| 項目 | 内容 |
|---|---|
| OS | Windows11 22H2 |
| CPU | Intel Core i7-11700k |
| メモリ | DDR4-3200 64GB |
| GPU | NVIDIA RTX3090Ti |
Windows11のWSL2が起動できていることを前提に解説していきます。
また、今回WSL2にインストールしたディストリビューションはUbuntu-20.04です。
CUDAのインストール
まずWSL2を起動し、nvidia-smiコマンドが使えるか確認します。
Windowsホスト側にグラフィックドライバーが入っていれば、WSL2上でもグラボを認識します。
$ nvidia-smi
Tue Oct 18 19:00:00 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.76.02 Driver Version: 517.48 CUDA Version: 11.7 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 NVIDIA GeForce ... On | 00000000:01:00.0 On | Off |
| 0% 45C P8 17W / 450W | 1267MiB / 24564MiB | 5% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| No running processes found |
+-----------------------------------------------------------------------------+
うまく動かない場合は、Windowsのグラフィックドライバーを更新しましょう。
ここでは説明を割愛しますが、GeForce Experienceから更新するのが楽です。
次にWSL上にCUDAをインストールします。
CUDA Toolkitの公式ページを参考に、CUDAをインストールします。
$ 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.6.0/local_installers/cuda-repo-wsl-ubuntu-11-6-local_11.6.0-1_amd64.deb
$ sudo dpkg -i cuda-repo-wsl-ubuntu-11-6-local_11.6.0-1_amd64.deb
$ sudo apt-key add /var/cuda-repo-wsl-ubuntu-11-6-local/7fa2af80.pub
$ sudo apt-get update
$ sudo apt-get -y install cuda
追記：CUDA 12.xは正常に動かない可能性があるので、CUDA 11.xをダウンロードしてください。
cuDNNのインストール
cuDNNをインストールします。cuDNNのダウンロードページからWindowsにダウンロードします。
ダウンロードにはNVIDIA DEVELOPERへの登録が必要です。
Windowsの「ダウンロード」に保存した場合、以下のコマンドでインストールできます。違う場所に保存した場合は適宜書き換えてください。
$ sudo dpkg -i /mnt/c/Users/<ユーザー名>/Downloads/<cuDNNのファイル名>.deb
その後、To install the key, run this command:と表示され、その下にコマンドが表示されます。
これをコピーして実行しましょう。
cuDNNをインストールしただけでは環境変数が設定されていないので、以下のコマンドで環境変数を設定します。
$ vi ~/.bashrc
末尾に以下のコマンドを追記します。
以下のコードをコピーし、.bashrcの一番下まで移動したあと、oを押し、右クリックで貼り付けます。
あとはEscを押して:wqを押せば保存できます。
export PATH="/usr/local/cuda/bin:$PATH"
export LD_LIBRARY_PATH="/usr/local/cuda/lib64:$LD_LIBRARY_PATH"
Ubuntuを再起動し、nvccコマンドが使えるか確認します。
$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Sep_21_10:33:58_PDT_2022
Cuda compilation tools, release 11.8, V11.8.89
Build cuda_11.8.r11.8/compiler.31833905_0
細かい数字は一致していなくても大丈夫なので、このように表示されたら成功です。
Anacondaのインストール
次に、Anacondaをインストールします。
Anaconda公式サイトから、Linuxタブにある64-Bit (x86) Installerをダウンロードします。
Windowsの「ダウンロード」に保存した場合、以下のコマンドでインストールできます。違う場所に保存した場合は適宜書き換えてください。
$ bash /mnt/c/Users/<ユーザー名>/Downloads/<Anacondaのファイル名>.sh
実行すると、以下の文章が表示されます。「続行するにはENTERキーを押してください」とのことなので、Enterを押します。
In order to continue the installation process, please review the license
agreement.
Please, press ENTER to continue
>>>
Spaceを何回か押して下までスクロールします。
利用規約に同意するか聞かれます。同意する場合はyesと入力し、Enterを押します。
Do you accept the license terms? [yes|no]
[no] >>>
この場所にインストールして良いかを聞かれるので、Enterを押します。
Anaconda3 will now be installed into this location:
/home/<Ubuntuユーザー名>/anaconda3
- Press ENTER to confirm the location
- Press CTRL-C to abort the installation
- Or specify a different location below
[/home/<Ubuntuユーザー名>/anaconda3] >>>
インストールが始まります。しばらく時間がかかります。
インストールが完了すると、conda initを実行してAnacondaを初期化してよいか聞かれるので、yesと入力し、Enterを押します。
installation finished.
Do you wish the installer to initialize Anaconda3
by running conda init? [yes|no]
[no] >>>
Ubuntuを再起動し、conda環境が有効化されていることを確認します。
$ conda info -e
# conda environments:
#
base * /home/<Ubuntuユーザー名>/anaconda3
仮想環境の作成
Dreamboothを実行する仮想環境を作成します。
以下のコマンドで仮想環境が作成できます。
$ conda create -n dreambooth python=3.10
yと入力し、Enterを押します。
The following NEW packages will be INSTALLED:
_libgcc_mutex pkgs/main/linux-64::_libgcc_mutex-0.1-main
_openmp_mutex pkgs/main/linux-64::_openmp_mutex-5.1-1_gnu
ca-certificates pkgs/main/linux-64::ca-certificates-2022.07.19-h06a4308_0
certifi pkgs/main/linux-64::certifi-2022.9.24-py39h06a4308_0
ld_impl_linux-64 pkgs/main/linux-64::ld_impl_linux-64-2.38-h1181459_1
libffi pkgs/main/linux-64::libffi-3.3-he6710b0_2
libgcc-ng pkgs/main/linux-64::libgcc-ng-11.2.0-h1234567_1
libgomp pkgs/main/linux-64::libgomp-11.2.0-h1234567_1
libstdcxx-ng pkgs/main/linux-64::libstdcxx-ng-11.2.0-h1234567_1
ncurses pkgs/main/linux-64::ncurses-6.3-h5eee18b_3
openssl pkgs/main/linux-64::openssl-1.1.1q-h7f8727e_0
pip pkgs/main/linux-64::pip-22.2.2-py39h06a4308_0
python pkgs/main/linux-64::python-3.9.13-haa1d7c7_2
readline pkgs/main/linux-64::readline-8.1.2-h7f8727e_1
setuptools pkgs/main/linux-64::setuptools-63.4.1-py39h06a4308_0
sqlite pkgs/main/linux-64::sqlite-3.39.3-h5082296_0
tk pkgs/main/linux-64::tk-8.6.12-h1ccaba5_0
tzdata pkgs/main/noarch::tzdata-2022e-h04d1e81_0
wheel pkgs/main/noarch::wheel-0.37.1-pyhd3eb1b0_0
xz pkgs/main/linux-64::xz-5.2.6-h5eee18b_0
zlib pkgs/main/linux-64::zlib-1.2.12-h5eee18b_3
Proceed ([y]/n)?
インストールが完了したら、念のためUbuntuを再起動しておきましょう。
仮想環境の起動
仮想環境を作ったので、起動しましょう。
$ conda activate dreambooth
$ conda info -e
アスタリスクの位置がdreamboothに変わっていることを確認してください。
PyTorchのインストール
では、仮想環境にPyTorchをインストールします。
公式のPyTorchサイトを参考に、コマンドを調べます。
ただし、注意点としてCondaではなくPipを選択してください。
この後、pipでライブラリをインストールするのですが、condaとpipを併用すると仮想環境が壊れる危険性があるため、pipで統一します。
この記事に沿って進めた場合、以下のコマンドになります。
$ pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116
インストールが終わったら、PyTorchからCUDAを利用できるか、以下のコマンドで確認しましょう。
$ python
>>> import torch
>>> torch.cuda.is_available()
True
Trueと表示されれば成功、なんかエラーが表示されたら失敗です。
Dreamboothのインストール
ようやく本題です。
まず、以下のコマンドでDreamboothをcloneします。
$ cd && git clone https://github.com/JoePenna/Dreambooth-Stable-Diffusion
$ cd Dreambooth-Stable-Diffusion
次に、必要なパッケージをインストールします。
$ pip install omegaconf
$ pip install einops
$ pip install pytorch-lightning==1.6.5
$ pip install test-tube
$ pip install transformers
$ pip install kornia
$ pip install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
$ pip install -e git+https://github.com/openai/CLIP.git@main#egg=clip
$ pip install setuptools==59.5.0
$ pip install pillow==9.0.1
$ pip install torchmetrics==0.6.0
$ pip install -e .
$ pip install protobuf==3.19.5
$ pip install gdown
$ pip install pydrive
$ pip install -qq diffusers["training"]==0.3.0 transformers ftfy
$ pip install -qq "ipywidgets>=7,<8"
$ pip install huggingface_hub
$ pip install ipywidgets==7.7.1
$ pip install captionizer==1.0.1
モデルの準備
ダウンロード
追加学習を行いたいモデルをダウンロードします。
今回はWaifu-Diffusionに追加学習を行う想定で解説します。
以下のリンクにアクセスします。
今回はDreamboothでトレーニングするため、wd-v1-3-full.ckptをダウンロードします。
上記のリンクから、「Float 32 Full Weights」→「download」でダウンロードできます。
ダウンロードしたら、名前をmodel.ckptにリネームします。
モデルデータの配置
ダウンロードしたモデルデータを配置します。
今回も同様に「ダウンロード」に保存した前提で進めます。
$ sudo cp /mnt/c/Users/<ユーザー名>/Downloads/model.ckpt ~/Dreambooth-Stable-Diffusion
画像の準備
最後に、正則化画像と学習元画像の準備を行います。
まずはフォルダを作成します。
$ cd ~/Dreambooth-Stable-Diffusion
$ mkdir regularization_images
$ mkdir training_images
正則化画像
正則化画像はAIに生成させます。
AUTOMATIC1111やNMKDなどGUIで操作できるソフトがおすすめです。
導入方法は解説している人が多数いるので割愛します。
私は、waifu-diffusionに1girlと入力し、1000枚ほど生成しました。(200枚程度で十分らしいです）
正則化画像を用意したら、先程作成したregularization_imagesフォルダにコピーします。
以下のコマンドをエクスプローラのアドレスバーに入力すると、WSL2内のフォルダにアクセスできるので、home→<Ubuntuユーザー名>→Dreambooth-Stable-Diffusion→regularization_imagesと進み、先程生成した正則化画像をすべてコピーします。
\\wsl$\
学習元画像
学習させたい画像は自分で用意してください。
画像の解像度は512×512に揃える必要があるので、手動でリサイズするか、一括リサイズソフトを使いましょう。
用意したら先程と同様の手順で、エクスプローラからtraining_imagesフォルダに学習元画像をすべてコピーします。
学習の実行
以下のコマンドで学習を始めることができます。
VRAMを22GBほど消費するので、余計なソフトは終了しておきましょう。
$ python "main.py" \
--base configs/stable-diffusion/v1-finetune_unfrozen.yaml \
-t \
--actual_resume "model.ckpt" \
--reg_data_root "regularization_images" \
-n "test" \
--gpus 0, \
--data_root "training_images" \
--max_training_steps 1000 \
--class_word "1girl" \
--token "sks" \
--no-test
Ubuntuを再起動した場合は、以下のコマンドを打ってから学習を開始しましょう。
$ conda activate dreambooth
$ cd ~/Dreambooth-Stable-Diffusion
2023/01/19追記
現在モデルのロードに失敗することがあります。
その場合はmain.pyの修正を修正しましょう。
main.pyを開き、30～31行目を以下のコードで上書きしてください。
sd = torch.load(ckpt, map_location="cpu")
# sd = pl_sd["state_dict"]
学習内容の調整
上記の学習開始コマンドのうち、いくつか調整が必要な項目があります。
・--n：トレーニングを識別するための名前です。自由に付けて良いです。
・--max_training_steps：トレーニングを終了するステップ数を指定します。目安としては学習元画像の100倍～150倍ほどです。
・--class_word：正則化画像の生成に使用したプロンプトを記入しましょう。
・--token：基本的に"sks"で問題ありませんが、既に"sks"で学習済みのckptに再度学習を行いたい場合などは別の単語を指定します。
学習結果
Dreambooth-Stable-Diffusion\logs\training_images...~\checkpointsにあるlast.ckptファイルが学習結果です。
これをAUTOMATIC1111に入れればそのまま使用できます。
参考
Discussion