Apple Silicon（M1 / M2）MacのGPUでStable DiffusionのLoRAを学習

2023年5月6日 2023年6月1日

Apple Silicon（M1 / M2）MacのGPU（Metal / mps）を使ってStable DiffusionのLoRAを学習させる。

環境構築
- インストール
- Acceleratorの設定
学習用のデータを準備
学習の実行
- 学習時に途中経過のサンプル画像を生成したい場合
- Pytorchバージョンの注意
学習したLoRAを使用する

環境構築

インストール

Miniforgeとかは入っていて、Pythonの仮想環境も構築済みとする。
(ここでは、PythonバージョンはPython 3.10.6で実施した)
仮想環境に入って以下のように環境構築を行う。

# metal用のpytorchをインストール
pip3 install torch torchvision

# metal用のtensorflowをインストール
conda install -c apple tensorflow-deps
pip3 install tensorflow-macos tensorflow-metal tensorboard

# diffusersのインストール
pip install diffusers

# LoRAの学習用コードをダウンロード
git clone https://github.com/bmaltais/kohya_ss.git
cd kohya_ss

# requirements.txt の tensorflow, tensorboard はコメントアウトする
pip install -r requirements.txt

Acceleratorの設定

# acceleratorの設定をする
accelerate config

＃以下のように答える

--------------------------------------------------------------------------------
In which compute environment are you running?
This machine                                                                    
--------------------------------------------------------------------------------
Which type of machine are you using?                                            
No distributed training                                                         
Do you want to run your training on CPU only (even if a GPU / Apple Silicon device is available)? [yes/NO]:NO                                                   
Do you wish to optimize your script with torch dynamo?[yes/NO]:NO               
--------------------------------------------------------------------------------
Do you wish to use FP16 or BF16 (mixed precision)?
no                                                                            
accelerate configuration saved at ~/.cache/huggingface/accelerate/default_config.yaml

学習用のデータを準備

詳しくは、配布元の説明を参照すること。

作成場所はどこでも良いが、配布元の説明にあるように、ディレクトリ構造と命名規則は守る。
画像のファイル名や種類は適当にする。

学習用のディレクトリを作成する（作成場所・名前は自由/今回はtrainとした）
学習用のディレクトリ配下に学習画像を格納するディレクトリを作る
名前は、「<繰り返し回数>_<学習させる概念を表すワード> <学習させる概念を包括するワード>」とする
例えば、自分の飼っている犬の画像をpochiという単語に学習させるなら、「20_pochi dog」のようにする
学習データを格納するディレクトリ配下に学習させる画像を入れる（画像のファイル名や種類は適当で良い）
必要であれば正規化用のディレクトリを作成する（作成場所・名前は自由/今回はregとした）
正規化画像があると、例えばpochiを学習させる際にdogがpochiに影響されるのを防いでくれるらしい
正規化用のディレクトリ配下に正規化画像を格納するディレクトリを作る
名前は、「<繰り返し回数>_<学習させる概念を包括するワード>」とする
例えば、自分の買っている犬の画像をpochiという単語に学習させるなら、「20_dog」のようにする
学習元のモデルに<学習させる概念を包括するワード>（今回はdog）をプロンプトに入力して、正規化画像を生成する
正規化画像を格納するディレクトリ配下に正規化画像を入れる（画像のファイル名や種類は適当で良い）
LoRAの出力用のディレクトリを作成する（作成場所・名前は自由/今回はoutputsとした）

.
├── train
│   └── 20_pochi dog
│       ├── AAA1.jpg
│       ├── AAA2.jpg
│       ├── AAA3.jpg
├── reg
│   └── 20_dog
│       ├── BBB1.png
│       ├── BBB2.png
│       ├── BBB3.png
│       ├── BBB4.png
│       ├── BBB5.png
└── outputs

学習の実行

