diff --git a/README-ja.md b/README-ja.md index 064464c0..47aaf16e 100644 --- a/README-ja.md +++ b/README-ja.md @@ -16,9 +16,10 @@ GUIやPowerShellスクリプトなど、より使いやすくする機能が[bma 当リポジトリ内およびnote.comに記事がありますのでそちらをご覧ください（将来的にはすべてこちらへ移すかもしれません）。 +* [学習について、共通編](./train_README-ja.md) : データ整備やオプションなど + * [データセット設定](./config_README-ja.md) * [DreamBoothの学習について](./train_db_README-ja.md) * [fine-tuningのガイド](./fine_tune_README_ja.md): -BLIPによるキャプショニングと、DeepDanbooruまたはWD14 taggerによるタグ付けを含みます * [LoRAの学習について](./train_network_README-ja.md) * [Textual Inversionの学習について](./train_ti_README-ja.md) * note.com [画像生成スクリプト](https://note.com/kohya_ss/n/n2693183a798e) @@ -131,6 +132,8 @@ pip install --use-pep517 --upgrade -r requirements.txt LoRAの実装は[cloneofsimo氏のリポジトリ](https://github.com/cloneofsimo/lora)を基にしたものです。感謝申し上げます。 +Conv2d 3x3への拡大は [cloneofsimo氏](https://github.com/cloneofsimo/lora) が最初にリリースし、KohakuBlueleaf氏が [LoCon](https://github.com/KohakuBlueleaf/LoCon) でその有効性を明らかにしたものです。KohakuBlueleaf氏に深く感謝します。 + ## ライセンス スクリプトのライセンスはASL 2.0ですが（Diffusersおよびcloneofsimo氏のリポジトリ由来のものも同様）、一部他のライセンスのコードを含みます。 diff --git a/README.md b/README.md index f8cc375e..aaf371cb 100644 --- a/README.md +++ b/README.md @@ -28,9 +28,10 @@ The scripts are tested with PyTorch 1.12.1 and 1.13.0, Diffusers 0.10.2. All documents are in Japanese currently. +* [Training guide - common](./train_README-ja.md) : data preparation, options etc... + * [Dataset config](./config_README-ja.md) * [DreamBooth training guide](./train_db_README-ja.md) * [Step by Step fine-tuning guide](./fine_tune_README_ja.md): -Including BLIP captioning and tagging by DeepDanbooru or WD14 tagger * [training LoRA](./train_network_README-ja.md) * [training Textual Inversion](./train_ti_README-ja.md) * note.com [Image generation](https://note.com/kohya_ss/n/n2693183a798e) @@ -110,11 +111,13 @@ Once the commands have completed successfully you should be ready to use the new ## Credits -The implementation for LoRA is based on [cloneofsimo's repo](https://github.com/cloneofsimo/lora). Thank you for great work!!! +The implementation for LoRA is based on [cloneofsimo's repo](https://github.com/cloneofsimo/lora). Thank you for great work! + +The LoRA expansion to Conv2d 3x3 was initially released by cloneofsimo and its effectiveness was demonstrated at [LoCon](https://github.com/KohakuBlueleaf/LoCon) by KohakuBlueleaf. Thank you so much KohakuBlueleaf! ## License -The majority of scripts is licensed under ASL 2.0 (including codes from Diffusers, cloneofsimo's), however portions of the project are available under separate license terms: +The majority of scripts is licensed under ASL 2.0 (including codes from Diffusers, cloneofsimo's and LoCon), however portions of the project are available under separate license terms: [Memory Efficient Attention Pytorch](https://github.com/lucidrains/memory-efficient-attention-pytorch): MIT @@ -124,22 +127,34 @@ The majority of scripts is licensed under ASL 2.0 (including codes from Diffuser ## Change History -- 2 Mar. 2023, 2023/3/2: +- 9 Mar. 2023, 2023/3/9: - There may be problems due to major changes. If you cannot revert back to the previous version when problems occur, please do not update for a while. - - Dependencies are updated, Please [upgrade](#upgrade) the repo. - - Add detail dataset config feature by extra config file. Thanks to fur0ut0 for this great contribution! - - Documentation is [here](./config_README-ja.md) (only in Japanese currently.) - - Specify ``.toml`` file with ``--dataset_config`` option. - - The previous options for dataset can be used as is. - - There might be a bug due to the large scale of update, please report any problems if you find. - - Add feature to generate sample images in the middle of training for each training scripts. - - ``--sample_every_n_steps`` and ``--sample_every_n_epochs`` options: frequency to generate. - - ``--sample_prompts`` option: the file contains prompts (each line generates one image.) - - The prompt is subset of ``gen_img_diffusers.py``. The prompt options ``w, h, d, l, s, n`` are supported. - - ``--sample_sampler`` option: sampler (scheduler) for generating, such as ddim or k_euler. See help for useable samplers. - - Add ``--tokenizer_cache_dir`` to each training and generation scripts to cache Tokenizer locally from Diffusers. - - Scripts will support offline training/generation after caching. - - Support letents upscaling for highres. fix, and VAE batch size in ``gen_img_diffusers.py`` (no documentation yet.) + - Minimum metadata (module name, dim, alpha and network_args) is recorded even with `--no_metadata`, issue https://github.com/kohya-ss/sd-scripts/issues/254 + - `train_network.py` supports LoRA for Conv2d-3x3 (extended to conv2d with a kernel size not 1x1). + - Same as a current version of [LoCon](https://github.com/KohakuBlueleaf/LoCon). __Thank you very much KohakuBlueleaf for your help!__ + - LoCon will be enhanced in the future. Compatibility for future versions is not guaranteed. + - Specify `--network_args` option like: `--network_args "conv_dim=4" "conv_alpha=1"` + - [Additional Networks extension](https://github.com/kohya-ss/sd-webui-additional-networks) version 0.5.0 or later is required to use 'LoRA for Conv2d-3x3' in Stable Diffusion web UI. + - __Stable Diffusion web UI built-in LoRA does not support 'LoRA for Conv2d-3x3' now. Consider carefully whether or not to use it.__ + - Merging/extracting scripts also support LoRA for Conv2d-3x3. + - Free CUDA memory after sample generation to reduce VRAM usage, issue https://github.com/kohya-ss/sd-scripts/issues/260 + - Empty caption doesn't cause error now, issue https://github.com/kohya-ss/sd-scripts/issues/258 + - Fix sample generation is crashing in Textual Inversion training when using templates, or if height/width is not divisible by 8. + - Update documents (Japanese only). + + - 大きく変更したため不具合があるかもしれません。問題が起きた時にスクリプトを前のバージョンに戻せない場合は、しばらく更新を控えてください。 + - 最低限のメタデータ（module name, dim, alpha および network_args）が `--no_metadata` オプション指定時にも記録されます。issue https://github.com/kohya-ss/sd-scripts/issues/254 + - `train_network.py` で LoRAの Conv2d-3x3 拡張に対応しました（カーネルサイズ1x1以外のConv2dにも対象範囲を拡大します）。 + - 現在のバージョンの [LoCon](https://github.com/KohakuBlueleaf/LoCon) と同一の仕様です。__KohakuBlueleaf氏のご支援に深く感謝します。__ + - LoCon が将来的に拡張された場合、それらのバージョンでの互換性は保証できません。 + - `--network_args` オプションを `--network_args "conv_dim=4" "conv_alpha=1"` のように指定してください。 + - Stable Diffusion web UI での使用には [Additional Networks extension](https://github.com/kohya-ss/sd-webui-additional-networks) のversion 0.5.0 以降が必要です。 + - __Stable Diffusion web UI の LoRA 機能は LoRAの Conv2d-3x3 拡張に対応していないようです。使用するか否か慎重にご検討ください。__ + - マージ、抽出のスクリプトについても LoRA の Conv2d-3x3 拡張に対応しました. + - サンプル画像生成後にCUDAメモリを解放しVRAM使用量を削減しました。 issue https://github.com/kohya-ss/sd-scripts/issues/260 + - 空のキャプションが使えるようになりました。 issue https://github.com/kohya-ss/sd-scripts/issues/258 + - Textual Inversion 学習でテンプレートを使ったとき、height/width が 8 で割り切れなかったときにサンプル画像生成がクラッシュするのを修正しました。 + - ドキュメント類を更新しました。 - Sample image generation: A prompt file might look like this, for example @@ -163,22 +178,6 @@ The majority of scripts is licensed under ASL 2.0 (including codes from Diffuser The prompt weighting such as `( )` and `[ ]` are not working. - - 大きく変更したため不具合があるかもしれません。問題が起きた時にスクリプトを前のバージョンに戻せない場合は、しばらく更新を控えてください。 - - ライブラリを更新しました。[アップグレード](https://github.com/kohya-ss/sd-scripts/blob/main/README-ja.md#%E3%82%A2%E3%83%83%E3%83%97%E3%82%B0%E3%83%AC%E3%83%BC%E3%83%89)に従って更新してください。 - - 設定ファイルによるデータセット定義機能を追加しました。素晴らしいPRを提供していただいた fur0ut0 氏に感謝します。 - - ドキュメントは[こちら](./config_README-ja.md)。 - - ``--dataset_config`` オプションで ``.toml`` ファイルを指定してください。 - - 今までのオプションはそのまま使えます。 - - 大規模なアップデートのため、もし不具合がありましたらご報告ください。 - - 学習の途中でサンプル画像を生成する機能を各学習スクリプトに追加しました。 - - ``--sample_every_n_steps`` と ``--sample_every_n_epochs`` オプション：生成頻度を指定 - - ``--sample_prompts`` オプション：プロンプトを記述したファイルを指定（1行ごとに1枚の画像を生成） - - プロンプトには ``gen_img_diffusers.py`` のプロンプトオプションの一部、 ``w, h, d, l, s, n`` が使えます。 - - ``--sample_sampler`` オプション：ddim や k_euler などの sampler (scheduler) を指定します。使用できる sampler についてはヘルプをご覧ください。 - - ``--tokenizer_cache_dir`` オプションを各学習スクリプトおよび生成スクリプトに追加しました。Diffusers から Tokenizer を取得してきてろーかるに保存します。 - - 一度キャッシュしておくことでオフライン学習、生成ができるかもしれません。 - - ``gen_img_diffusers.py`` で highres. fix での letents upscaling と VAE のバッチサイズ指定に対応しました。 - - サンプル画像生成： プロンプトファイルは例えば以下のようになります。 diff --git a/fine_tune_README_ja.md b/fine_tune_README_ja.md index 9dcd34af..686947c9 100644 --- a/fine_tune_README_ja.md +++ b/fine_tune_README_ja.md @@ -1,6 +1,9 @@ -NovelAIの提案した学習手法、自動キャプションニング、タグ付け、Windows＋VRAM 12GB（v1.4/1.5の場合）環境等に対応したfine tuningです。 +NovelAIの提案した学習手法、自動キャプションニング、タグ付け、Windows＋VRAM 12GB（SD v1.xの場合）環境等に対応したfine tuningです。ここでfine tuningとは、モデルを画像とキャプションで学習することを指します（LoRAやTextual Inversion、Hypernetworksは含みません） + +[学習についての共通ドキュメント](./train_README-ja.md) もあわせてご覧ください。 + +# 概要 -## 概要 Diffusersを用いてStable DiffusionのU-Netのfine tuningを行います。NovelAIの記事にある以下の改善に対応しています（Aspect Ratio BucketingについてはNovelAIのコードを参考にしましたが、最終的なコードはすべてオリジナルです）。 * CLIP（Text Encoder）の最後の層ではなく最後から二番目の層の出力を用いる。 @@ -13,19 +16,24 @@ Diffusersを用いてStable DiffusionのU-Netのfine tuningを行います。Nov デフォルトではText Encoderの学習は行いません。モデル全体のfine tuningではU-Netだけを学習するのが一般的なようです（NovelAIもそのようです）。オプション指定でText Encoderも学習対象とできます。 -## 追加機能について -### CLIPの出力の変更 +# 追加機能について + +## CLIPの出力の変更 + プロンプトを画像に反映するため、テキストの特徴量への変換を行うのがCLIP（Text Encoder）です。Stable DiffusionではCLIPの最後の層の出力を用いていますが、それを最後から二番目の層の出力を用いるよう変更できます。NovelAIによると、これによりより正確にプロンプトが反映されるようになるとのことです。 元のまま、最後の層の出力を用いることも可能です。 + ※Stable Diffusion 2.0では最後から二番目の層をデフォルトで使います。clip_skipオプションを指定しないでください。 -### 正方形以外の解像度での学習 +## 正方形以外の解像度での学習 + Stable Diffusionは512\*512で学習されていますが、それに加えて256\*1024や384\*640といった解像度でも学習します。これによりトリミングされる部分が減り、より正しくプロンプトと画像の関係が学習されることが期待されます。 学習解像度はパラメータとして与えられた解像度の面積（＝メモリ使用量）を超えない範囲で、64ピクセル単位で縦横に調整、作成されます。 機械学習では入力サイズをすべて統一するのが一般的ですが、特に制約があるわけではなく、実際は同一のバッチ内で統一されていれば大丈夫です。NovelAIの言うbucketingは、あらかじめ教師データを、アスペクト比に応じた学習解像度ごとに分類しておくことを指しているようです。そしてバッチを各bucket内の画像で作成することで、バッチの画像サイズを統一します。 -### トークン長の75から225への拡張 +## トークン長の75から225への拡張 + Stable Diffusionでは最大75トークン（開始・終了を含むと77トークン）ですが、それを225トークンまで拡張します。 ただしCLIPが受け付ける最大長は75トークンですので、225トークンの場合、単純に三分割してCLIPを呼び出してから結果を連結しています。 @@ -33,296 +41,67 @@ Stable Diffusionでは最大75トークン（開始・終了を含むと77トー ※Automatic1111氏のWeb UIではカンマを意識して分割、といったこともしているようですが、私の場合はそこまでしておらず単純な分割です。 -## 環境整備 +# 学習の手順 -このリポジトリの[README](./README-ja.md)を参照してください。 +あらかじめこのリポジトリのREADMEを参照し、環境整備を行ってください。 -## 教師データの用意 - -学習させたい画像データを用意し、任意のフォルダに入れてください。リサイズ等の事前の準備は必要ありません。 -ただし学習解像度よりもサイズが小さい画像については、超解像などで品質を保ったまま拡大しておくことをお勧めします。 - -複数の教師データフォルダにも対応しています。前処理をそれぞれのフォルダに対して実行する形となります。 - -たとえば以下のように画像を格納します。 - -![教師データフォルダのスクショ](https://user-images.githubusercontent.com/52813779/208907739-8e89d5fa-6ca8-4b60-8927-f484d2a9ae04.png) - -## 自動キャプショニング -キャプションを使わずタグだけで学習する場合はスキップしてください。 - -また手動でキャプションを用意する場合、キャプションは教師データ画像と同じディレクトリに、同じファイル名、拡張子.caption等で用意してください。各ファイルは1行のみのテキストファイルとします。 - -### BLIPによるキャプショニング - -最新版ではBLIPのダウンロード、重みのダウンロード、仮想環境の追加は不要になりました。そのままで動作します。 - -finetuneフォルダ内のmake_captions.pyを実行します。 - -``` -python finetune\make_captions.py --batch_size <バッチサイズ> <教師データフォルダ> -``` - -バッチサイズ8、教師データを親フォルダのtrain_dataに置いた場合、以下のようになります。 - -``` -python finetune\make_captions.py --batch_size 8 ..\train_data -``` - -キャプションファイルが教師データ画像と同じディレクトリに、同じファイル名、拡張子.captionで作成されます。 - -batch_sizeはGPUのVRAM容量に応じて増減してください。大きいほうが速くなります（VRAM 12GBでももう少し増やせると思います）。 -max_lengthオプションでキャプションの最大長を指定できます。デフォルトは75です。モデルをトークン長225で学習する場合には長くしても良いかもしれません。 -caption_extensionオプションでキャプションの拡張子を変更できます。デフォルトは.captionです（.txtにすると後述のDeepDanbooruと競合します）。 - -複数の教師データフォルダがある場合には、それぞれのフォルダに対して実行してください。 - -なお、推論にランダム性があるため、実行するたびに結果が変わります。固定する場合には--seedオプションで「--seed 42」のように乱数seedを指定してください。 - -その他のオプションは--helpでヘルプをご参照ください（パラメータの意味についてはドキュメントがまとまっていないようで、ソースを見るしかないようです）。 - -デフォルトでは拡張子.captionでキャプションファイルが生成されます。 - -![captionが生成されたフォルダ](https://user-images.githubusercontent.com/52813779/208908845-48a9d36c-f6ee-4dae-af71-9ab462d1459e.png) - -たとえば以下のようなキャプションが付きます。 - -![キャプションと画像](https://user-images.githubusercontent.com/52813779/208908947-af936957-5d73-4339-b6c8-945a52857373.png) - -## DeepDanbooruによるタグ付け -danbooruタグのタグ付け自体を行わない場合は「キャプションとタグ情報の前処理」に進んでください。 - -タグ付けはDeepDanbooruまたはWD14Taggerで行います。WD14Taggerのほうが精度が良いようです。WD14Taggerでタグ付けする場合は、次の章へ進んでください。 - -### 環境整備 -DeepDanbooru https://github.com/KichangKim/DeepDanbooru を作業フォルダにcloneしてくるか、zipをダウンロードして展開します。私はzipで展開しました。 -またDeepDanbooruのReleasesのページ https://github.com/KichangKim/DeepDanbooru/releases の「DeepDanbooru Pretrained Model v3-20211112-sgd-e28」のAssetsから、deepdanbooru-v3-20211112-sgd-e28.zipをダウンロードしてきてDeepDanbooruのフォルダに展開します。 - -以下からダウンロードします。Assetsをクリックして開き、そこからダウンロードします。 - -![DeepDanbooruダウンロードページ](https://user-images.githubusercontent.com/52813779/208909417-10e597df-7085-41ee-bd06-3e856a1339df.png) - -以下のようなこういうディレクトリ構造にしてください - -![DeepDanbooruのディレクトリ構造](https://user-images.githubusercontent.com/52813779/208909486-38935d8b-8dc6-43f1-84d3-fef99bc471aa.png) - -Diffusersの環境に必要なライブラリをインストールします。DeepDanbooruのフォルダに移動してインストールします（実質的にはtensorflow-ioが追加されるだけだと思います）。 - -``` -pip install -r requirements.txt -``` - -続いてDeepDanbooru自体をインストールします。 - -``` -pip install . -``` - -以上でタグ付けの環境整備は完了です。 - -### タグ付けの実施 -DeepDanbooruのフォルダに移動し、deepdanbooruを実行してタグ付けを行います。 - -``` -deepdanbooru evaluate <教師データフォルダ> --project-path deepdanbooru-v3-20211112-sgd-e28 --allow-folder --save-txt -``` - -教師データを親フォルダのtrain_dataに置いた場合、以下のようになります。 - -``` -deepdanbooru evaluate ../train_data --project-path deepdanbooru-v3-20211112-sgd-e28 --allow-folder --save-txt -``` - -タグファイルが教師データ画像と同じディレクトリに、同じファイル名、拡張子.txtで作成されます。1件ずつ処理されるためわりと遅いです。 - -複数の教師データフォルダがある場合には、それぞれのフォルダに対して実行してください。 - -以下のように生成されます。 - -![DeepDanbooruの生成ファイル](https://user-images.githubusercontent.com/52813779/208909855-d21b9c98-f2d3-4283-8238-5b0e5aad6691.png) - -こんな感じにタグが付きます（すごい情報量……）。 - -![DeepDanbooruタグと画像](https://user-images.githubusercontent.com/52813779/208909908-a7920174-266e-48d5-aaef-940aba709519.png) - -## WD14Taggerによるタグ付け -DeepDanbooruの代わりにWD14Taggerを用いる手順です。 - -Automatic1111氏のWebUIで使用しているtaggerを利用します。こちらのgithubページ（https://github.com/toriato/stable-diffusion-webui-wd14-tagger#mrsmilingwolfs-model-aka-waifu-diffusion-14-tagger ）の情報を参考にさせていただきました。 - -最初の環境整備で必要なモジュールはインストール済みです。また重みはHugging Faceから自動的にダウンロードしてきます。 - -### タグ付けの実施 -スクリプトを実行してタグ付けを行います。 -``` -python tag_images_by_wd14_tagger.py --batch_size <バッチサイズ> <教師データフォルダ> -``` - -教師データを親フォルダのtrain_dataに置いた場合、以下のようになります。 -``` -python tag_images_by_wd14_tagger.py --batch_size 4 ..\train_data -``` - -初回起動時にはモデルファイルがwd14_tagger_modelフォルダに自動的にダウンロードされます（フォルダはオプションで変えられます）。以下のようになります。 - -![ダウンロードされたファイル](https://user-images.githubusercontent.com/52813779/208910447-f7eb0582-90d6-49d3-a666-2b508c7d1842.png) - -タグファイルが教師データ画像と同じディレクトリに、同じファイル名、拡張子.txtで作成されます。 - -![生成されたタグファイル](https://user-images.githubusercontent.com/52813779/208910534-ea514373-1185-4b7d-9ae3-61eb50bc294e.png) - -![タグと画像](https://user-images.githubusercontent.com/52813779/208910599-29070c15-7639-474f-b3e4-06bd5a3df29e.png) - -threshオプションで、判定されたタグのconfidence（確信度）がいくつ以上でタグをつけるかが指定できます。デフォルトはWD14Taggerのサンプルと同じ0.35です。値を下げるとより多くのタグが付与されますが、精度は下がります。 -batch_sizeはGPUのVRAM容量に応じて増減してください。大きいほうが速くなります（VRAM 12GBでももう少し増やせると思います）。caption_extensionオプションでタグファイルの拡張子を変更できます。デフォルトは.txtです。 -model_dirオプションでモデルの保存先フォルダを指定できます。 -またforce_downloadオプションを指定すると保存先フォルダがあってもモデルを再ダウンロードします。 - -複数の教師データフォルダがある場合には、それぞれのフォルダに対して実行してください。 - -## キャプションとタグ情報の前処理 - -スクリプトから処理しやすいようにキャプションとタグをメタデータとしてひとつのファイルにまとめます。 - -### キャプションの前処理 - -キャプションをメタデータに入れるには、作業フォルダ内で以下を実行してください（キャプションを学習に使わない場合は実行不要です）（実際は1行で記述します、以下同様）。 - -``` -python merge_captions_to_metadata.py <教師データフォルダ> -　 --in_json <読み込むメタデータファイル名> - <メタデータファイル名> -``` - -メタデータファイル名は任意の名前です。 -教師データがtrain_data、読み込むメタデータファイルなし、メタデータファイルがmeta_cap.jsonの場合、以下のようになります。 - -``` -python merge_captions_to_metadata.py train_data meta_cap.json -``` - -caption_extensionオプションでキャプションの拡張子を指定できます。 - -複数の教師データフォルダがある場合には、full_path引数を指定してください（メタデータにフルパスで情報を持つようになります）。そして、それぞれのフォルダに対して実行してください。 - -``` -python merge_captions_to_metadata.py --full_path - train_data1 meta_cap1.json -python merge_captions_to_metadata.py --full_path --in_json meta_cap1.json - train_data2 meta_cap2.json -``` - -in_jsonを省略すると書き込み先メタデータファイルがあるとそこから読み込み、そこに上書きします。 - -__※in_jsonオプションと書き込み先を都度書き換えて、別のメタデータファイルへ書き出すようにすると安全です。__ - -### タグの前処理 - -同様にタグもメタデータにまとめます（タグを学習に使わない場合は実行不要です）。 -``` -python merge_dd_tags_to_metadata.py <教師データフォルダ> - --in_json <読み込むメタデータファイル名> - <書き込むメタデータファイル名> -``` - -先と同じディレクトリ構成で、meta_cap.jsonを読み、meta_cap_dd.jsonに書きだす場合、以下となります。 -``` -python merge_dd_tags_to_metadata.py train_data --in_json meta_cap.json meta_cap_dd.json -``` - -複数の教師データフォルダがある場合には、full_path引数を指定してください。そして、それぞれのフォルダに対して実行してください。 - -``` -python merge_dd_tags_to_metadata.py --full_path --in_json meta_cap2.json - train_data1 meta_cap_dd1.json -python merge_dd_tags_to_metadata.py --full_path --in_json meta_cap_dd1.json - train_data2 meta_cap_dd2.json -``` - -in_jsonを省略すると書き込み先メタデータファイルがあるとそこから読み込み、そこに上書きします。 - -__※in_jsonオプションと書き込み先を都度書き換えて、別のメタデータファイルへ書き出すようにすると安全です。__ - -### キャプションとタグのクリーニング -ここまででメタデータファイルにキャプションとDeepDanbooruのタグがまとめられています。ただ自動キャプショニングにしたキャプションは表記ゆれなどがあり微妙（※）ですし、タグにはアンダースコアが含まれていたりratingが付いていたりしますので（DeepDanbooruの場合）、エディタの置換機能などを用いてキャプションとタグのクリーニングをしたほうがいいでしょう。 - -※たとえばアニメ絵の少女を学習する場合、キャプションにはgirl/girls/woman/womenなどのばらつきがあります。また「anime girl」なども単に「girl」としたほうが適切かもしれません。 - -クリーニング用のスクリプトが用意してありますので、スクリプトの内容を状況に応じて編集してお使いください。 - -（教師データフォルダの指定は不要になりました。メタデータ内の全データをクリーニングします。） - -``` -python clean_captions_and_tags.py <読み込むメタデータファイル名> <書き込むメタデータファイル名> -``` - ---in_jsonは付きませんのでご注意ください。たとえば次のようになります。 - -``` -python clean_captions_and_tags.py meta_cap_dd.json meta_clean.json -``` - -以上でキャプションとタグの前処理は完了です。 - -## latentsの事前取得 - -学習を高速に進めるためあらかじめ画像の潜在表現を取得しディスクに保存しておきます。あわせてbucketing（教師データをアスペクト比に応じて分類する）を行います。 - -作業フォルダで以下のように入力してください。 -``` -python prepare_buckets_latents.py <教師データフォルダ> - <読み込むメタデータファイル名> <書き込むメタデータファイル名> - - --batch_size <バッチサイズ> - --max_resolution <解像度 幅,高さ> - --mixed_precision <精度> -``` - -モデルがmodel.ckpt、バッチサイズ4、学習解像度は512\*512、精度no（float32）で、meta_clean.jsonからメタデータを読み込み、meta_lat.jsonに書き込む場合、以下のようになります。 - -``` -python prepare_buckets_latents.py - train_data meta_clean.json meta_lat.json model.ckpt - --batch_size 4 --max_resolution 512,512 --mixed_precision no -``` - -教師データフォルダにnumpyのnpz形式でlatentsが保存されます。 - -Stable Diffusion 2.0のモデルを読み込む場合は--v2オプションを指定してください（--v_parameterizationは不要です）。 - -解像度の最小サイズを--min_bucket_resoオプションで、最大サイズを--max_bucket_resoで指定できます。デフォルトはそれぞれ256、1024です。たとえば最小サイズに384を指定すると、256\*1024や320\*768などの解像度は使わなくなります。 -解像度を768\*768のように大きくした場合、最大サイズに1280などを指定すると良いでしょう。 - ---flip_augオプションを指定すると左右反転のaugmentation（データ拡張）を行います。疑似的にデータ量を二倍に増やすことができますが、データが左右対称でない場合に指定すると（例えばキャラクタの外見、髪型など）学習がうまく行かなくなります。 -（反転した画像についてもlatentsを取得し、\*\_flip.npzファイルを保存する単純な実装です。fline_tune.pyには特にオプション指定は必要ありません。\_flip付きのファイルがある場合、flip付き・なしのファイルを、ランダムに読み込みます。） - -バッチサイズはVRAM 12GBでももう少し増やせるかもしれません。 -解像度は64で割り切れる数字で、"幅,高さ"で指定します。解像度はfine tuning時のメモリサイズに直結します。VRAM 12GBでは512,512が限界と思われます（※）。16GBなら512,704や512,768まで上げられるかもしれません。なお256,256等にしてもVRAM 8GBでは厳しいようです（パラメータやoptimizerなどは解像度に関係せず一定のメモリが必要なため）。 - -※batch size 1の学習で12GB VRAM、640,640で動いたとの報告もありました。 - -以下のようにbucketingの結果が表示されます。 - -![bucketingの結果](https://user-images.githubusercontent.com/52813779/208911419-71c00fbb-2ce6-49d5-89b5-b78d7715e441.png) - -複数の教師データフォルダがある場合には、full_path引数を指定してください。そして、それぞれのフォルダに対して実行してください。 -``` -python prepare_buckets_latents.py --full_path - train_data1 meta_clean.json meta_lat1.json model.ckpt - --batch_size 4 --max_resolution 512,512 --mixed_precision no - -python prepare_buckets_latents.py --full_path - train_data2 meta_lat1.json meta_lat2.json model.ckpt - --batch_size 4 --max_resolution 512,512 --mixed_precision no - -``` -読み込み元と書き込み先を同じにすることも可能ですが別々の方が安全です。 - -__※引数を都度書き換えて、別のメタデータファイルに書き込むと安全です。__ +## データの準備 +[学習データの準備について](./train_README-ja.md) を参照してください。fine tuningではメタデータを用いるfine tuning方式のみ対応しています。 ## 学習の実行 -たとえば以下のように実行します。以下は省メモリ化のための設定です。 +たとえば以下のように実行します。以下は省メモリ化のための設定です。それぞれの行を必要に応じて書き換えてください。 + +``` +accelerate launch --num_cpu_threads_per_process 1 fine_tune.py + --pretrained_model_name_or_path=<.ckptまたは.safetensordまたはDiffusers版モデルのディレクトリ> + --output_dir=<学習したモデルの出力先フォルダ> + --output_name=<学習したモデル出力時のファイル名> + --dataset_config=<データ準備で作成した.tomlファイル> + --save_model_as=safetensors + --learning_rate=5e-6 --max_train_steps=10000 + --use_8bit_adam --xformers --gradient_checkpointing + --mixed_precision=fp16 +``` + +`num_cpu_threads_per_process` には通常は1を指定するとよいようです。 + +`pretrained_model_name_or_path` に追加学習を行う元となるモデルを指定します。Stable Diffusionのcheckpointファイル（.ckptまたは.safetensors）、Diffusersのローカルディスクにあるモデルディレクトリ、DiffusersのモデルID（"stabilityai/stable-diffusion-2"など）が指定できます。 + +`output_dir` に学習後のモデルを保存するフォルダを指定します。`output_name` にモデルのファイル名を拡張子を除いて指定します。`save_model_as` でsafetensors形式での保存を指定しています。 + +`dataset_config` に `.toml` ファイルを指定します。ファイル内でのバッチサイズ指定は、当初はメモリ消費を抑えるために `1` としてください。 + +学習させるステップ数 `max_train_steps` を10000とします。学習率 `learning_rate` はここでは5e-6を指定しています。 + +省メモリ化のため `mixed_precision="fp16"` を指定します（RTX30 シリーズ以降では `bf16` も指定できます。環境整備時にaccelerateに行った設定と合わせてください）。また `gradient_checkpointing` を指定します。 + +オプティマイザ（モデルを学習データにあうように最適化＝学習させるクラス）にメモリ消費の少ない 8bit AdamW を使うため、 `optimizer_type="AdamW8bit"` を指定します。 + +`xformers` オプションを指定し、xformersのCrossAttentionを用います。xformersをインストールしていない場合やエラーとなる場合（環境にもよりますが `mixed_precision="no"` の場合など）、代わりに `mem_eff_attn` オプションを指定すると省メモリ版CrossAttentionを使用します（速度は遅くなります）。 + +ある程度メモリがある場合は、`.toml` ファイルを編集してバッチサイズをたとえば `4` くらいに増やしてください（高速化と精度向上の可能性があります）。 + +### よく使われるオプションについて + +以下の場合にはオプションに関するドキュメントを参照してください。 + +- Stable Diffusion 2.xまたはそこからの派生モデルを学習する +- clip skipを2以上を前提としたモデルを学習する +- 75トークンを超えたキャプションで学習する + +### バッチサイズについて + +モデル全体を学習するためLoRA等の学習に比べるとメモリ消費量は多くなります（DreamBoothと同じ）。 + +### 学習率について + +1e-6から5e-6程度が一般的なようです。他のfine tuningの例なども参照してみてください。 + +### 以前の形式のデータセット指定をした場合のコマンドライン + +解像度やバッチサイズをオプションで指定します。コマンドラインの例は以下の通りです。 + ``` accelerate launch --num_cpu_threads_per_process 1 fine_tune.py --pretrained_model_name_or_path=model.ckpt @@ -336,76 +115,7 @@ accelerate launch --num_cpu_threads_per_process 1 fine_tune.py --save_every_n_epochs=4 ``` -accelerateのnum_cpu_threads_per_processには通常は1を指定するとよいようです。 - -pretrained_model_name_or_pathに学習対象のモデルを指定します（Stable DiffusionのcheckpointかDiffusersのモデル）。Stable Diffusionのcheckpointは.ckptと.safetensorsに対応しています（拡張子で自動判定）。 - -in_jsonにlatentをキャッシュしたときのメタデータファイルを指定します。 - -train_data_dirに教師データのフォルダを、output_dirに学習後のモデルの出力先フォルダを指定します。 - -shuffle_captionを指定すると、キャプション、タグをカンマ区切りされた単位でシャッフルして学習します（Waifu Diffusion v1.3で行っている手法です）。 -（先頭のトークンのいくつかをシャッフルせずに固定できます。その他のオプションのkeep_tokensをご覧ください。） - -train_batch_sizeにバッチサイズを指定します。VRAM 12GBでは1か2程度を指定してください。解像度によっても指定可能な数は変わってきます。 -学習に使用される実際のデータ量は「バッチサイズ×ステップ数」です。バッチサイズを増やした時には、それに応じてステップ数を下げることが可能です。 - -learning_rateに学習率を指定します。たとえばWaifu Diffusion v1.3は5e-6のようです。 -max_train_stepsにステップ数を指定します。 - -use_8bit_adamを指定すると8-bit Adam Optimizerを使用します。省メモリ化、高速化されますが精度は下がる可能性があります。 - -xformersを指定するとCrossAttentionを置換して省メモリ化、高速化します。 -※11/9時点ではfloat32の学習ではxformersがエラーになるため、bf16/fp16を使うか、代わりにmem_eff_attnを指定して省メモリ版CrossAttentionを使ってください（速度はxformersに劣ります）。 - -gradient_checkpointingで勾配の途中保存を有効にします。速度は遅くなりますが使用メモリ量が減ります。 - -mixed_precisionで混合精度を使うか否かを指定します。"fp16"または"bf16"を指定すると省メモリになりますが精度は劣ります。 -"fp16"と"bf16"は使用メモリ量はほぼ同じで、bf16の方が学習結果は良くなるとの話もあります（試した範囲ではあまり違いは感じられませんでした）。 -"no"を指定すると使用しません（float32になります）。 - -※bf16で学習したcheckpointをAUTOMATIC1111氏のWeb UIで読み込むとエラーになるようです。これはデータ型のbfloat16がWeb UIのモデルsafety checkerでエラーとなるためのようです。save_precisionオプションを指定してfp16またはfloat32形式で保存してください。またはsafetensors形式で保管しても良さそうです。 - -save_every_n_epochsを指定するとそのエポックだけ経過するたびに学習中のモデルを保存します。 - -### Stable Diffusion 2.0対応 -Hugging Faceのstable-diffusion-2-baseを使う場合は--v2オプションを、stable-diffusion-2または768-v-ema.ckptを使う場合は--v2と--v_parameterizationの両方のオプションを指定してください。 - -### メモリに余裕がある場合に精度や速度を上げる -まずgradient_checkpointingを外すと速度が上がります。ただし設定できるバッチサイズが減りますので、精度と速度のバランスを見ながら設定してください。 - -バッチサイズを増やすと速度、精度が上がります。メモリが足りる範囲で、1データ当たりの速度を確認しながら増やしてください（メモリがぎりぎりになるとかえって速度が落ちることがあります）。 - -### 使用するCLIP出力の変更 -clip_skipオプションに2を指定すると、後ろから二番目の層の出力を用います。1またはオプション省略時は最後の層を用います。 -学習したモデルはAutomatic1111氏のWeb UIで推論できるはずです。 - -※SD2.0はデフォルトで後ろから二番目の層を使うため、SD2.0の学習では指定しないでください。 - -学習対象のモデルがもともと二番目の層を使うように学習されている場合は、2を指定するとよいでしょう。 - -そうではなく最後の層を使用していた場合はモデル全体がそれを前提に学習されています。そのため改めて二番目の層を使用して学習すると、望ましい学習結果を得るにはある程度の枚数の教師データ、長めの学習が必要になるかもしれません。 - -### トークン長の拡張 -max_token_lengthに150または225を指定することでトークン長を拡張して学習できます。 -学習したモデルはAutomatic1111氏のWeb UIで推論できるはずです。 - -clip_skipと同様に、モデルの学習状態と異なる長さで学習するには、ある程度の教師データ枚数、長めの学習時間が必要になると思われます。 - -### 学習ログの保存 -logging_dirオプションにログ保存先フォルダを指定してください。TensorBoard形式のログが保存されます。 - -たとえば--logging_dir=logsと指定すると、作業フォルダにlogsフォルダが作成され、その中の日時フォルダにログが保存されます。 -また--log_prefixオプションを指定すると、日時の前に指定した文字列が追加されます。「--logging_dir=logs --log_prefix=fine_tune_style1」などとして識別用にお使いください。 - -TensorBoardでログを確認するには、別のコマンドプロンプトを開き、作業フォルダで以下のように入力します（tensorboardはDiffusersのインストール時にあわせてインストールされると思いますが、もし入っていないならpip install tensorboardで入れてください）。 -``` -tensorboard --logdir=logs -``` - -### Hypernetworkの学習 -別の記事で解説予定です。 - + -### その他のオプション +# fine tuning特有のその他の主なオプション -#### keep_tokens -数値を指定するとキャプションの先頭から、指定した数だけのトークン（カンマ区切りの文字列）をシャッフルせず固定します。 +すべてのオプションについては別文書を参照してください。 -キャプションとタグが両方ある場合、学習時のプロンプトは「キャプション,タグ1,タグ2……」のように連結されますので、「--keep_tokens=1」とすれば、学習時にキャプションが必ず先頭に来るようになります。 - -#### dataset_repeats -データセットの枚数が極端に少ない場合、epochがすぐに終わってしまうため（epochの区切りで少し時間が掛かります）、数値を指定してデータを何倍かしてepochを長めにしてください。 - -#### train_text_encoder +## `train_text_encoder` Text Encoderも学習対象とします。メモリ使用量が若干増加します。 通常のfine tuningではText Encoderは学習対象としませんが（恐らくText Encoderの出力に従うようにU-Netを学習するため）、学習データ数が少ない場合には、DreamBoothのようにText Encoder側に学習させるのも有効的なようです。 -#### save_precision -checkpoint保存時のデータ形式をfloat、fp16、bf16から指定できます（未指定時は学習中のデータ形式と同じ）。ディスク容量が節約できますがモデルによる生成結果は変わってきます。またfloatやfp16を指定すると、1111氏のWeb UIでも読めるようになるはずです。 - -※VAEについては元のcheckpointのデータ形式のままになりますので、fp16でもモデルサイズが2GB強まで小さくならない場合があります。 - -#### save_model_as -モデルの保存形式を指定します。ckpt、safetensors、diffusers、diffusers_safetensorsのいずれかを指定してください。 - -Stable Diffusion形式（ckptまたはsafetensors）を読み込み、Diffusers形式で保存する場合、不足する情報はHugging Faceからv1.5またはv2.1の情報を落としてきて補完します。 - -#### use_safetensors -このオプションを指定するとsafetensors形式でcheckpointを保存します。保存形式はデフォルト（読み込んだ形式と同じ）になります。 - -#### save_stateとresume -save_stateオプションで、途中保存時および最終保存時に、checkpointに加えてoptimizer等の学習状態をフォルダに保存します。これにより中断してから学習再開したときの精度低下が避けられます（optimizerは状態を持ちながら最適化をしていくため、その状態がリセットされると再び初期状態から最適化を行わなくてはなりません）。なお、Accelerateの仕様でステップ数は保存されません。 - -スクリプト起動時、resumeオプションで状態の保存されたフォルダを指定すると再開できます。 - -学習状態は一回の保存あたり5GB程度になりますのでディスク容量にご注意ください。 - -#### gradient_accumulation_steps -指定したステップ数だけまとめて勾配を更新します。バッチサイズを増やすのと同様の効果がありますが、メモリを若干消費します。 - -※Accelerateの仕様で学習モデルが複数の場合には対応していないとのことですので、Text Encoderを学習対象にして、このオプションに2以上の値を指定するとエラーになるかもしれません。 - -#### lr_scheduler / lr_warmup_steps -lr_schedulerオプションで学習率のスケジューラをlinear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmupから選べます。デフォルトはconstantです。 - -lr_warmup_stepsでスケジューラのウォームアップ（だんだん学習率を変えていく）ステップ数を指定できます。詳細については各自お調べください。 - -#### diffusers_xformers +## `diffusers_xformers` スクリプト独自のxformers置換機能ではなくDiffusersのxformers機能を利用します。Hypernetworkの学習はできなくなります。 diff --git a/gen_img_diffusers.py b/gen_img_diffusers.py index 6bab0bb8..8a185170 100644 --- a/gen_img_diffusers.py +++ b/gen_img_diffusers.py @@ -1649,10 +1649,11 @@ def get_unweighted_text_embeddings( if pad == eos: # v1 text_input_chunk[:, -1] = text_input[0, -1] else: # v2 - if text_input_chunk[:, -1] != eos and text_input_chunk[:, -1] != pad: # 最後に普通の文字がある - text_input_chunk[:, -1] = eos - if text_input_chunk[:, 1] == pad: # BOSだけであとはPAD - text_input_chunk[:, 1] = eos + for j in range(len(text_input_chunk)): + if text_input_chunk[j, -1] != eos and text_input_chunk[j, -1] != pad: # 最後に普通の文字がある + text_input_chunk[j, -1] = eos + if text_input_chunk[j, 1] == pad: # BOSだけであとはPAD + text_input_chunk[j, 1] = eos if clip_skip is None or clip_skip == 1: text_embedding = pipe.text_encoder(text_input_chunk)[0] @@ -2276,13 +2277,26 @@ def main(args): mask_images = l # 画像サイズにオプション指定があるときはリサイズする - if init_images is not None and args.W is not None and args.H is not None: - print(f"resize img2img source images to {args.W}*{args.H}") - init_images = resize_images(init_images, (args.W, args.H)) + if args.W is not None and args.H is not None: + if init_images is not None: + print(f"resize img2img source images to {args.W}*{args.H}") + init_images = resize_images(init_images, (args.W, args.H)) if mask_images is not None: print(f"resize img2img mask images to {args.W}*{args.H}") mask_images = resize_images(mask_images, (args.W, args.H)) + if networks and mask_images: + # mask を領域情報として流用する、現在は1枚だけ対応 + # TODO 複数のnetwork classの混在時の考慮 + print("use mask as region") + # import cv2 + # for i in range(3): + # cv2.imshow("msk", np.array(mask_images[0])[:,:,i]) + # cv2.waitKey() + # cv2.destroyAllWindows() + networks[0].__class__.set_regions(networks, np.array(mask_images[0])) + mask_images = None + prev_image = None # for VGG16 guided if args.guide_image_path is not None: print(f"load image for CLIP/VGG16/ControlNet guidance: {args.guide_image_path}") diff --git a/library/model_util.py b/library/model_util.py index 1a715c8f..d1020c05 100644 --- a/library/model_util.py +++ b/library/model_util.py @@ -4,7 +4,7 @@ import math import os import torch -from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig +from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig, logging from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNet2DConditionModel from safetensors.torch import load_file, save_file @@ -916,7 +916,11 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, dtype=None): info = text_model.load_state_dict(converted_text_encoder_checkpoint) else: converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v1(state_dict) + + logging.set_verbosity_error() # don't show annoying warning text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14") + logging.set_verbosity_warning() + info = text_model.load_state_dict(converted_text_encoder_checkpoint) print("loading text encoder:", info) diff --git a/library/train_util.py b/library/train_util.py index 75176e13..6af1abee 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -7,13 +7,13 @@ import re import shutil import time from typing import ( - Dict, - List, - NamedTuple, - Optional, - Sequence, - Tuple, - Union, + Dict, + List, + NamedTuple, + Optional, + Sequence, + Tuple, + Union, ) from accelerate import Accelerator import glob @@ -214,24 +214,24 @@ class AugHelper: def __init__(self): # prepare all possible augmentators color_aug_method = albu.OneOf([ - albu.HueSaturationValue(8, 0, 0, p=.5), - albu.RandomGamma((95, 105), p=.5), + albu.HueSaturationValue(8, 0, 0, p=.5), + albu.RandomGamma((95, 105), p=.5), ], p=.33) flip_aug_method = albu.HorizontalFlip(p=0.5) # key: (use_color_aug, use_flip_aug) self.augmentors = { - (True, True): albu.Compose([ - color_aug_method, - flip_aug_method, - ], p=1.), - (True, False): albu.Compose([ - color_aug_method, - ], p=1.), - (False, True): albu.Compose([ - flip_aug_method, - ], p=1.), - (False, False): None + (True, True): albu.Compose([ + color_aug_method, + flip_aug_method, + ], p=1.), + (True, False): albu.Compose([ + color_aug_method, + ], p=1.), + (False, True): albu.Compose([ + flip_aug_method, + ], p=1.), + (False, False): None } def get_augmentor(self, use_color_aug: bool, use_flip_aug: bool) -> Optional[albu.Compose]: @@ -260,7 +260,7 @@ class DreamBoothSubset(BaseSubset): assert image_dir is not None, "image_dir must be specified / image_dirは指定が必須です" super().__init__(image_dir, num_repeats, shuffle_caption, keep_tokens, color_aug, flip_aug, - face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate) + face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate) self.is_reg = is_reg self.class_tokens = class_tokens @@ -271,12 +271,13 @@ class DreamBoothSubset(BaseSubset): return NotImplemented return self.image_dir == other.image_dir + class FineTuningSubset(BaseSubset): def __init__(self, image_dir, metadata_file: str, num_repeats, shuffle_caption, keep_tokens, color_aug, flip_aug, face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate) -> None: assert metadata_file is not None, "metadata_file must be specified / metadata_fileは指定が必須です" super().__init__(image_dir, num_repeats, shuffle_caption, keep_tokens, color_aug, flip_aug, - face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate) + face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate) self.metadata_file = metadata_file @@ -285,6 +286,7 @@ class FineTuningSubset(BaseSubset): return NotImplemented return self.metadata_file == other.metadata_file + class BaseDataset(torch.utils.data.Dataset): def __init__(self, tokenizer: CLIPTokenizer, max_token_length: int, resolution: Optional[Tuple[int, int]], debug_dataset: bool) -> None: super().__init__() @@ -804,7 +806,7 @@ class DreamBoothDataset(BaseDataset): captions.append("") else: captions.append(subset.class_tokens if cap_for_img is None else cap_for_img) - + self.set_tag_frequency(os.path.basename(subset.image_dir), captions) # タグ頻度を記録 return img_paths, captions @@ -815,11 +817,13 @@ class DreamBoothDataset(BaseDataset): reg_infos: List[ImageInfo] = [] for subset in subsets: if subset.num_repeats < 1: - print(f"ignore subset with image_dir='{subset.image_dir}': num_repeats is less than 1 / num_repeatsが1を下回っているためサブセットを無視します: {subset.num_repeats}") + print( + f"ignore subset with image_dir='{subset.image_dir}': num_repeats is less than 1 / num_repeatsが1を下回っているためサブセットを無視します: {subset.num_repeats}") continue if subset in self.subsets: - print(f"ignore duplicated subset with image_dir='{subset.image_dir}': use the first one / 既にサブセットが登録されているため、重複した後発のサブセットを無視します") + print( + f"ignore duplicated subset with image_dir='{subset.image_dir}': use the first one / 既にサブセットが登録されているため、重複した後発のサブセットを無視します") continue img_paths, captions = load_dreambooth_dir(subset) @@ -881,11 +885,13 @@ class FineTuningDataset(BaseDataset): for subset in subsets: if subset.num_repeats < 1: - print(f"ignore subset with metadata_file='{subset.metadata_file}': num_repeats is less than 1 / num_repeatsが1を下回っているためサブセットを無視します: {subset.num_repeats}") + print( + f"ignore subset with metadata_file='{subset.metadata_file}': num_repeats is less than 1 / num_repeatsが1を下回っているためサブセットを無視します: {subset.num_repeats}") continue if subset in self.subsets: - print(f"ignore duplicated subset with metadata_file='{subset.metadata_file}': use the first one / 既にサブセットが登録されているため、重複した後発のサブセットを無視します") + print( + f"ignore duplicated subset with metadata_file='{subset.metadata_file}': use the first one / 既にサブセットが登録されているため、重複した後発のサブセットを無視します") continue # メタデータを読み込む @@ -918,7 +924,9 @@ class FineTuningDataset(BaseDataset): elif tags is not None and len(tags) > 0: caption = caption + ', ' + tags tags_list.append(tags) - assert caption is not None and len(caption) > 0, f"caption or tag is required / キャプションまたはタグは必須です:{abs_path}" + + if caption is None: + caption = "" image_info = ImageInfo(image_key, subset.num_repeats, caption, False, abs_path) image_info.image_size = img_md.get('train_resolution') @@ -937,7 +945,7 @@ class FineTuningDataset(BaseDataset): self.subsets.append(subset) # check existence of all npz files - use_npz_latents = all([not(subset.color_aug or subset.random_crop) for subset in self.subsets]) + use_npz_latents = all([not (subset.color_aug or subset.random_crop) for subset in self.subsets]) if use_npz_latents: flip_aug_in_subset = False npz_any = False @@ -2201,7 +2209,7 @@ def sample_images(accelerator, args: argparse.Namespace, epoch, steps, device, v if epoch is None or epoch % args.sample_every_n_epochs != 0: return else: - if steps % args.sample_every_n_steps != 0: + if steps % args.sample_every_n_steps != 0 or epoch is not None: # steps is not divisible or end of epoch return print(f"generating sample images at step / サンプル画像生成 ステップ: {steps}") @@ -2209,8 +2217,6 @@ def sample_images(accelerator, args: argparse.Namespace, epoch, steps, device, v print(f"No prompt file / プロンプトファイルがありません: {args.sample_prompts}") return - # ここでCUDAのキャッシュクリアとかしたほうがいいのか…… - org_vae_device = vae.device # CPUにいるはず vae.to(device) @@ -2346,7 +2352,9 @@ def sample_images(accelerator, args: argparse.Namespace, epoch, steps, device, v prompt = prompt.replace(prompt_replacement[0], prompt_replacement[1]) if negative_prompt is not None: negative_prompt = negative_prompt.replace(prompt_replacement[0], prompt_replacement[1]) - + + height = max(64, height - height % 8) # round to divisible by 8 + width = max(64, width - width % 8) # round to divisible by 8 image = pipeline(prompt, height, width, sample_steps, scale, negative_prompt).images[0] ts_str = time.strftime('%Y%m%d%H%M%S', time.localtime()) @@ -2356,6 +2364,10 @@ def sample_images(accelerator, args: argparse.Namespace, epoch, steps, device, v image.save(os.path.join(save_dir, img_filename)) + # clear pipeline and cache to reduce vram usage + del pipeline + torch.cuda.empty_cache() + torch.set_rng_state(rng_state) torch.cuda.set_rng_state(cuda_rng_state) vae.to(org_vae_device) diff --git a/networks/check_lora_weights.py b/networks/check_lora_weights.py index 4ee3f575..6bd9ccd9 100644 --- a/networks/check_lora_weights.py +++ b/networks/check_lora_weights.py @@ -21,7 +21,7 @@ def main(file): for key, value in values: value = value.to(torch.float32) - print(f"{key},{torch.mean(torch.abs(value))},{torch.min(torch.abs(value))}") + print(f"{key},{str(tuple(value.size())).replace(', ', '-')},{torch.mean(torch.abs(value))},{torch.min(torch.abs(value))}") if __name__ == '__main__': diff --git a/networks/extract_lora_from_models.py b/networks/extract_lora_from_models.py index 84d705cf..5d77b9e5 100644 --- a/networks/extract_lora_from_models.py +++ b/networks/extract_lora_from_models.py @@ -45,8 +45,13 @@ def svd(args): text_encoder_t, _, unet_t = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.model_tuned) # create LoRA network to extract weights: Use dim (rank) as alpha - lora_network_o = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_o, unet_o) - lora_network_t = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_t, unet_t) + if args.conv_dim is None: + kwargs = {} + else: + kwargs = {"conv_dim": args.conv_dim, "conv_alpha": args.conv_dim} + + lora_network_o = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_o, unet_o, **kwargs) + lora_network_t = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_t, unet_t, **kwargs) assert len(lora_network_o.text_encoder_loras) == len( lora_network_t.text_encoder_loras), f"model version is different (SD1.x vs SD2.x) / それぞれのモデルのバージョンが違います（SD1.xベースとSD2.xベース） " @@ -85,13 +90,27 @@ def svd(args): # make LoRA with svd print("calculating by svd") - rank = args.dim lora_weights = {} with torch.no_grad(): for lora_name, mat in tqdm(list(diffs.items())): + # if args.conv_dim is None, diffs do not include LoRAs for conv2d-3x3 conv2d = (len(mat.size()) == 4) + kernel_size = None if not conv2d else mat.size()[2:4] + conv2d_3x3 = conv2d and kernel_size != (1, 1) + + rank = args.dim if not conv2d_3x3 or args.conv_dim is None else args.conv_dim + out_dim, in_dim = mat.size()[0:2] + + if args.device: + mat = mat.to(args.device) + # print(mat.size(), mat.device, rank, in_dim, out_dim) + rank = min(rank, in_dim, out_dim) # LoRA rank cannot exceed the original dim + if conv2d: - mat = mat.squeeze() + if conv2d_3x3: + mat = mat.flatten(start_dim=1) + else: + mat = mat.squeeze() U, S, Vh = torch.linalg.svd(mat) @@ -108,6 +127,13 @@ def svd(args): U = U.clamp(low_val, hi_val) Vh = Vh.clamp(low_val, hi_val) + if conv2d: + U = U.reshape(out_dim, rank, 1, 1) + Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1]) + + U = U.to("cpu").contiguous() + Vh = Vh.to("cpu").contiguous() + lora_weights[lora_name] = (U, Vh) # make state dict for LoRA @@ -124,8 +150,8 @@ def svd(args): weights = lora_weights[lora_name][i] # print(key, i, weights.size(), lora_sd[key].size()) - if len(lora_sd[key].size()) == 4: - weights = weights.unsqueeze(2).unsqueeze(3) + # if len(lora_sd[key].size()) == 4: + # weights = weights.unsqueeze(2).unsqueeze(3) assert weights.size() == lora_sd[key].size(), f"size unmatch: {key}" lora_sd[key] = weights @@ -139,7 +165,7 @@ def svd(args): os.makedirs(dir_name, exist_ok=True) # minimum metadata - metadata = {"ss_network_dim": str(args.dim), "ss_network_alpha": str(args.dim)} + metadata = {"ss_network_module": "networks.lora", "ss_network_dim": str(args.dim), "ss_network_alpha": str(args.dim)} lora_network_o.save_weights(args.save_to, save_dtype, metadata) print(f"LoRA weights are saved to: {args.save_to}") @@ -158,6 +184,8 @@ if __name__ == '__main__': parser.add_argument("--save_to", type=str, default=None, help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors") parser.add_argument("--dim", type=int, default=4, help="dimension (rank) of LoRA (default 4) / LoRAの次元数（rank）（デフォルト4）") + parser.add_argument("--conv_dim", type=int, default=None, + help="dimension (rank) of LoRA for Conv2d-3x3 (default None, disabled) / LoRAのConv2d-3x3の次元数（rank）（デフォルトNone、適用なし）") parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う") args = parser.parse_args() diff --git a/networks/lora.py b/networks/lora.py index 24b107ba..c0181c02 100644 --- a/networks/lora.py +++ b/networks/lora.py @@ -6,6 +6,7 @@ import math import os from typing import List +import numpy as np import torch from library import train_util @@ -25,8 +26,16 @@ class LoRAModule(torch.nn.Module): if org_module.__class__.__name__ == 'Conv2d': in_dim = org_module.in_channels out_dim = org_module.out_channels - self.lora_down = torch.nn.Conv2d(in_dim, lora_dim, (1, 1), bias=False) - self.lora_up = torch.nn.Conv2d(lora_dim, out_dim, (1, 1), bias=False) + + self.lora_dim = min(self.lora_dim, in_dim, out_dim) + if self.lora_dim != lora_dim: + print(f"{lora_name} dim (rank) is changed to: {self.lora_dim}") + + kernel_size = org_module.kernel_size + stride = org_module.stride + padding = org_module.padding + self.lora_down = torch.nn.Conv2d(in_dim, self.lora_dim, kernel_size, stride, padding, bias=False) + self.lora_up = torch.nn.Conv2d(self.lora_dim, out_dim, (1, 1), (1, 1), bias=False) else: in_dim = org_module.in_features out_dim = org_module.out_features @@ -45,20 +54,98 @@ class LoRAModule(torch.nn.Module): self.multiplier = multiplier self.org_module = org_module # remove in applying + self.region = None + self.region_mask = None def apply_to(self): self.org_forward = self.org_module.forward self.org_module.forward = self.forward del self.org_module + def set_region(self, region): + self.region = region + self.region_mask = None + def forward(self, x): - return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale + if self.region is None: + return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale + + # regional LoRA FIXME same as additional-network extension + if x.size()[1] % 77 == 0: + # print(f"LoRA for context: {self.lora_name}") + self.region = None + return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale + + # calculate region mask first time + if self.region_mask is None: + if len(x.size()) == 4: + h, w = x.size()[2:4] + else: + seq_len = x.size()[1] + ratio = math.sqrt((self.region.size()[0] * self.region.size()[1]) / seq_len) + h = int(self.region.size()[0] / ratio + .5) + w = seq_len // h + + r = self.region.to(x.device) + if r.dtype == torch.bfloat16: + r = r.to(torch.float) + r = r.unsqueeze(0).unsqueeze(1) + # print(self.lora_name, self.region.size(), x.size(), r.size(), h, w) + r = torch.nn.functional.interpolate(r, (h, w), mode='bilinear') + r = r.to(x.dtype) + + if len(x.size()) == 3: + r = torch.reshape(r, (1, x.size()[1], -1)) + + self.region_mask = r + + return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale * self.region_mask def create_network(multiplier, network_dim, network_alpha, vae, text_encoder, unet, **kwargs): if network_dim is None: network_dim = 4 # default - network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim, alpha=network_alpha) + + # extract dim/alpha for conv2d, and block dim + conv_dim = kwargs.get('conv_dim', None) + conv_alpha = kwargs.get('conv_alpha', None) + if conv_dim is not None: + conv_dim = int(conv_dim) + if conv_alpha is None: + conv_alpha = 1.0 + else: + conv_alpha = float(conv_alpha) + + """ + block_dims = kwargs.get("block_dims") + block_alphas = None + + if block_dims is not None: + block_dims = [int(d) for d in block_dims.split(',')] + assert len(block_dims) == NUM_BLOCKS, f"Number of block dimensions is not same to {NUM_BLOCKS}" + block_alphas = kwargs.get("block_alphas") + if block_alphas is None: + block_alphas = [1] * len(block_dims) + else: + block_alphas = [int(a) for a in block_alphas(',')] + assert len(block_alphas) == NUM_BLOCKS, f"Number of block alphas is not same to {NUM_BLOCKS}" + + conv_block_dims = kwargs.get("conv_block_dims") + conv_block_alphas = None + + if conv_block_dims is not None: + conv_block_dims = [int(d) for d in conv_block_dims.split(',')] + assert len(conv_block_dims) == NUM_BLOCKS, f"Number of block dimensions is not same to {NUM_BLOCKS}" + conv_block_alphas = kwargs.get("conv_block_alphas") + if conv_block_alphas is None: + conv_block_alphas = [1] * len(conv_block_dims) + else: + conv_block_alphas = [int(a) for a in conv_block_alphas(',')] + assert len(conv_block_alphas) == NUM_BLOCKS, f"Number of block alphas is not same to {NUM_BLOCKS}" + """ + + network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim, + alpha=network_alpha, conv_lora_dim=conv_dim, conv_alpha=conv_alpha) return network @@ -69,45 +156,88 @@ def create_network_from_weights(multiplier, file, vae, text_encoder, unet, **kwa else: weights_sd = torch.load(file, map_location='cpu') - # get dim (rank) - network_alpha = None - network_dim = None + # get dim/alpha mapping + modules_dim = {} + modules_alpha = {} for key, value in weights_sd.items(): - if network_alpha is None and 'alpha' in key: - network_alpha = value - if network_dim is None and 'lora_down' in key and len(value.size()) == 2: - network_dim = value.size()[0] + if '.' not in key: + continue - if network_alpha is None: - network_alpha = network_dim + lora_name = key.split('.')[0] + if 'alpha' in key: + modules_alpha[lora_name] = value + elif 'lora_down' in key: + dim = value.size()[0] + modules_dim[lora_name] = dim + # print(lora_name, value.size(), dim) - network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim, alpha=network_alpha) + # support old LoRA without alpha + for key in modules_dim.keys(): + if key not in modules_alpha: + modules_alpha = modules_dim[key] + + network = LoRANetwork(text_encoder, unet, multiplier=multiplier, modules_dim=modules_dim, modules_alpha=modules_alpha) network.weights_sd = weights_sd return network class LoRANetwork(torch.nn.Module): - UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"] + # is it possible to apply conv_in and conv_out? + UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention", "ResnetBlock2D", "Downsample2D", "Upsample2D"] TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"] LORA_PREFIX_UNET = 'lora_unet' LORA_PREFIX_TEXT_ENCODER = 'lora_te' - def __init__(self, text_encoder, unet, multiplier=1.0, lora_dim=4, alpha=1) -> None: + def __init__(self, text_encoder, unet, multiplier=1.0, lora_dim=4, alpha=1, conv_lora_dim=None, conv_alpha=None, modules_dim=None, modules_alpha=None) -> None: super().__init__() self.multiplier = multiplier + self.lora_dim = lora_dim self.alpha = alpha + self.conv_lora_dim = conv_lora_dim + self.conv_alpha = conv_alpha + + if modules_dim is not None: + print(f"create LoRA network from weights") + else: + print(f"create LoRA network. base dim (rank): {lora_dim}, alpha: {alpha}") + + self.apply_to_conv2d_3x3 = self.conv_lora_dim is not None + if self.apply_to_conv2d_3x3: + if self.conv_alpha is None: + self.conv_alpha = self.alpha + print(f"apply LoRA to Conv2d with kernel size (3,3). dim (rank): {self.conv_lora_dim}, alpha: {self.conv_alpha}") # create module instances def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> List[LoRAModule]: loras = [] for name, module in root_module.named_modules(): if module.__class__.__name__ in target_replace_modules: + # TODO get block index here for child_name, child_module in module.named_modules(): - if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)): + is_linear = child_module.__class__.__name__ == "Linear" + is_conv2d = child_module.__class__.__name__ == "Conv2d" + is_conv2d_1x1 = is_conv2d and child_module.kernel_size == (1, 1) + if is_linear or is_conv2d: lora_name = prefix + '.' + name + '.' + child_name lora_name = lora_name.replace('.', '_') - lora = LoRAModule(lora_name, child_module, self.multiplier, self.lora_dim, self.alpha) + + if modules_dim is not None: + if lora_name not in modules_dim: + continue # no LoRA module in this weights file + dim = modules_dim[lora_name] + alpha = modules_alpha[lora_name] + else: + if is_linear or is_conv2d_1x1: + dim = self.lora_dim + alpha = self.alpha + elif self.apply_to_conv2d_3x3: + dim = self.conv_lora_dim + alpha = self.conv_alpha + else: + continue + + lora = LoRAModule(lora_name, child_module, self.multiplier, dim, alpha) loras.append(lora) return loras @@ -130,7 +260,7 @@ class LoRANetwork(torch.nn.Module): self.multiplier = multiplier for lora in self.text_encoder_loras + self.unet_loras: lora.multiplier = self.multiplier - + def load_weights(self, file): if os.path.splitext(file)[1] == '.safetensors': from safetensors.torch import load_file, safe_open @@ -240,3 +370,18 @@ class LoRANetwork(torch.nn.Module): save_file(state_dict, file, metadata) else: torch.save(state_dict, file) + + @staticmethod + def set_regions(networks, image): + image = image.astype(np.float32) / 255.0 + for i, network in enumerate(networks[:3]): + # NOTE: consider averaging overwrapping area + region = image[:, :, i] + if region.max() == 0: + continue + region = torch.tensor(region) + network.set_region(region) + + def set_region(self, region): + for lora in self.unet_loras: + lora.set_region(region) diff --git a/networks/merge_lora.py b/networks/merge_lora.py index 09aea7b2..09dee4de 100644 --- a/networks/merge_lora.py +++ b/networks/merge_lora.py @@ -48,7 +48,7 @@ def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype): for name, module in root_module.named_modules(): if module.__class__.__name__ in target_replace_modules: for child_name, child_module in module.named_modules(): - if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)): + if child_module.__class__.__name__ == "Linear" or child_module.__class__.__name__ == "Conv2d": lora_name = prefix + '.' + name + '.' + child_name lora_name = lora_name.replace('.', '_') name_to_module[lora_name] = child_module @@ -80,13 +80,19 @@ def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype): # W <- W + U * D weight = module.weight + # print(module_name, down_weight.size(), up_weight.size()) if len(weight.size()) == 2: # linear weight = weight + ratio * (up_weight @ down_weight) * scale - else: - # conv2d + elif down_weight.size()[2:4] == (1, 1): + # conv2d 1x1 weight = weight + ratio * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2) ).unsqueeze(2).unsqueeze(3) * scale + else: + # conv2d 3x3 + conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3) + # print(conved.size(), weight.size(), module.stride, module.padding) + weight = weight + ratio * conved * scale module.weight = torch.nn.Parameter(weight) @@ -123,7 +129,7 @@ def merge_lora_models(models, ratios, merge_dtype): alphas[lora_module_name] = alpha if lora_module_name not in base_alphas: base_alphas[lora_module_name] = alpha - + print(f"dim: {list(set(dims.values()))}, alpha: {list(set(alphas.values()))}") # merge @@ -145,7 +151,7 @@ def merge_lora_models(models, ratios, merge_dtype): merged_sd[key] = merged_sd[key] + lora_sd[key] * scale else: merged_sd[key] = lora_sd[key] * scale - + # set alpha to sd for lora_module_name, alpha in base_alphas.items(): key = lora_module_name + ".alpha" diff --git a/networks/svd_merge_lora.py b/networks/svd_merge_lora.py index c0448fcb..c8e39b80 100644 --- a/networks/svd_merge_lora.py +++ b/networks/svd_merge_lora.py @@ -35,7 +35,8 @@ def save_to_file(file_name, model, state_dict, dtype): torch.save(model, file_name) -def merge_lora_models(models, ratios, new_rank, device, merge_dtype): +def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dtype): + print(f"new rank: {new_rank}, new conv rank: {new_conv_rank}") merged_sd = {} for model, ratio in zip(models, ratios): print(f"loading: {model}") @@ -58,11 +59,12 @@ def merge_lora_models(models, ratios, new_rank, device, merge_dtype): in_dim = down_weight.size()[1] out_dim = up_weight.size()[0] conv2d = len(down_weight.size()) == 4 - print(lora_module_name, network_dim, alpha, in_dim, out_dim) + kernel_size = None if not conv2d else down_weight.size()[2:4] + # print(lora_module_name, network_dim, alpha, in_dim, out_dim, kernel_size) # make original weight if not exist if lora_module_name not in merged_sd: - weight = torch.zeros((out_dim, in_dim, 1, 1) if conv2d else (out_dim, in_dim), dtype=merge_dtype) + weight = torch.zeros((out_dim, in_dim, *kernel_size) if conv2d else (out_dim, in_dim), dtype=merge_dtype) if device: weight = weight.to(device) else: @@ -77,9 +79,12 @@ def merge_lora_models(models, ratios, new_rank, device, merge_dtype): scale = (alpha / network_dim) if not conv2d: # linear weight = weight + ratio * (up_weight @ down_weight) * scale - else: + elif kernel_size == (1, 1): weight = weight + ratio * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2) ).unsqueeze(2).unsqueeze(3) * scale + else: + conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3) + weight = weight + ratio * conved * scale merged_sd[lora_module_name] = weight @@ -89,16 +94,25 @@ def merge_lora_models(models, ratios, new_rank, device, merge_dtype): with torch.no_grad(): for lora_module_name, mat in tqdm(list(merged_sd.items())): conv2d = (len(mat.size()) == 4) + kernel_size = None if not conv2d else mat.size()[2:4] + conv2d_3x3 = conv2d and kernel_size != (1, 1) + out_dim, in_dim = mat.size()[0:2] + if conv2d: - mat = mat.squeeze() + if conv2d_3x3: + mat = mat.flatten(start_dim=1) + else: + mat = mat.squeeze() + + module_new_rank = new_conv_rank if conv2d_3x3 else new_rank U, S, Vh = torch.linalg.svd(mat) - U = U[:, :new_rank] - S = S[:new_rank] + U = U[:, :module_new_rank] + S = S[:module_new_rank] U = U @ torch.diag(S) - Vh = Vh[:new_rank, :] + Vh = Vh[:module_new_rank, :] dist = torch.cat([U.flatten(), Vh.flatten()]) hi_val = torch.quantile(dist, CLAMP_QUANTILE) @@ -107,16 +121,16 @@ def merge_lora_models(models, ratios, new_rank, device, merge_dtype): U = U.clamp(low_val, hi_val) Vh = Vh.clamp(low_val, hi_val) + if conv2d: + U = U.reshape(out_dim, module_new_rank, 1, 1) + Vh = Vh.reshape(module_new_rank, in_dim, kernel_size[0], kernel_size[1]) + up_weight = U down_weight = Vh - if conv2d: - up_weight = up_weight.unsqueeze(2).unsqueeze(3) - down_weight = down_weight.unsqueeze(2).unsqueeze(3) - merged_lora_sd[lora_module_name + '.lora_up.weight'] = up_weight.to("cpu").contiguous() merged_lora_sd[lora_module_name + '.lora_down.weight'] = down_weight.to("cpu").contiguous() - merged_lora_sd[lora_module_name + '.alpha'] = torch.tensor(new_rank) + merged_lora_sd[lora_module_name + '.alpha'] = torch.tensor(module_new_rank) return merged_lora_sd @@ -138,7 +152,8 @@ def merge(args): if save_dtype is None: save_dtype = merge_dtype - state_dict = merge_lora_models(args.models, args.ratios, args.new_rank, args.device, merge_dtype) + new_conv_rank = args.new_conv_rank if args.new_conv_rank is not None else args.new_rank + state_dict = merge_lora_models(args.models, args.ratios, args.new_rank, new_conv_rank, args.device, merge_dtype) print(f"saving model to: {args.save_to}") save_to_file(args.save_to, state_dict, state_dict, save_dtype) @@ -158,6 +173,8 @@ if __name__ == '__main__': help="ratios for each model / それぞれのLoRAモデルの比率") parser.add_argument("--new_rank", type=int, default=4, help="Specify rank of output LoRA / 出力するLoRAのrank (dim)") + parser.add_argument("--new_conv_rank", type=int, default=None, + help="Specify rank of output LoRA for Conv2d 3x3, None for same as new_rank / 出力するConv2D 3x3 LoRAのrank (dim)、Noneでnew_rankと同じ") parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う") args = parser.parse_args() diff --git a/train_README-ja.md b/train_README-ja.md index bf0d9f9c..479f9604 100644 --- a/train_README-ja.md +++ b/train_README-ja.md @@ -1,4 +1,8 @@ -当リポジトリではモデルのfine tuning、DreamBooth、およびLoRAとTextual Inversionの学習をサポートします。この文書ではそれらに共通する、学習データの準備方法やスクリプトオプションについて説明します。 +__ドキュメント更新中のため記述に誤りがあるかもしれません。__ + +# 学習について、共通編 + +当リポジトリではモデルのfine tuning、DreamBooth、およびLoRAとTextual Inversionの学習をサポートします。この文書ではそれらに共通する、学習データの準備方法やオプション等について説明します。 # 概要 @@ -8,15 +12,14 @@ 以下について説明します。 1. 学習データの準備について（設定ファイルを用いる新形式） -1. Aspect Ratio Bucketingについて +1. 学習で使われる用語のごく簡単な解説 1. 以前の指定形式（設定ファイルを用いずコマンドラインから指定） +1. 学習途中のサンプル画像生成 +1. 各スクリプトで共通の、よく使われるオプション 1. fine tuning 方式のメタデータ準備：キャプションニングなど 1.だけ実行すればとりあえず学習は可能です（学習については各スクリプトのドキュメントを参照）。2.以降は必要に応じて参照してください。 - # 学習データの準備について @@ -36,7 +39,7 @@ 1. fine tuning方式（正則化画像使用不可） - あらかじめキャプションをメタデータファイルにまとめます。タグとキャプションを分けて管理したり、学習を高速化するためlatentsを事前キャッシュしたりなどの機能をサポートします（いずれも別文書で説明しています）。 + あらかじめキャプションをメタデータファイルにまとめます。タグとキャプションを分けて管理したり、学習を高速化するためlatentsを事前キャッシュしたりなどの機能をサポートします（いずれも別文書で説明しています）。（fine tuning方式という名前ですが fine tuning 以外でも使えます。） 学習したいものと使用できる指定方法の組み合わせは以下の通りです。 @@ -124,7 +127,7 @@ batch_size = 4 # バッチサイズ num_repeats = 1 # 正則化画像の繰り返し回数、基本的には1でよい ``` -基本的には以下を場所のみ書き換えれば学習できます。 +基本的には以下の場所のみ書き換えれば学習できます。 1. 学習解像度 @@ -132,7 +135,7 @@ batch_size = 4 # バッチサイズ 1. バッチサイズ - 同時に何件のデータを学習するかを指定します。GPUのVRAMサイズ、学習解像度によって変わってきます。またfine tuning/DreamBooth/LoRA等でも変わってきますので、詳しくは各スクリプトの説明をご覧ください。 + 同時に何件のデータを学習するかを指定します。GPUのVRAMサイズ、学習解像度によって変わってきます。詳しくは後述します。またfine tuning/DreamBooth/LoRA等でも変わってきますので各スクリプトの説明もご覧ください。 1. フォルダ指定 @@ -248,7 +251,45 @@ batch_size = 4 # バッチサイズ それぞれのドキュメントを参考に学習を行ってください。 -# Aspect Ratio Bucketing について +# 学習で使われる用語のごく簡単な解説 + +細かいことは省略していますし私も完全には理解していないため、詳しくは各自お調べください。 + +## fine tuning（ファインチューニング） + +モデルを学習して微調整することを指します。使われ方によって意味が異なってきますが、狭義のfine tuningはStable Diffusionの場合、モデルを画像とキャプションで学習することです。DreamBoothは狭義のfine tuningのひとつの特殊なやり方と言えます。広義のfine tuningは、LoRAやTextual Inversion、Hypernetworksなどを含み、モデルを学習することすべてを含みます。 + +## ステップ + +ざっくりいうと学習データで1回計算すると1ステップです。「学習データのキャプションを今のモデルに流してみて、出てくる画像を学習データの画像と比較し、学習データに近づくようにモデルをわずかに変更する」のが1ステップです。 + +## バッチサイズ + +バッチサイズは1ステップで何件のデータをまとめて計算するかを指定する値です。まとめて計算するため速度は相対的に向上します。また一般的には精度も高くなるといわれています。 + +`バッチサイズ×ステップ数` が学習に使われるデータの件数になります。そのため、バッチサイズを増やした分だけステップ数を減らすとよいでしょう。 + +（ただし、たとえば「バッチサイズ1で1600ステップ」と「バッチサイズ4で400ステップ」は同じ結果にはなりません。同じ学習率の場合、一般的には後者のほうが学習不足になります。学習率を多少大きくするか（たとえば `2e-6` など）、ステップ数をたとえば500ステップにするなどして工夫してください。） + +バッチサイズを大きくするとその分だけGPUメモリを消費します。メモリが足りなくなるとエラーになりますし、エラーにならないギリギリでは学習速度が低下します。タスクマネージャーや `nvidia-smi` コマンドで使用メモリ量を確認しながら調整するとよいでしょう。 + +なお、バッチは「一塊のデータ」位の意味です。 + +## 学習率 + +ざっくりいうと1ステップごとにどのくらい変化させるかを表します。大きな値を指定するとそれだけ速く学習が進みますが、変化しすぎてモデルが壊れたり、最適な状態にまで至れない場合があります。小さい値を指定すると学習速度は遅くなり、また最適な状態にやはり至れない場合があります。 + +fine tuning、DreamBoooth、LoRAそれぞれで大きく異なり、また学習データや学習させたいモデル、バッチサイズやステップ数によっても変わってきます。一般的な値から初めて学習状態を見ながら増減してください。 + +デフォルトでは学習全体を通して学習率は固定です。スケジューラの指定で学習率をどう変化させるか決められますので、それらによっても結果は変わってきます。 + +## エポック（epoch） + +学習データが一通り学習されると（データが一周すると）1 epochです。繰り返し回数を指定した場合は、その繰り返し後のデータが一周すると1 epochです。 + +1 epochのステップ数は、基本的には `データ件数÷バッチサイズ` ですが、Aspect Ratio Bucketing を使うと微妙に増えます（異なるbucketのデータは同じバッチにできないため、ステップ数が増えます）。 + +## Aspect Ratio Bucketing Stable Diffusion のv1は512\*512で学習されていますが、それに加えて256\*1024や384\*640といった解像度でも学習します。これによりトリミングされる部分が減り、より正しくキャプションと画像の関係が学習されることが期待されます。 @@ -260,11 +301,15 @@ Stable Diffusion のv1は512\*512で学習されていますが、それに加 機械学習では入力サイズをすべて統一するのが一般的ですが、特に制約があるわけではなく、実際は同一のバッチ内で統一されていれば大丈夫です。NovelAIの言うbucketingは、あらかじめ教師データを、アスペクト比に応じた学習解像度ごとに分類しておくことを指しているようです。そしてバッチを各bucket内の画像で作成することで、バッチの画像サイズを統一します。 -# 以前のデータ指定方法 +# 以前の指定形式（設定ファイルを用いずコマンドラインから指定） -フォルダ名で繰り返し回数を指定する方法です。 +`.toml` ファイルを指定せずコマンドラインオプションで指定する方法です。DreamBooth class+identifier方式、DreamBooth キャプション方式、fine tuning方式があります。 -## step 1. 学習用画像の準備 +## DreamBooth、class+identifier方式 + +フォルダ名で繰り返し回数を指定します。また `train_data_dir` オプションと `reg_data_dir` オプションを用います。 + +### step 1. 学習用画像の準備 学習用画像を格納するフォルダを作成します。 __さらにその中に__ 、以下の名前でディレクトリを作成します。 @@ -294,15 +339,7 @@ classがひとつで対象が複数の場合、正則化画像フォルダはひ - reg_girls - 1_1girl -### DreamBoothでキャプションを使う - -学習用画像、正則化画像のフォルダに、画像と同じファイル名で、拡張子.caption（オプションで変えられます）のファイルを置くと、そのファイルからキャプションを読み込みプロンプトとして学習します。 - -※それらの画像の学習に、フォルダ名（identifier class）は使用されなくなります。 - -キャプションファイルの拡張子はデフォルトで.captionです。学習スクリプトの `--caption_extension` オプションで変更できます。`--shuffle_caption` オプションで学習時のキャプションについて、カンマ区切りの各部分をシャッフルしながら学習します。 - -## step 2. 正則化画像の準備 +### step 2. 正則化画像の準備 正則化画像を使う場合の手順です。 @@ -313,16 +350,288 @@ classがひとつで対象が複数の場合、正則化画像フォルダはひ ![image](https://user-images.githubusercontent.com/52813779/210770897-329758e5-3675-49f1-b345-c135f1725832.png) -## step 3. 学習の実行 +### step 3. 学習の実行 各学習スクリプトを実行します。 `--train_data_dir` オプションで前述の学習用データのフォルダを（__画像を含むフォルダではなく、その親フォルダ__）、`--reg_data_dir` オプションで正則化画像のフォルダ（__画像を含むフォルダではなく、その親フォルダ__）を指定してください。 - # メタデータファイルの作成 diff --git a/train_db_README-ja.md b/train_db_README-ja.md index 85ae35aa..0d0747bb 100644 --- a/train_db_README-ja.md +++ b/train_db_README-ja.md @@ -1,75 +1,104 @@ -DreamBoothのガイドです。LoRA等の追加ネットワークの学習にも同じ手順を使います。 +DreamBoothのガイドです。 + +[学習についての共通ドキュメント](./train_README-ja.md) もあわせてご覧ください。 # 概要 +DreamBoothとは、画像生成モデルに特定の主題を追加学習し、それを特定の識別子で生成する技術です。[論文はこちら](https://arxiv.org/abs/2208.12242)。 + +具体的には、Stable Diffusionのモデルにキャラや画風などを学ばせ、それを `shs` のような特定の単語で呼び出せる（生成画像に出現させる）ことができます。 + +スクリプトは[DiffusersのDreamBooth](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth)を元にしていますが、以下のような機能追加を行っています（いくつかの機能は元のスクリプト側もその後対応しています）。 + スクリプトの主な機能は以下の通りです。 -- 8bit Adam optimizerおよびlatentのキャッシュによる省メモリ化（ShivamShrirao氏版と同様）。 +- 8bit Adam optimizerおよびlatentのキャッシュによる省メモリ化（[Shivam Shrirao氏版](https://github.com/ShivamShrirao/diffusers/tree/main/examples/dreambooth)と同様）。 - xformersによる省メモリ化。 - 512x512だけではなく任意サイズでの学習。 - augmentationによる品質の向上。 - DreamBoothだけではなくText Encoder+U-Netのfine tuningに対応。 -- StableDiffusion形式でのモデルの読み書き。 +- Stable Diffusion形式でのモデルの読み書き。 - Aspect Ratio Bucketing。 - Stable Diffusion v2.0対応。 # 学習の手順 -## step 1. 環境整備 +あらかじめこのリポジトリのREADMEを参照し、環境整備を行ってください。 -このリポジトリのREADMEを参照してください。 +## データの準備 +[学習データの準備について](./train_README-ja.md) を参照してください。 -## step 2. identifierとclassを決める +## 学習の実行 -学ばせたい対象を結びつける単語identifierと、対象の属するclassを決めます。 - -（instanceなどいろいろな呼び方がありますが、とりあえず元の論文に合わせます。） - -以下ごく簡単に説明します（詳しくは調べてください）。 - -classは学習対象の一般的な種別です。たとえば特定の犬種を学ばせる場合には、classはdogになります。アニメキャラならモデルによりboyやgirl、1boyや1girlになるでしょう。 - -identifierは学習対象を識別して学習するためのものです。任意の単語で構いませんが、元論文によると「tokinizerで1トークンになる3文字以下でレアな単語」が良いとのことです。 - -identifierとclassを使い、たとえば「shs dog」などでモデルを学習することで、学習させたい対象をclassから識別して学習できます。 - -画像生成時には「shs dog」とすれば学ばせた犬種の画像が生成されます。 - -（identifierとして私が最近使っているものを参考までに挙げると、``shs sts scs cpc coc cic msm usu ici lvl cic dii muk ori hru rik koo yos wny`` などです。） - -## step 3. 学習用画像の準備 -学習用画像を格納するフォルダを作成します。 __さらにその中に__ 、以下の名前でディレクトリを作成します。 +スクリプトを実行します。最大限、メモリを節約したコマンドは以下のようになります（実際には1行で入力します）。それぞれの行を必要に応じて書き換えてください。12GB程度のVRAMで動作するようです。 ``` -<繰り返し回数>_ +accelerate launch --num_cpu_threads_per_process 1 train_db.py + --pretrained_model_name_or_path=<.ckptまたは.safetensordまたはDiffusers版モデルのディレクトリ> + --dataset_config=<データ準備で作成した.tomlファイル> + --output_dir=<学習したモデルの出力先フォルダ> + --output_name=<学習したモデル出力時のファイル名> + --save_model_as=safetensors + --prior_loss_weight=1.0 + --max_train_steps=1600 + --learning_rate=1e-6 + --optimizer_type="AdamW8bit" + --xformers + --mixed_precision="fp16" + --cache_latents + --gradient_checkpointing ``` -間の``_``を忘れないでください。 +`num_cpu_threads_per_process` には通常は1を指定するとよいようです。 -繰り返し回数は、正則化画像と枚数を合わせるために指定します（後述します）。 +`pretrained_model_name_or_path` に追加学習を行う元となるモデルを指定します。Stable Diffusionのcheckpointファイル（.ckptまたは.safetensors）、Diffusersのローカルディスクにあるモデルディレクトリ、DiffusersのモデルID（"stabilityai/stable-diffusion-2"など）が指定できます。 -たとえば「sls frog」というプロンプトで、データを20回繰り返す場合、「20_sls frog」となります。以下のようになります。 +`output_dir` に学習後のモデルを保存するフォルダを指定します。`output_name` にモデルのファイル名を拡張子を除いて指定します。`save_model_as` でsafetensors形式での保存を指定しています。 -![image](https://user-images.githubusercontent.com/52813779/210770636-1c851377-5936-4c15-90b7-8ac8ad6c2074.png) +`dataset_config` に `.toml` ファイルを指定します。ファイル内でのバッチサイズ指定は、当初はメモリ消費を抑えるために `1` としてください。 -## step 4. 正則化画像の準備 -正則化画像を使う場合の手順です。使わずに学習することもできます（正則化画像を使わないと区別ができなくなるので対象class全体が影響を受けます）。 +`prior_loss_weight` は正則化画像のlossの重みです。通常は1.0を指定します。 -正則化画像を格納するフォルダを作成します。 __さらにその中に__ ``<繰り返し回数>_`` という名前でディレクトリを作成します。 +学習させるステップ数 `max_train_steps` を1600とします。学習率 `learning_rate` はここでは1e-6を指定しています。 -たとえば「frog」というプロンプトで、データを繰り返さない（1回だけ）場合、以下のようになります。 +省メモリ化のため `mixed_precision="fp16"` を指定します（RTX30 シリーズ以降では `bf16` も指定できます。環境整備時にaccelerateに行った設定と合わせてください）。また `gradient_checkpointing` を指定します。 -![image](https://user-images.githubusercontent.com/52813779/210770897-329758e5-3675-49f1-b345-c135f1725832.png) +オプティマイザ（モデルを学習データにあうように最適化＝学習させるクラス）にメモリ消費の少ない 8bit AdamW を使うため、 `optimizer_type="AdamW8bit"` を指定します。 -繰り返し回数は「 __学習用画像の繰り返し回数×学習用画像の枚数≧正則化画像の繰り返し回数×正則化画像の枚数__ 」となるように指定してください。 +`xformers` オプションを指定し、xformersのCrossAttentionを用います。xformersをインストールしていない場合やエラーとなる場合（環境にもよりますが `mixed_precision="no"` の場合など）、代わりに `mem_eff_attn` オプションを指定すると省メモリ版CrossAttentionを使用します（速度は遅くなります）。 -（1 epochのデータ数が「学習用画像の繰り返し回数×学習用画像の枚数」となります。正則化画像の枚数がそれより多いと、余った部分の正則化画像は使用されません。） +省メモリ化のため `cache_latents` オプションを指定してVAEの出力をキャッシュします。 -## step 5. 学習の実行 -スクリプトを実行します。最大限、メモリを節約したコマンドは以下のようになります（実際には1行で入力します）。 +ある程度メモリがある場合は、`.toml` ファイルを編集してバッチサイズをたとえば `4` くらいに増やしてください（高速化と精度向上の可能性があります）。また `cache_latents` を外すことで augmentation が可能になります。 -※LoRA等の追加ネットワークを学習する場合のコマンドは ``train_db.py`` ではなく ``train_network.py`` となります。また追加でnetwork_\*オプションが必要となりますので、LoRAのガイドを参照してください。 +### よく使われるオプションについて + +以下の場合には [学習の共通ドキュメント](./train_README-ja.md) の「よく使われるオプション」を参照してください。 + +- Stable Diffusion 2.xまたはそこからの派生モデルを学習する +- clip skipを2以上を前提としたモデルを学習する +- 75トークンを超えたキャプションで学習する + +### DreamBoothでのステップ数について + +当スクリプトでは省メモリ化のため、ステップ当たりの学習回数が元のスクリプトの半分になっています（対象の画像と正則化画像を同一のバッチではなく別のバッチに分割して学習するため）。 + +元のDiffusers版やXavierXiao氏のStable Diffusion版とほぼ同じ学習を行うには、ステップ数を倍にしてください。 + +（学習画像と正則化画像をまとめてから shuffle するため厳密にはデータの順番が変わってしまいますが、学習には大きな影響はないと思います。） + +### DreamBoothでのバッチサイズについて + +モデル全体を学習するためLoRA等の学習に比べるとメモリ消費量は多くなります（fine tuningと同じ）。 + +### 学習率について + +Diffusers版では5e-6ですがStable Diffusion版は1e-6ですので、上のサンプルでは1e-6を指定しています。 + +### 以前の形式のデータセット指定をした場合のコマンドライン + +解像度やバッチサイズをオプションで指定します。コマンドラインの例は以下の通りです。 ``` accelerate launch --num_cpu_threads_per_process 1 train_db.py @@ -77,6 +106,7 @@ accelerate launch --num_cpu_threads_per_process 1 train_db.py --train_data_dir=<学習用データのディレクトリ> --reg_data_dir=<正則化画像のディレクトリ> --output_dir=<学習したモデルの出力先ディレクトリ> + --output_name=<学習したモデル出力時のファイル名> --prior_loss_weight=1.0 --resolution=512 --train_batch_size=1 @@ -89,43 +119,33 @@ accelerate launch --num_cpu_threads_per_process 1 train_db.py --gradient_checkpointing ``` -num_cpu_threads_per_processには通常は1を指定するとよいようです。 +## 学習したモデルで画像生成する -pretrained_model_name_or_pathに追加学習を行う元となるモデルを指定します。Stable Diffusionのcheckpointファイル（.ckptまたは.safetensors）、Diffusersのローカルディスクにあるモデルディレクトリ、DiffusersのモデルID（"stabilityai/stable-diffusion-2"など）が指定できます。学習後のモデルの保存形式はデフォルトでは元のモデルと同じになります（save_model_asオプションで変更できます）。 +学習が終わると指定したフォルダに指定した名前でsafetensorsファイルが出力されます。 -prior_loss_weightは正則化画像のlossの重みです。通常は1.0を指定します。 +v1.4/1.5およびその他の派生モデルの場合、このモデルでAutomatic1111氏のWebUIなどで推論できます。models\Stable-diffusionフォルダに置いてください。 -resolutionは画像のサイズ（解像度、幅と高さ）になります。bucketing（後述）を用いない場合、学習用画像、正則化画像はこのサイズとしてください。 +v2.xモデルでWebUIで画像生成する場合、モデルの仕様が記述された.yamlファイルが別途必要になります。v2.x baseの場合はv2-inference.yamlを、768/vの場合はv2-inference-v.yamlを、同じフォルダに置き、拡張子の前の部分をモデルと同じ名前にしてください。 -train_batch_sizeは学習時のバッチサイズです。max_train_stepsを1600とします。学習率learning_rateは、diffusers版では5e-6ですがStableDiffusion版は1e-6ですのでここでは1e-6を指定しています。 +![image](https://user-images.githubusercontent.com/52813779/210776915-061d79c3-6582-42c2-8884-8b91d2f07313.png) -省メモリ化のためmixed_precision="bf16"（または"fp16"）、およびgradient_checkpointing を指定します。 +各yamlファイルは[Stability AIのSD2.0のリポジトリ](https://github.com/Stability-AI/stablediffusion/tree/main/configs/stable-diffusion)にあります。 -xformersオプションを指定し、xformersのCrossAttentionを用います。xformersをインストールしていない場合、エラーとなる場合（mixed_precisionなしの場合、私の環境ではエラーとなりました）、代わりにmem_eff_attnオプションを指定すると省メモリ版CrossAttentionを使用します（速度は遅くなります）。 +# DreamBooth特有のその他の主なオプション -省メモリ化のためcache_latentsオプションを指定してVAEの出力をキャッシュします。 +すべてのオプションについては別文書を参照してください。 -ある程度メモリがある場合はたとえば以下のように指定します。 +## Text Encoderの学習を途中から行わない --stop_text_encoder_training -``` -accelerate launch --num_cpu_threads_per_process 8 train_db.py - --pretrained_model_name_or_path=<.ckptまたは.safetensordまたはDiffusers版モデルのディレクトリ> - --train_data_dir=<学習用データのディレクトリ> - --reg_data_dir=<正則化画像のディレクトリ> - --output_dir=<学習したモデルの出力先ディレクトリ> - --prior_loss_weight=1.0 - --resolution=512 - --train_batch_size=4 - --learning_rate=1e-6 - --max_train_steps=400 - --use_8bit_adam - --xformers - --mixed_precision="bf16" - --cache_latents -``` +stop_text_encoder_trainingオプションに数値を指定すると、そのステップ数以降はText Encoderの学習を行わずU-Netだけ学習します。場合によっては精度の向上が期待できるかもしれません。 -gradient_checkpointingを外し高速化します（メモリ使用量は増えます）。バッチサイズを増やし、高速化と精度向上を図ります。 +（恐らくText Encoderだけ先に過学習することがあり、それを防げるのではないかと推測していますが、詳細な影響は不明です。） +## Tokenizerのパディングをしない --no_token_padding +no_token_paddingオプションを指定するとTokenizerの出力をpaddingしません（Diffusers版の旧DreamBoothと同じ動きになります）。 + + + diff --git a/train_network.py b/train_network.py index 387590b3..cf64c894 100644 --- a/train_network.py +++ b/train_network.py @@ -427,10 +427,13 @@ def train(args): "ss_bucket_info": json.dumps(dataset.bucket_info), }) - # uncomment if another network is added - # for key, value in net_kwargs.items(): - # metadata["ss_arg_" + key] = value + # add extra args + if args.network_args: + metadata["ss_network_args"] = json.dumps(net_kwargs) + # for key, value in net_kwargs.items(): + # metadata["ss_arg_" + key] = value + # model name and hash if args.pretrained_model_name_or_path is not None: sd_model_name = args.pretrained_model_name_or_path if os.path.exists(sd_model_name): @@ -449,6 +452,13 @@ def train(args): metadata = {k: str(v) for k, v in metadata.items()} + # make minimum metadata for filtering + minimum_keys = ["ss_network_module", "ss_network_dim", "ss_network_alpha", "ss_network_args"] + minimum_metadata = {} + for key in minimum_keys: + if key in metadata: + minimum_metadata[key] = metadata[key] + progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps") global_step = 0 @@ -564,7 +574,7 @@ def train(args): ckpt_file = os.path.join(args.output_dir, ckpt_name) metadata["ss_training_finished_at"] = str(time.time()) print(f"saving checkpoint: {ckpt_file}") - unwrap_model(network).save_weights(ckpt_file, save_dtype, None if args.no_metadata else metadata) + unwrap_model(network).save_weights(ckpt_file, save_dtype, minimum_metadata if args.no_metadata else metadata) def remove_old_func(old_epoch_no): old_ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, old_epoch_no) + '.' + args.save_model_as @@ -603,7 +613,7 @@ def train(args): ckpt_file = os.path.join(args.output_dir, ckpt_name) print(f"save trained model to {ckpt_file}") - network.save_weights(ckpt_file, save_dtype, None if args.no_metadata else metadata) + network.save_weights(ckpt_file, save_dtype, minimum_metadata if args.no_metadata else metadata) print("model saved.") diff --git a/train_network_README-ja.md b/train_network_README-ja.md index 4a507ffc..4a79a6f7 100644 --- a/train_network_README-ja.md +++ b/train_network_README-ja.md @@ -1,118 +1,99 @@ -## LoRAの学習について +# LoRAの学習について [LoRA: Low-Rank Adaptation of Large Language Models](https://arxiv.org/abs/2106.09685)（arxiv）、[LoRA](https://github.com/microsoft/LoRA)（github）をStable Diffusionに適用したものです。 [cloneofsimo氏のリポジトリ](https://github.com/cloneofsimo/lora)を大いに参考にさせていただきました。ありがとうございます。 +通常のLoRAは Linear およぴカーネルサイズ 1x1 の Conv2d にのみ適用されますが、カーネルサイズ 3x3 のConv2dに適用を拡大することもできます。 + +Conv2d 3x3への拡大は [cloneofsimo氏](https://github.com/cloneofsimo/lora) が最初にリリースし、KohakuBlueleaf氏が [LoCon](https://github.com/KohakuBlueleaf/LoCon) でその有効性を明らかにしたものです。KohakuBlueleaf氏に深く感謝します。 + 8GB VRAMでもぎりぎり動作するようです。 +[学習についての共通ドキュメント](./train_README-ja.md) もあわせてご覧ください。 + ## 学習したモデルに関する注意 cloneofsimo氏のリポジトリ、およびd8ahazard氏の[Dreambooth Extension for Stable-Diffusion-WebUI](https://github.com/d8ahazard/sd_dreambooth_extension)とは、現時点では互換性がありません。いくつかの機能拡張を行っているためです（後述）。 WebUI等で画像生成する場合には、学習したLoRAのモデルを学習元のStable Diffusionのモデルにこのリポジトリ内のスクリプトであらかじめマージしておくか、こちらの[WebUI用extension](https://github.com/kohya-ss/sd-webui-additional-networks)を使ってください。 -## 学習方法 +# 学習の手順 -train_network.pyを用います。 +あらかじめこのリポジトリのREADMEを参照し、環境整備を行ってください。 -DreamBoothの手法（identifier（sksなど）とclass、オプションで正則化画像を用いる）と、キャプションを用いるfine tuningの手法の両方で学習できます。 +## データの準備 -どちらの方法も既存のスクリプトとほぼ同じ方法で学習できます。異なる点については後述します。 +[学習データの準備について](./train_README-ja.md) を参照してください。 -### DreamBoothの手法を用いる場合 -[DreamBoothのガイド](./train_db_README-ja.md) を参照してデータを用意してください。 +## 学習の実行 -学習するとき、train_db.pyの代わりにtrain_network.pyを指定してください。そして「LoRAの学習のためのオプション」にあるようにLoRA関連のオプション（``network_dim``や``network_alpha``など）を追加してください。 +`train_network.py`を用います。 -ほぼすべてのオプション（Stable Diffusionのモデル保存関係を除く）が使えますが、stop_text_encoder_trainingはサポートしていません。 - -### キャプションを用いる場合 - -[fine-tuningのガイド](./fine_tune_README_ja.md) を参照し、各手順を実行してください。 - -学習するとき、fine_tune.pyの代わりにtrain_network.pyを指定してください。ほぼすべてのオプション（モデル保存関係を除く）がそのまま使えます。そして「LoRAの学習のためのオプション」にあるようにLoRA関連のオプション（``network_dim``や``network_alpha``など）を追加してください。 - -なお「latentsの事前取得」は行わなくても動作します。VAEから学習時（またはキャッシュ時）にlatentを取得するため学習速度は遅くなりますが、代わりにcolor_augが使えるようになります。 - -### LoRAの学習のためのオプション - -train_network.pyでは--network_moduleオプションに、学習対象のモジュール名を指定します。LoRAに対応するのはnetwork.loraとなりますので、それを指定してください。 +`train_network.py`では `--network_module` オプションに、学習対象のモジュール名を指定します。LoRAに対応するのはnetwork.loraとなりますので、それを指定してください。 なお学習率は通常のDreamBoothやfine tuningよりも高めの、1e-4程度を指定するとよいようです。 -以下はコマンドラインの例です（DreamBooth手法）。 +以下はコマンドラインの例です。 ``` accelerate launch --num_cpu_threads_per_process 1 train_network.py - --pretrained_model_name_or_path=..\models\model.ckpt - --train_data_dir=..\data\db\char1 --output_dir=..\lora_train1 - --reg_data_dir=..\data\db\reg1 --prior_loss_weight=1.0 - --resolution=448,640 --train_batch_size=1 --learning_rate=1e-4 - --max_train_steps=400 --optimizer_type=AdamW8bit --xformers --mixed_precision=fp16 - --save_every_n_epochs=1 --save_model_as=safetensors --clip_skip=2 --seed=42 --color_aug + --pretrained_model_name_or_path=<.ckptまたは.safetensordまたはDiffusers版モデルのディレクトリ> + --dataset_config=<データ準備で作成した.tomlファイル> + --output_dir=<学習したモデルの出力先フォルダ> + --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 ``` -（2023/2/22:オプティマイザの指定方法が変わりました。[こちら](#オプティマイザの指定について）をご覧ください。） - ---output_dirオプションで指定したフォルダに、LoRAのモデルが保存されます。 +`--output_dir` オプションで指定したフォルダに、LoRAのモデルが保存されます。他のオプション、オプティマイザ等については [学習の共通ドキュメント](./train_README-ja.md) の「よく使われるオプション」も参照してください。 その他、以下のオプションが指定できます。 -* --network_dim +* `--network_dim` * LoRAのRANKを指定します（``--networkdim=4``など）。省略時は4になります。数が多いほど表現力は増しますが、学習に必要なメモリ、時間は増えます。また闇雲に増やしても良くないようです。 -* --network_alpha +* `--network_alpha` * アンダーフローを防ぎ安定して学習するための ``alpha`` 値を指定します。デフォルトは1です。``network_dim``と同じ値を指定すると以前のバージョンと同じ動作になります。 -* --network_weights +* `--network_weights` * 学習前に学習済みのLoRAの重みを読み込み、そこから追加で学習します。 -* --network_train_unet_only +* `--network_train_unet_only` * U-Netに関連するLoRAモジュールのみ有効とします。fine tuning的な学習で指定するとよいかもしれません。 -* --network_train_text_encoder_only +* `--network_train_text_encoder_only` * Text Encoderに関連するLoRAモジュールのみ有効とします。Textual Inversion的な効果が期待できるかもしれません。 -* --unet_lr +* `--unet_lr` * U-Netに関連するLoRAモジュールに、通常の学習率（--learning_rateオプションで指定）とは異なる学習率を使う時に指定します。 -* --text_encoder_lr +* `--text_encoder_lr` * Text Encoderに関連するLoRAモジュールに、通常の学習率（--learning_rateオプションで指定）とは異なる学習率を使う時に指定します。Text Encoderのほうを若干低めの学習率（5e-5など）にしたほうが良い、という話もあるようです。 +* `--network_args` + * 複数の引数を指定できます。後述します。 ---network_train_unet_onlyと--network_train_text_encoder_onlyの両方とも未指定時（デフォルト）はText EncoderとU-Netの両方のLoRAモジュールを有効にします。 +`--network_train_unet_only` と `--network_train_text_encoder_only` の両方とも未指定時（デフォルト）はText EncoderとU-Netの両方のLoRAモジュールを有効にします。 -## オプティマイザの指定について +## LoRA を Conv2d に拡大して適用する ---optimizer_type オプションでオプティマイザの種類を指定します。以下が指定できます。 +通常のLoRAは Linear およぴカーネルサイズ 1x1 の Conv2d にのみ適用されますが、カーネルサイズ 3x3 のConv2dに適用を拡大することもできます。 -- AdamW : [torch.optim.AdamW](https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html) - - 過去のバージョンのオプション未指定時と同じ -- AdamW8bit : 引数は同上 - - 過去のバージョンの--use_8bit_adam指定時と同じ -- Lion : https://github.com/lucidrains/lion-pytorch - - 過去のバージョンの--use_lion_optimizer指定時と同じ -- SGDNesterov : [torch.optim.SGD](https://pytorch.org/docs/stable/generated/torch.optim.SGD.html), nesterov=True -- SGDNesterov8bit : 引数は同上 -- DAdaptation : https://github.com/facebookresearch/dadaptation -- AdaFactor : [Transformers AdaFactor](https://huggingface.co/docs/transformers/main_classes/optimizer_schedules) -- 任意のオプティマイザ +`--network_args` に以下のように指定してください。`conv_dim` で Conv2d (3x3) の rank を、`conv_alpha` で alpha を指定してください。 -オプティマイザのオプション引数は--optimizer_argsオプションで指定してください。key=valueの形式で、複数の値が指定できます。また、valueはカンマ区切りで複数の値が指定できます。たとえばAdamWオプティマイザに引数を指定する場合は、``--optimizer_args weight_decay=0.01 betas=.9,.999``のようになります。 +``` +--network_args "conv_dim=1" "conv_alpha=1" +``` -オプション引数を指定する場合は、それぞれのオプティマイザの仕様をご確認ください。 +以下のように alpha 省略時は1になります。 -一部のオプティマイザでは必須の引数があり、省略すると自動的に追加されます（SGDNesterovのmomentumなど）。コンソールの出力を確認してください。 - -D-Adaptationオプティマイザは学習率を自動調整します。学習率のオプションに指定した値は学習率そのものではなくD-Adaptationが決定した学習率の適用率になりますので、通常は1.0を指定してください。Text EncoderにU-Netの半分の学習率を指定したい場合は、``--text_encoder_lr=0.5 --unet_lr=1.0``と指定します。 - -AdaFactorオプティマイザはrelative_step=Trueを指定すると学習率を自動調整できます（省略時はデフォルトで追加されます）。自動調整する場合は学習率のスケジューラにはadafactor_schedulerが強制的に使用されます。またscale_parameterとwarmup_initを指定するとよいようです。 - -自動調整する場合のオプション指定はたとえば ``--optimizer_args "relative_step=True" "scale_parameter=True" "warmup_init=True"`` のようになります。 - -学習率を自動調整しない場合はオプション引数 ``relative_step=False`` を追加してください。その場合、学習率のスケジューラにはconstant_with_warmupが、また勾配のclip normをしないことが推奨されているようです。そのため引数は ``--optimizer_type=adafactor --optimizer_args "relative_step=False" --lr_scheduler="constant_with_warmup" --max_grad_norm=0.0`` のようになります。 - -### 任意のオプティマイザを使う - -``torch.optim`` のオプティマイザを使う場合にはクラス名のみを（``--optimizer_type=RMSprop``など）、他のモジュールのオプティマイザを使う時は「モジュール名.クラス名」を指定してください（``--optimizer_type=bitsandbytes.optim.lamb.LAMB``など）。 - -（内部でimportlibしているだけで動作は未確認です。必要ならパッケージをインストールしてください。） +``` +--network_args "conv_dim=1" +``` ## マージスクリプトについて @@ -176,6 +157,27 @@ v1で学習したLoRAとv2で学習したLoRA、rank（次元数）や``alpha`` * save_precision * モデル保存時の精度をfloat、fp16、bf16から指定できます。省略時はprecisionと同じ精度になります。 + +## 複数のrankが異なるLoRAのモデルをマージする + +複数のLoRAをひとつのLoRAで近似します（完全な再現はできません）。`svd_merge_lora.py`を用います。たとえば以下のようなコマンドラインになります。 + +``` +python networks\svd_merge_lora.py + --save_to ..\lora_train1\model-char1-style1-merged.safetensors + --models ..\lora_train1\last.safetensors ..\lora_train2\last.safetensors + --ratios 0.6 0.4 --new_rank 32 --device cuda +``` + +`merge_lora.py` と主なオプションは同一です。以下のオプションが追加されています。 + +- `--new_rank` + - 作成するLoRAのrankを指定します。 +- `--new_conv_rank` + - 作成する Conv2d 3x3 LoRA の rank を指定します。省略時は `new_rank` と同じになります。 +- `--device` + - `--device cuda`としてcudaを指定すると計算をGPU上で行います。処理が速くなります。 + ## 当リポジトリ内の画像生成スクリプトで生成する gen_img_diffusers.pyに、--network_module、--network_weightsの各オプションを追加してください。意味は学習時と同様です。 @@ -209,12 +211,14 @@ Text Encoderが二つのモデルで同じ場合にはLoRAはU-NetのみのLoRA ### その他のオプション -- --v2 +- `--v2` - v2.xのStable Diffusionモデルを使う場合に指定してください。 -- --device +- `--device` - ``--device cuda``としてcudaを指定すると計算をGPU上で行います。処理が速くなります（CPUでもそこまで遅くないため、せいぜい倍～数倍程度のようです）。 -- --save_precision +- `--save_precision` - LoRAの保存形式を"float", "fp16", "bf16"から指定します。省略時はfloatになります。 +- `--conv_dim` + - 指定するとLoRAの適用範囲を Conv2d 3x3 へ拡大します。Conv2d 3x3 の rank を指定します。 ## 画像リサイズスクリプト @@ -252,7 +256,7 @@ python tools\resize_images_to_resolution.py --max_resolution 512x512,384x384,256 ### cloneofsimo氏のリポジトリとの違い -12/25時点では、当リポジトリはLoRAの適用個所をText EncoderのMLP、U-NetのFFN、Transformerのin/out projectionに拡大し、表現力が増しています。ただその代わりメモリ使用量は増え、8GBぎりぎりになりました。 +2022/12/25時点では、当リポジトリはLoRAの適用個所をText EncoderのMLP、U-NetのFFN、Transformerのin/out projectionに拡大し、表現力が増しています。ただその代わりメモリ使用量は増え、8GBぎりぎりになりました。 またモジュール入れ替え機構は全く異なります。 diff --git a/train_textual_inversion.py b/train_textual_inversion.py index d91a78ff..34b7f092 100644 --- a/train_textual_inversion.py +++ b/train_textual_inversion.py @@ -181,6 +181,11 @@ def train(args): for tmpl in templates: captions.append(tmpl.format(replace_to)) train_dataset_group.add_replacement("", captions) + + if args.num_vectors_per_token > 1: + prompt_replacement = (args.token_string, replace_to) + else: + prompt_replacement = None else: if args.num_vectors_per_token > 1: replace_to = " ".join(token_strings) diff --git a/train_ti_README-ja.md b/train_ti_README-ja.md index 90989ec0..90873696 100644 --- a/train_ti_README-ja.md +++ b/train_ti_README-ja.md @@ -1,32 +1,41 @@ -## Textual Inversionの学習について +[Textual Inversion](https://textual-inversion.github.io/) の学習についての説明です。 -[Textual Inversion](https://textual-inversion.github.io/)です。実装に当たっては https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion を大いに参考にしました。 +[学習についての共通ドキュメント](./train_README-ja.md) もあわせてご覧ください。 -学習したモデルはWeb UIでもそのまま使えます。 +実装に当たっては https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion を大いに参考にしました。 -なお恐らくSD2.xにも対応していますが現時点では未テストです。 +学習したモデルはWeb UIでもそのまま使えます。なお恐らくSD2.xにも対応していますが現時点では未テストです。 -## 学習方法 +# 学習の手順 -``train_textual_inversion.py`` を用います。 +あらかじめこのリポジトリのREADMEを参照し、環境整備を行ってください。 -データの準備については ``train_network.py`` と全く同じですので、[そちらのドキュメント](./train_network_README-ja.md)を参照してください。 +## データの準備 -## オプション +[学習データの準備について](./train_README-ja.md) を参照してください。 -以下はコマンドラインの例です（DreamBooth手法）。 +## 学習の実行 + +``train_textual_inversion.py`` を用います。以下はコマンドラインの例です（DreamBooth手法）。 ``` accelerate launch --num_cpu_threads_per_process 1 train_textual_inversion.py - --pretrained_model_name_or_path=..\models\model.ckpt - --train_data_dir=..\data\db\char1 --output_dir=..\ti_train1 - --resolution=448,640 --train_batch_size=1 --learning_rate=1e-4 - --max_train_steps=400 --use_8bit_adam --xformers --mixed_precision=fp16 - --save_every_n_epochs=1 --save_model_as=safetensors --clip_skip=2 --seed=42 --color_aug + --dataset_config=<データ準備で作成した.tomlファイル> + --output_dir=<学習したモデルの出力先フォルダ> + --output_name=<学習したモデル出力時のファイル名> + --save_model_as=safetensors + --prior_loss_weight=1.0 + --max_train_steps=1600 + --learning_rate=1e-6 + --optimizer_type="AdamW8bit" + --xformers + --mixed_precision="fp16" + --cache_latents + --gradient_checkpointing --token_string=mychar4 --init_word=cute --num_vectors_per_token=4 ``` -``--token_string`` に学習時のトークン文字列を指定します。__学習時のプロンプトは、この文字列を含むようにしてください（token_stringがmychar4なら、``mychar4 1girl`` など）__。プロンプトのこの文字列の部分が、Textual Inversionの新しいtokenに置換されて学習されます。 +``--token_string`` に学習時のトークン文字列を指定します。__学習時のプロンプトは、この文字列を含むようにしてください（token_stringがmychar4なら、``mychar4 1girl`` など）__。プロンプトのこの文字列の部分が、Textual Inversionの新しいtokenに置換されて学習されます。DreamBooth, class+identifier形式のデータセットとして、`token_string` をトークン文字列にするのが最も簡単で確実です。 プロンプトにトークン文字列が含まれているかどうかは、``--debug_dataset`` で置換後のtoken idが表示されますので、以下のように ``49408`` 以降のtokenが存在するかどうかで確認できます。 @@ -47,14 +56,47 @@ tokenizerがすでに持っている単語（一般的な単語）は使用で ``--num_vectors_per_token`` にいくつのトークンをこの学習で使うかを指定します。多いほうが表現力が増しますが、その分多くのトークンを消費します。たとえばnum_vectors_per_token=8の場合、指定したトークン文字列は（一般的なプロンプトの77トークン制限のうち）8トークンを消費します。 +以上がTextual Inversionのための主なオプションです。以降は他の学習スクリプトと同様です。 -その他、以下のオプションが指定できます。 +`num_cpu_threads_per_process` には通常は1を指定するとよいようです。 -* --weights +`pretrained_model_name_or_path` に追加学習を行う元となるモデルを指定します。Stable Diffusionのcheckpointファイル（.ckptまたは.safetensors）、Diffusersのローカルディスクにあるモデルディレクトリ、DiffusersのモデルID（"stabilityai/stable-diffusion-2"など）が指定できます。 + +`output_dir` に学習後のモデルを保存するフォルダを指定します。`output_name` にモデルのファイル名を拡張子を除いて指定します。`save_model_as` でsafetensors形式での保存を指定しています。 + +`dataset_config` に `.toml` ファイルを指定します。ファイル内でのバッチサイズ指定は、当初はメモリ消費を抑えるために `1` としてください。 + +学習させるステップ数 `max_train_steps` を10000とします。学習率 `learning_rate` はここでは5e-6を指定しています。 + +省メモリ化のため `mixed_precision="fp16"` を指定します（RTX30 シリーズ以降では `bf16` も指定できます。環境整備時にaccelerateに行った設定と合わせてください）。また `gradient_checkpointing` を指定します。 + +オプティマイザ（モデルを学習データにあうように最適化＝学習させるクラス）にメモリ消費の少ない 8bit AdamW を使うため、 `optimizer_type="AdamW8bit"` を指定します。 + +`xformers` オプションを指定し、xformersのCrossAttentionを用います。xformersをインストールしていない場合やエラーとなる場合（環境にもよりますが `mixed_precision="no"` の場合など）、代わりに `mem_eff_attn` オプションを指定すると省メモリ版CrossAttentionを使用します（速度は遅くなります）。 + +ある程度メモリがある場合は、`.toml` ファイルを編集してバッチサイズをたとえば `8` くらいに増やしてください（高速化と精度向上の可能性があります）。 + +### よく使われるオプションについて + +以下の場合にはオプションに関するドキュメントを参照してください。 + +- Stable Diffusion 2.xまたはそこからの派生モデルを学習する +- clip skipを2以上を前提としたモデルを学習する +- 75トークンを超えたキャプションで学習する + +### Textual Inversionでのバッチサイズについて + +モデル全体を学習するDreamBoothやfine tuningに比べてメモリ使用量が少ないため、バッチサイズは大きめにできます。 + +# Textual Inversionのその他の主なオプション + +すべてのオプションについては別文書を参照してください。 + +* `--weights` * 学習前に学習済みのembeddingsを読み込み、そこから追加で学習します。 -* --use_object_template +* `--use_object_template` * キャプションではなく既定の物体用テンプレート文字列（``a photo of a {}``など）で学習します。公式実装と同じになります。キャプションは無視されます。 -* --use_style_template +* `--use_style_template` * キャプションではなく既定のスタイル用テンプレート文字列で学習します（``a painting in the style of {}``など）。公式実装と同じになります。キャプションは無視されます。 ## 当リポジトリ内の画像生成スクリプトで生成する