【お絵描きAIカスタム】おうちdreamboothやってみた【技術倉庫】
こんにちは！さおみです。
最近、お絵描きAIが流行っています。中でもぼくは、自宅のパソコンで動かせるお絵描きAIソフトに関心があります。
お絵描きAIは、モデル(絵柄を決めるデータ)を組み合わせて使います。モデルは通常、作るのに機材も時間もお金もかかるので、モデルを作成するのはなかなかたいへんであり、ぼくは他のひとが作ったモデルを使わせてもらっていました。
そんなある日、dreamboothという、お絵描きAIのモデルに新しいものを覚えさせるソフトがあるという話を聞きました。しかも改良のおかげで、一般向けのコンピューターでも(ギリギリ)動かせるようになったとのことです。
一般的には、Google Colabという、プログラミング言語のひとつであるPythonをクラウドで使えるサービスがあり、こちらで動かすのがおすすめです。ですがぼくは、
おうちパソコンでdreamboothを動かしてみたい！
と思い、挑戦してみました。
結果としてVRAM使用量を削れたのは19GBまで、また作業の途中でエラーに何度も出くわしました。中でも、bitsandbytesというソフトの近代化改修が(ぼくには)大変だったので、同じことをやってみようと思ったひとの参考になればと思い、記事を書いてみました。
やってみたことを書き連ねていったら、文字だらけでとんでもないことになってしまいました。書いてあることは大概一度以上つまずいているもので、記録として残します。必要な部分だけお読みいただければと思います。
【自宅パソコンがAI非対応の方におすすめ】
Google Colab ではじめる DreamBooth
https://note.com/npaka/n/n3b96d0b5d9c2
【出会い】
dreamboothの話を初めて聞いたとき「これはすごい！」と思ったのですが…
どうやらdreamboothはLinux専用らしい
えっ
改良のキモとなるbitsandbytesというソフトがLinux専用のようなのです(MacOSも一応入っていますがNVIDIA製品がサポート終了しています)。
“Requirements Linux distribution (Ubuntu, MacOS, etc.)”
https://github.com/TimDettmers/bitsandbytes
たまたま自宅のパソコンが条件に合っており、作業当時は夢中になっていて気づかなかったのですが、動作条件が
Linuxデスクトップパソコン+VRAM12GB以上のグラフィックカード
いやこれ条件マニアックすぎだろ…
(まあ、Windows機上の仮想マシン使用という手もあるかもしれませんが)
(追記)
と思っていたら、実際にWindowsで動かした方がいらっしゃいました。すごい…。
Diffusers版DreamboothをVRAM 12GBのWindows PCで動かす
https://note.com/kohya_ss/n/ne17e34dd51bf
というわけで、使用したパソコン環境は以下の通りです。
Ubuntu Linux 22.04.1
NVIDIA GeForce RTX 3090(VRAM 24GB)
Python 3.10 pip使用
今回特にたいへんだったのは以下の3つです。
・Pytorchまわりのバージョン合わせ
・現行バージョンのbitsandbytes導入
・Waifu-Diffusion v1.3(scheduler: LMSDiscreteScheduler)への対応
※記事中では一部手順を省略しています。ご了承ください。
【導入】
ShivamShrirao氏の改良版Dreanboothを使います。
https://github.com/ShivamShrirao/diffusers/tree/main/examples/dreambooth
今回は、ホームディレクトリ(Windowsでいうマイフォルダー)のすぐ下にdreamboothをインストールします。diffusersパッケージの一部として提供されています。
pip install git+https://github.com/ShivamShrirao/diffusers.git
dreamboothのディレクトリに移動します。
cd diffusers/examples/dreambooth/
Pythonのパッケージ(プログラムの部品)をインストールしていきますが、プログラムの部品を入れる場所をdreambooth専用に作ります(仮想環境といいます)。
ディレクトリの名前はvenvとします。
mkdir venv
python -m venv venv
仮想環境を有効にします。
source venv/bin/activate
※仮想環境を無効にするにはdeactivateを使います。
あと、あまりないとは思いますが、Python2環境を併用している方は、pythonおよびpipコマンドをそれぞれpython3, pip3としてください。
【2022.10.6以降は不要な作業 ここから】
diffusersのバージョンが0.3.0以下ですと今回使用するVRAM抑制のオプションが使えないため、先ほどダウンロードしたdiffusersをインストールします。
https://github.com/huggingface/diffusers/issues/656
cd ../../
python setup.py install
cd examples/dreambooth/
2022.10.6にdiffusersが0.4.0にバージョンアップされたので、通常のpipインストールで動作します。
【2022.10.6以降は不要な作業 ここまで】
diffusersをインストールします。
pip install diffusers
dreamboothのReadmeに沿って準備していきます。
https://github.com/ShivamShrirao/diffusers/tree/main/examples/dreambooth
必要なパッケージをインストールします。
pip install -U -r requirements.txt
プログラムの事前設定をします。
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): 2
How many different machines will you use (use more than 1 for multi-node training)? [1]: 1
Do you want to use DeepSpeed? [yes/NO]: no
Do you want to use FullyShardedDataParallel? [yes/NO]: no
How many GPU(s) should be used for distributed training? [1]:1
Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: no
mixed precisionをfp16にするとVRAM使用量が減るらしいのですが、今回は指定してもVRAM使用量があまり変わらなかったのでNOにしています。
モデルをダウンロードするため、Hugging FaceのREADアクセストークンを取得します。
https://huggingface.co/settings/tokens
(アカウント作成が必要です)
Hugging Faceにログインします。先ほど取得したアクセストークンを求められます。ブラウザーから端末にトークンを直接コピペするとうまく行かないことがあるので、トークンは一度gedit(Windowsのメモ帳的なアプリ)などに一度ペーストするとよさそうです。
huggingface-cli login
モデルをダウンロードします。モデルは10GBほどあるため、ディスク容量、通信容量に気をつけてください。
git lfs install
ここでは標準的なモデルを例にします。
git clone https://huggingface.co/CompVis/stable-diffusion-v1-4
・Pytorchまわりのバージョン合わせ
続いてグラフィックカードまわりのソフトをダウンロードしますが、以下のソフトでバージョンを合わせる必要があります。
・グラフィックカードドライバー
・CUDA
・Pytorch
以上の中でいちばんバージョン指定がきびしいのがPytorchです。通常のpipでインストールするとCUDA非対応版になってしまうので、Pytorch公式サイトで確認します。
https://pytorch.org/get-started/locally/
CUDAのバージョン11.6に対応したPytorchが2022.10.20現在最新なので、こちらをインストールします。
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116
続いて、NVIDIA公式サイトからCUDAをダウンロードします。後述するbitsandbytes用にCUDAのインストール場所を明示する必要があるのですが、この操作に対応するのがNVIDIA公式サイト版です。
https://developer.nvidia.com/cuda-toolkit-archive
バージョン11.6.2, Linux, x86_64, Ubuntuを選択しました。
が、Ubuntu 22.04用のファイルがないので、20.04用のrunfile(local)を以下コマンドでダウンロードおよびインストールします。
wget https://developer.download.nvidia.com/compute/cuda/11.6.2/local_installers/cuda_11.6.2_510.47.03_linux.run
sudo sh cuda_11.6.2_510.47.03_linux.run
手元の環境ですとドライバーをインストールしようとしたらエラーが出たので、runfileではCUDAのみをインストールしました。
Toolkit: Installed in /usr/local/cuda-11.6/
Please make sure that
- PATH includes /usr/local/cuda-11.6/bin
- LD_LIBRARY_PATH includes /usr/local/cuda-11.6/lib64, or, add /usr/local/cuda-11.6/lib64 to /etc/ld.so.conf and run ldconfig as root
また、以上のようにインストールログが出たので、CUDAの場所を環境変数に追加します。
(ちなみに、bitsandbytesは実行時にLD_LIBRARY_PATHの中にあるlibcudart.soを探しにいきます)
(参考) https://qiita.com/JeJeNeNo/items/05e148a325192004e2cd
echo 'export PATH=/usr/local/cuda-11.6/bin:${PATH}' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.6/lib64:${LD_LIBRARY_PATH}' >> ~/.bashrc
source ~/.bashrc
CUDAのrunfileバージョンが510だったので、それに合わせてグラフィックドライバーをインストールします。
sudo apt install nvidia-driver-510
sudo apt install nvidia-utils-510
本当はドライバーとCUDAでインストール元を統一したかったです。
ここでコンピューターを再起動します。
・現行バージョンのbitsandbytes導入
【作業不要 ここから】
GPUの計算でVRAM使用量を減らせるbitsandbytesですが、古いバージョンでよければ、以下のコマンドでインストールできます。
pip install bitsandbytes-cuda116
【作業不要 ここまで】
ですが、上記パッケージは2022年8月で開発が止まっているので、直近のバージョンをインストールするには以下のコマンドを使います。
$ pip install bitsandbytes
(パッケージによるバージョン指定がなくなったので、環境変数のLD_LIBRARY_PATHを手がかりにします)
bitsandbytesのインストール作業自体は以上なのですが、最初ぼくはCUDAをaptからインストールしたため、エラーの対応方法がわからず苦労しました。runfileからCUDAをインストールし、環境変数を指定することで解決しました。
シェルスクリプトから実行します。以下は、実行ファイルlaunch.shの例です。
export MODEL_NAME="/home/username/diffusers/examples/dreambooth/stable-diffusion"
export INSTANCE_DIR="/home/username/diffusers/examples/dreambooth/instance"
export CLASS_DIR="/home/username/diffusers/examples/dreambooth/class"
export OUTPUT_DIR="/home/username/diffusers/examples/dreambooth/output"
export LD_LIBRARY_PATH="/usr/local/cuda/lib64"
accelerate launch train_dreambooth.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--class_data_dir=$CLASS_DIR \
--output_dir=$OUTPUT_DIR \
--with_prior_preservation --prior_loss_weight=1.0 \
--instance_prompt="sks face" \
--class_prompt="girl face" \
--resolution=512 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 --gradient_checkpointing \
--use_8bit_adam \
--learning_rate=5e-6 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--num_class_images=200 \
--sample_batch_size=4 \
--max_train_steps=800
オプションはバックスラッシュ(半角の¥マーク)で区切ります。
MODEL_NAMEには、学習元となるモデルのディレクトリを入れます。モデルには、feature_extractor, safety_checker, scheduler, text_encoder, tokenizer, unet, vae, model_index.jsonが含まれています。
INSTANCE_DIRには学習させたい画像を入れます。
CLASS_DIRには、元になるプロンプトで生成された絵が入ります。学習前に自動生成してくれますが、自分で準備してもかまいません。--num_class_imagesオプションが生成または準備する画像の枚数です。
--use_8bit_adamと--gradient_checkpointingが、VRAM使用量を減らすのに関わるオプションです。
--class_promptには学習させたいものに近いプロンプト、--instance_promptは学習後に使いたいプロンプトを入れます。女の子キャラクターの顔を学習させたかったので学習前を"girl face", 学習後を"sks face"としています。”sks”というのは他の単語とかぶりにくいオリジナル要素の例らしく、sksをキャラクターの名前や特徴に替えてもかまいません。
--learning_rateは「学習率」というパラメーターです。”5e-6”は5×10のマイナス6乗を意味し、0.000005のことです。学習率を小さくすると、きめ細かい学習ができる反面、学習に時間がかかります。学習率を大きくすると、学習が早い半面、最後の細かい詰めが効きにくいようです。
--max_train_stepsはモデルを学習させる量です。
学習が成功すると、OUTPUT_DIRに学習後のモデルが入ります。
シェルスクリプトファイルに実行権限を付けて、実行します。
chmod +x launch.sh
./launch-sh
ぼくの環境では、画像200枚生成に約10分、学習に約10分の計約20分かかりました。
結果の確認およびチェックポイントファイルの作成に、jupyter notebookを使います。今回使ったPython環境に切り替えるために、ipywidgetsパッケージをインストールします。
pip install notebook
pip install ipywidgets
jupyter notebook
DreamBooth_Stable_Diffusion.ipynbを立ち上げ、メニューのカーネル→カーネルの変更から作成した仮想環境を選択します。Inferenceセクションでは作成したモデルの動作確認、Convert weights to ckpt to use in web UIs like AUTOMATIC1111.セクションではモデルをチェックポイントファイルに変換できます。
エラーが出る場合は、作成したモデルに以下のディレクトリを揃えると解決することがあります。
feature_extractor, safety_checker, scheduler, text_encoder, tokenizer, unet, vae
作成元のディレクトリから不足分をコピーし、作成元のmodel_index.jsonを参考に作成したモデルのmodel_index.jsonに追記します。
また、推論(Inference)の際に出力画像が真っ黒になってしまう場合は、画像がNSFW(Not Safe For Work、公共の場にふさわしくない的な)判定をされていると考えられます。
StableDiffusionPipeline.from_pretrainedからwith autocast("cuda"), torch.inference_mode():までの間に以下のコードを加えるとセーフチェックを回避できます。
def dummy(images, **kwargs): return images, False
pipe.safety_checker = dummy
(参考)
https://zenn.dev/platina/scraps/3259add3d014c9
なお、Google Colabでこのセーフチェックを回避してNSFW画像を生成させるとアカウントが使用禁止になるという話も聞いております。このコードの追加は自分のパソコンで実行してください。
・Waifu-Diffusion v1.3(scheduler: LMSDiscreteScheduler)への対応
Waifu Diffusion 1.3ではスケジューラーが変更になっています。対応方法は以下の2種類です。ただ、方法2は手元環境で動作したものの出力絵がステンドグラスのようになってしまったので、個人的には方法1がおすすめです。
【方法1 ここから】
schedulerディレクトリの中身をStable Diffusionのものと置き換えます。model_index.jsonファイルも併せて書き換えます。
【変更前】
"scheduler": [
"diffusers",
"LMSDiscreteScheduler"
],
【変更後】
"scheduler": [
"diffusers",
"PNDMScheduler"
],
【方法1 ここまで】
【方法2 ここから】
train_dreambooth.pyファイルを書き換えます。
【変更前】
from diffusers import StableDiffusionPipeline, DDIMScheduler
(中略)
pipe = StableDiffusionPipeline.from_pretrained(
"hakurei/waifu-diffusion",
torch_dtype=torch.float16,
revision="fp16",
scheduler=DDIMScheduler(
beta_start=0.00085,
beta_end=0.012,
beta_schedule="scaled_linear",
clip_sample=False,
set_alpha_to_one=False,
),
).to('cuda')
【変更後】
from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler
(中略)
pipe = StableDiffusionPipeline.from_pretrained(
"hakurei/waifu-diffusion",
torch_dtype=torch.float16,
scheduler=LMSDiscreteScheduler(
beta_start=0.00085,
beta_end=0.012,
beta_schedule='scaled_linear',
num_train_timesteps=1000
),
).to('cuda')
【方法2 ここまで】
(参考) Waifu Diffusion 1.3が動かない場合の対応
https://self-development.info/waifu-diffusion-1-3%E3%81%8C%E5%8B%95%E3%81%8B%E3%81%AA%E3%81%84%E5%A0%B4%E5%90%88%E3%81%AE%E5%AF%BE%E5%BF%9C/
https://huggingface.co/hakurei/waifu-diffusion/commit/2dff9dab944d77470b545a1bbe86fffb08d54912
お越しいただき、ありがとうございました！