kohya_gui.pyを使いたいが、UI周りのエラーが出て解消する気力がなかったので、とりあえずCUIを使う。
(コピペしてzsh: no such file or directoryが出たら\の後ろに余分な空白が無いか確認する)

accelerate launch --num_cpu_threads_per_process 4 train_network.py \
    --pretrained_model_name_or_path="your_base_model_path/model_name.safetensors" \
    --train_data_dir="kohya_ss/data/train" \
    --reg_data_dir="kohya_ss/data/reg" \
    --prior_loss_weight=1.0 --resolution 512 \
    --output_dir="kohya_ss/data/outputs" \
    --output_name=your_lora_name \
    --train_batch_size=2 --learning_rate=1e-4 --max_train_epochs 4 \
    --seed 42 --save_model_as=safetensors --save_every_n_epochs=1 \
    --max_data_loader_n_workers=1 \
    --network_module=networks.lora --network_dim=4 \
    --training_comment="activate by pochi dog" \
    --enable_bucket --bucket_no_upscale \
    --optimizer_type="AdamW" --mixed_precision="no" --save_precision="float"

–pretrained_model_name_or_path：学習元のモデルへのパス
–train_data_dir：学習用ディレクトリへのパス（今回はtrainディレクトリ）
–reg_data_dir：正規化用ディレクトリへのパス（今回はregディレクトリ）
–max_train_epochs：エポック数
–output_name：LoRAの名前（自由）
–enable_bucket –bucket_no_upscale：学習に使う画像が大きすぎると怒られるので設定した
–train_batch_size：ミニバッチサイズ（メモリに合わせて変更する）
–optimizer_type：オプティマイザーとか好きなのを使いたいなら変更する

上記の引数はとりあえず動くように調整したものなので、学習を効率的に行うには要調整。
ただし、fp16とかAdamW8bitとかを使おうとすると多分動かないので注意。
また、実行環境がM1Max（64GB）なのでメモリ周りは調整が必要かも…

その他、引数については配布元のページを参照すること。
https://github.com/kohya-ss/sd-scripts/blob/main/train_README-ja.md
https://github.com/kohya-ss/sd-scripts/blob/main/train_network_README-ja.md

学習時に途中経過のサンプル画像を生成したい場合

学習を開始するコマンドに引数を追加する

--sample_every_n_epochs=1 \
--sample_sampler="k_euler_a" \
--sample_prompts="kohya_ss/data/prompt.txt"

生成用のプロンプトをテキストファイルに保存しておく。

# prompt 1
pochi, run --w 512 --h 512 --d 1 --l 7.5 --s 20

生成用のコードがcuda用なので、kohya_ss/library/train_util.py内にあるcudaと書かれている部分をmpsに書き換えてmps用に改変する。（テキストエディタの置換であっさりと通った）

変更箇所の詳細

以下の箇所を変更しました。

① 2939行目

② 3330行目

③ 3387行目

④ 3434行目・3437行目

ただし、ダウンロードしたコードのバージョンによってはうまくいかない可能性があります。
ちなみに、筆者の環境でダウンロードしているkohya_ssレポジトリのコミットのハッシュは 52dd50fb91841ecdbb3c7b89df3ca9944ec00216 です。
その場合は、最新のコードに合わせて修正するか、以下のように筆者の環境とバージョンを揃えてコードを修正してみてください。

git clone https://github.com/bmaltais/kohya_ss.git
cd kohya_ss
git branch test_branch 52dd50fb91841ecdbb3c7b89df3ca9944ec00216
git checkout test_branch

Pytorchバージョンの注意

Pytorch1.13.0前後のバージョンを使用していると、LossがNaNになるというバグがある。
https://github.com/pytorch/pytorch/issues/88331

現在は修正されているので、アップグレードする。

pip3 install -U torch torchvision torchaudio

学習したLoRAを使用する

学習したモデルはoutput_dirで指定したディレクトリ（ここではoutputs）に保存されている。
あとは、web-uiなどを使って普通にLoRAを使用すればよい。