Update to latest sd-script release
add gui support for sample config
This commit is contained in:
bmaltais
parent
fccb1c3359
commit
7249b0baa8
12
README-ja.md
12
README-ja.md
@ -38,18 +38,6 @@ PowerShellを使う場合、venvを使えるようにするためには以下の
|
||||
- 「Set-ExecutionPolicy Unrestricted」と入力し、Yと答えます。
|
||||
- 管理者のPowerShellを閉じます。
|
||||
|
||||
## Ubuntu環境でのインストール
|
||||
|
||||
```
|
||||
git clone https://github.com/kohya-ss/sd-scripts.git
|
||||
cd sd-scripts
|
||||
bash ubuntu_setup.sh
|
||||
```
|
||||
|
||||
をコマンドプロンプトで実行し、tkをインストールし、accelerateの質問をWindowsと同じように答えます。
|
||||
|
||||
`./gui.sh`でGUIを実行します。
|
||||
|
||||
## Windows環境でのインストール
|
||||
|
||||
以下の例ではPyTorchは1.12.1／CUDA 11.6版をインストールします。CUDA 11.3版やPyTorch 1.13を使う場合は適宜書き換えください。
|
||||
|
||||
16
README.md
16
README.md
@ -176,6 +176,22 @@ This will store your a backup file with your current locally installed pip packa
|
||||
|
||||
## Change History
|
||||
|
||||
* 2023/03/05 (v21.2.0):
|
||||
- There may be problems due to major changes. If you cannot revert back to a previous version when problems occur (`git checkout <release name>`).
|
||||
- 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](https://github-com.translate.goog/kohya-ss/sd-scripts/blob/main/config_README-ja.md) (only in Japanese currently.)
|
||||
- Specify `.toml` file with `--dataset_config` option.
|
||||
- The options supported under the previous release can be used as is instead of the `.toml` config file.
|
||||
- There might be bugs due to the large scale of update, please report any problems if you find at https://github.com/kohya-ss/sd-scripts/issues.
|
||||
- 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.)
|
||||
* 2023/03/05 (v21.1.5):
|
||||
- Add replace underscore with space option to WD14 captioning. Thanks @sALTaccount!
|
||||
- Improve how custom preset is set and handles.
|
||||
|
||||
279
config_README-ja.md
Normal file
279
config_README-ja.md
Normal file
@ -0,0 +1,279 @@
|
||||
For non-Japanese speakers: this README is provided only in Japanese in the current state. Sorry for inconvenience. We will provide English version in the near future.
|
||||
|
||||
`--dataset_config` で渡すことができる設定ファイルに関する説明です。
|
||||
|
||||
## 概要
|
||||
|
||||
設定ファイルを渡すことにより、ユーザが細かい設定を行えるようにします。
|
||||
|
||||
* 複数のデータセットが設定可能になります
|
||||
* 例えば `resolution` をデータセットごとに設定して、それらを混合して学習できます。
|
||||
* DreamBooth の手法と fine tuning の手法の両方に対応している学習方法では、DreamBooth 方式と fine tuning 方式のデータセットを混合することが可能です。
|
||||
* サブセットごとに設定を変更することが可能になります
|
||||
* データセットを画像ディレクトリ別またはメタデータ別に分割したものがサブセットです。いくつかのサブセットが集まってデータセットを構成します。
|
||||
* `keep_tokens` や `flip_aug` 等のオプションはサブセットごとに設定可能です。一方、`resolution` や `batch_size` といったオプションはデータセットごとに設定可能で、同じデータセットに属するサブセットでは値が共通になります。詳しくは後述します。
|
||||
|
||||
設定ファイルの形式は JSON か TOML を利用できます。記述のしやすさを考えると [TOML](https://toml.io/ja/v1.0.0-rc.2) を利用するのがオススメです。以下、TOML の利用を前提に説明します。
|
||||
|
||||
TOML で記述した設定ファイルの例です。
|
||||
|
||||
```toml
|
||||
[general]
|
||||
shuffle_caption = true
|
||||
caption_extension = '.txt'
|
||||
keep_tokens = 1
|
||||
|
||||
# これは DreamBooth 方式のデータセット
|
||||
[[datasets]]
|
||||
resolution = 512
|
||||
batch_size = 4
|
||||
keep_tokens = 2
|
||||
|
||||
[[datasets.subsets]]
|
||||
image_dir = 'C:\hoge'
|
||||
class_tokens = 'hoge girl'
|
||||
# このサブセットは keep_tokens = 2 （所属する datasets の値が使われる）
|
||||
|
||||
[[datasets.subsets]]
|
||||
image_dir = 'C:\fuga'
|
||||
class_tokens = 'fuga boy'
|
||||
keep_tokens = 3
|
||||
|
||||
[[datasets.subsets]]
|
||||
is_reg = true
|
||||
image_dir = 'C:\reg'
|
||||
class_tokens = 'human'
|
||||
keep_tokens = 1
|
||||
|
||||
# これは fine tuning 方式のデータセット
|
||||
[[datasets]]
|
||||
resolution = [768, 768]
|
||||
batch_size = 2
|
||||
|
||||
[[datasets.subsets]]
|
||||
image_dir = 'C:\piyo'
|
||||
metadata_file = 'C:\piyo\piyo_md.json'
|
||||
# このサブセットは keep_tokens = 1 （general の値が使われる）
|
||||
```
|
||||
|
||||
この例では、3 つのディレクトリを DreamBooth 方式のデータセットとして 512x512 (batch size 4) で学習させ、1 つのディレクトリを fine tuning 方式のデータセットとして 768x768 (batch size 2) で学習させることになります。
|
||||
|
||||
## データセット・サブセットに関する設定
|
||||
|
||||
データセット・サブセットに関する設定は、登録可能な箇所がいくつかに分かれています。
|
||||
|
||||
* `[general]`
|
||||
* 全データセットまたは全サブセットに適用されるオプションを指定する箇所です。
|
||||
* データセットごとの設定及びサブセットごとの設定に同名のオプションが存在していた場合には、データセット・サブセットごとの設定が優先されます。
|
||||
* `[[datasets]]`
|
||||
* `datasets` はデータセットに関する設定の登録箇所になります。各データセットに個別に適用されるオプションを指定する箇所です。
|
||||
* サブセットごとの設定が存在していた場合には、サブセットごとの設定が優先されます。
|
||||
* `[[datasets.subsets]]`
|
||||
* `datasets.subsets` はサブセットに関する設定の登録箇所になります。各サブセットに個別に適用されるオプションを指定する箇所です。
|
||||
|
||||
先程の例における、画像ディレクトリと登録箇所の対応に関するイメージ図です。
|
||||
|
||||
```
|
||||
C:\
|
||||
├─ hoge -> [[datasets.subsets]] No.1 ┐ ┐
|
||||
├─ fuga -> [[datasets.subsets]] No.2 |-> [[datasets]] No.1 |-> [general]
|
||||
├─ reg -> [[datasets.subsets]] No.3 ┘ |
|
||||
└─ piyo -> [[datasets.subsets]] No.4 --> [[datasets]] No.2 ┘
|
||||
```
|
||||
|
||||
画像ディレクトリがそれぞれ1つの `[[datasets.subsets]]` に対応しています。そして `[[datasets.subsets]]` が1つ以上組み合わさって1つの `[[datasets]]` を構成します。`[general]` には全ての `[[datasets]]`, `[[datasets.subsets]]` が属します。
|
||||
|
||||
登録箇所ごとに指定可能なオプションは異なりますが、同名のオプションが指定された場合は下位の登録箇所にある値が優先されます。先程の例の `keep_tokens` オプションの扱われ方を確認してもらうと理解しやすいかと思います。
|
||||
|
||||
加えて、学習方法が対応している手法によっても指定可能なオプションが変化します。
|
||||
|
||||
* DreamBooth 方式専用のオプション
|
||||
* fine tuning 方式専用のオプション
|
||||
* caption dropout の手法が使える場合のオプション
|
||||
|
||||
DreamBooth の手法と fine tuning の手法の両方とも利用可能な学習方法では、両者を併用することができます。
|
||||
併用する際の注意点として、DreamBooth 方式なのか fine tuning 方式なのかはデータセット単位で判別を行っているため、同じデータセット中に DreamBooth 方式のサブセットと fine tuning 方式のサブセットを混在させることはできません。
|
||||
つまり、これらを併用したい場合には異なる方式のサブセットが異なるデータセットに所属するように設定する必要があります。
|
||||
|
||||
プログラムの挙動としては、後述する `metadata_file` オプションが存在していたら fine tuning 方式のサブセットだと判断します。
|
||||
そのため、同一のデータセットに所属するサブセットについて言うと、「全てが `metadata_file` オプションを持つ」か「全てが `metadata_file` オプションを持たない」かのどちらかになっていれば問題ありません。
|
||||
|
||||
以下、利用可能なオプションを説明します。コマンドライン引数と名称が同一のオプションについては、基本的に説明を割愛します。他の README を参照してください。
|
||||
|
||||
### 全学習方法で共通のオプション
|
||||
|
||||
学習方法によらずに指定可能なオプションです。
|
||||
|
||||
#### データセット向けオプション
|
||||
|
||||
データセットの設定に関わるオプションです。`datasets.subsets` には記述できません。
|
||||
|
||||
| オプション名 | 設定例 | `[general]` | `[[datasets]]` |
|
||||
| ---- | ---- | ---- | ---- |
|
||||
| `batch_size` | `1` | o | o |
|
||||
| `bucket_no_upscale` | `true` | o | o |
|
||||
| `bucket_reso_steps` | `64` | o | o |
|
||||
| `enable_bucket` | `true` | o | o |
|
||||
| `max_bucket_reso` | `1024` | o | o |
|
||||
| `min_bucket_reso` | `128` | o | o |
|
||||
| `resolution` | `256`, `[512, 512]` | o | o |
|
||||
|
||||
* `batch_size`
|
||||
* コマンドライン引数の `--train_batch_size` と同等です。
|
||||
|
||||
これらの設定はデータセットごとに固定です。
|
||||
つまり、データセットに所属するサブセットはこれらの設定を共有することになります。
|
||||
例えば解像度が異なるデータセットを用意したい場合は、上に挙げた例のように別々のデータセットとして定義すれば別々の解像度を設定可能です。
|
||||
|
||||
#### サブセット向けオプション
|
||||
|
||||
サブセットの設定に関わるオプションです。
|
||||
|
||||
| オプション名 | 設定例 | `[general]` | `[[datasets]]` | `[[dataset.subsets]]` |
|
||||
| ---- | ---- | ---- | ---- | ---- |
|
||||
| `color_aug` | `false` | o | o | o |
|
||||
| `face_crop_aug_range` | `[1.0, 3.0]` | o | o | o |
|
||||
| `flip_aug` | `true` | o | o | o |
|
||||
| `keep_tokens` | `2` | o | o | o |
|
||||
| `num_repeats` | `10` | o | o | o |
|
||||
| `random_crop` | `false` | o | o | o |
|
||||
| `shuffle_caption` | `true` | o | o | o |
|
||||
|
||||
* `num_repeats`
|
||||
* サブセットの画像の繰り返し回数を指定します。fine tuning における `--dataset_repeats` に相当しますが、`num_repeats` はどの学習方法でも指定可能です。
|
||||
|
||||
### DreamBooth 方式専用のオプション
|
||||
|
||||
DreamBooth 方式のオプションは、サブセット向けオプションのみ存在します。
|
||||
|
||||
#### サブセット向けオプション
|
||||
|
||||
DreamBooth 方式のサブセットの設定に関わるオプションです。
|
||||
|
||||
| オプション名 | 設定例 | `[general]` | `[[datasets]]` | `[[dataset.subsets]]` |
|
||||
| ---- | ---- | ---- | ---- | ---- |
|
||||
| `image_dir` | `‘C:\hoge’` | - | - | o（必須） |
|
||||
| `caption_extension` | `".txt"` | o | o | o |
|
||||
| `class_tokens` | `“sks girl”` | - | - | o |
|
||||
| `is_reg` | `false` | - | - | o |
|
||||
|
||||
まず注意点として、 `image_dir` には画像ファイルが直下に置かれているパスを指定する必要があります。従来の DreamBooth の手法ではサブディレクトリに画像を置く必要がありましたが、そちらとは仕様に互換性がありません。また、`5_cat` のようなフォルダ名にしても、画像の繰り返し回数とクラス名は反映されません。これらを個別に設定したい場合、`num_repeats` と `class_tokens` で明示的に指定する必要があることに注意してください。
|
||||
|
||||
* `image_dir`
|
||||
* 画像ディレクトリのパスを指定します。指定必須オプションです。
|
||||
* 画像はディレクトリ直下に置かれている必要があります。
|
||||
* `class_tokens`
|
||||
* クラストークンを設定します。
|
||||
* 画像に対応する caption ファイルが存在しない場合にのみ学習時に利用されます。利用するかどうかの判定は画像ごとに行います。`class_tokens` を指定しなかった場合に caption ファイルも見つからなかった場合にはエラーになります。
|
||||
* `is_reg`
|
||||
* サブセットの画像が正規化用かどうかを指定します。指定しなかった場合は `false` として、つまり正規化画像ではないとして扱います。
|
||||
|
||||
### fine tuning 方式専用のオプション
|
||||
|
||||
fine tuning 方式のオプションは、サブセット向けオプションのみ存在します。
|
||||
|
||||
#### サブセット向けオプション
|
||||
|
||||
fine tuning 方式のサブセットの設定に関わるオプションです。
|
||||
|
||||
| オプション名 | 設定例 | `[general]` | `[[datasets]]` | `[[dataset.subsets]]` |
|
||||
| ---- | ---- | ---- | ---- | ---- |
|
||||
| `image_dir` | `‘C:\hoge’` | - | - | o |
|
||||
| `metadata_file` | `'C:\piyo\piyo_md.json'` | - | - | o（必須） |
|
||||
|
||||
* `image_dir`
|
||||
* 画像ディレクトリのパスを指定します。DreamBooth の手法の方とは異なり指定は必須ではありませんが、設定することを推奨します。
|
||||
* 指定する必要がない状況としては、メタデータファイルの生成時に `--full_path` を付与して実行していた場合です。
|
||||
* 画像はディレクトリ直下に置かれている必要があります。
|
||||
* `metadata_file`
|
||||
* サブセットで利用されるメタデータファイルのパスを指定します。指定必須オプションです。
|
||||
* コマンドライン引数の `--in_json` と同等です。
|
||||
* サブセットごとにメタデータファイルを指定する必要がある仕様上、ディレクトリを跨いだメタデータを1つのメタデータファイルとして作成することは避けた方が良いでしょう。画像ディレクトリごとにメタデータファイルを用意し、それらを別々のサブセットとして登録することを強く推奨します。
|
||||
|
||||
### caption dropout の手法が使える場合に指定可能なオプション
|
||||
|
||||
caption dropout の手法が使える場合のオプションは、サブセット向けオプションのみ存在します。
|
||||
DreamBooth 方式か fine tuning 方式かに関わらず、caption dropout に対応している学習方法であれば指定可能です。
|
||||
|
||||
#### サブセット向けオプション
|
||||
|
||||
caption dropout が使えるサブセットの設定に関わるオプションです。
|
||||
|
||||
| オプション名 | `[general]` | `[[datasets]]` | `[[dataset.subsets]]` |
|
||||
| ---- | ---- | ---- | ---- |
|
||||
| `caption_dropout_every_n_epochs` | o | o | o |
|
||||
| `caption_dropout_rate` | o | o | o |
|
||||
| `caption_tag_dropout_rate` | o | o | o |
|
||||
|
||||
## 重複したサブセットが存在する時の挙動
|
||||
|
||||
DreamBooth 方式のデータセットの場合、その中にある `image_dir` が同一のサブセットは重複していると見なされます。
|
||||
fine tuning 方式のデータセットの場合は、その中にある `metadata_file` が同一のサブセットは重複していると見なされます。
|
||||
データセット中に重複したサブセットが存在する場合、2個目以降は無視されます。
|
||||
|
||||
一方、異なるデータセットに所属している場合は、重複しているとは見なされません。
|
||||
例えば、以下のように同一の `image_dir` を持つサブセットを別々のデータセットに入れた場合には、重複していないと見なします。
|
||||
これは、同じ画像でも異なる解像度で学習したい場合に役立ちます。
|
||||
|
||||
```toml
|
||||
# 別々のデータセットに存在している場合は重複とは見なされず、両方とも学習に使われる
|
||||
|
||||
[[datasets]]
|
||||
resolution = 512
|
||||
|
||||
[[datasets.subsets]]
|
||||
image_dir = 'C:\hoge'
|
||||
|
||||
[[datasets]]
|
||||
resolution = 768
|
||||
|
||||
[[datasets.subsets]]
|
||||
image_dir = 'C:\hoge'
|
||||
```
|
||||
|
||||
## コマンドライン引数との併用
|
||||
|
||||
設定ファイルのオプションの中には、コマンドライン引数のオプションと役割が重複しているものがあります。
|
||||
|
||||
以下に挙げるコマンドライン引数のオプションは、設定ファイルを渡した場合には無視されます。
|
||||
|
||||
* `--train_data_dir`
|
||||
* `--reg_data_dir`
|
||||
* `--in_json`
|
||||
|
||||
以下に挙げるコマンドライン引数のオプションは、コマンドライン引数と設定ファイルで同時に指定された場合、コマンドライン引数の値よりも設定ファイルの値が優先されます。特に断りがなければ同名のオプションとなります。
|
||||
|
||||
| コマンドライン引数のオプション | 優先される設定ファイルのオプション |
|
||||
| ---------------------------------- | ---------------------------------- |
|
||||
| `--bucket_no_upscale` | |
|
||||
| `--bucket_reso_steps` | |
|
||||
| `--caption_dropout_every_n_epochs` | |
|
||||
| `--caption_dropout_rate` | |
|
||||
| `--caption_extension` | |
|
||||
| `--caption_tag_dropout_rate` | |
|
||||
| `--color_aug` | |
|
||||
| `--dataset_repeats` | `num_repeats` |
|
||||
| `--enable_bucket` | |
|
||||
| `--face_crop_aug_range` | |
|
||||
| `--flip_aug` | |
|
||||
| `--keep_tokens` | |
|
||||
| `--min_bucket_reso` | |
|
||||
| `--random_crop` | |
|
||||
| `--resolution` | |
|
||||
| `--shuffle_caption` | |
|
||||
| `--train_batch_size` | `batch_size` |
|
||||
|
||||
## エラーの手引き
|
||||
|
||||
現在、外部ライブラリを利用して設定ファイルの記述が正しいかどうかをチェックしているのですが、整備が行き届いておらずエラーメッセージがわかりづらいという問題があります。
|
||||
将来的にはこの問題の改善に取り組む予定です。
|
||||
|
||||
次善策として、頻出のエラーとその対処法について載せておきます。
|
||||
正しいはずなのにエラーが出る場合、エラー内容がどうしても分からない場合は、バグかもしれないのでご連絡ください。
|
||||
|
||||
* `voluptuous.error.MultipleInvalid: required key not provided @ ...`: 指定必須のオプションが指定されていないというエラーです。指定を忘れているか、オプション名を間違って記述している可能性が高いです。
|
||||
* `...` の箇所にはエラーが発生した場所が載っています。例えば `voluptuous.error.MultipleInvalid: required key not provided @ data['datasets'][0]['subsets'][0]['image_dir']` のようなエラーが出たら、0 番目の `datasets` 中の 0 番目の `subsets` の設定に `image_dir` が存在しないということになります。
|
||||
* `voluptuous.error.MultipleInvalid: expected int for dictionary value @ ...`: 指定する値の形式が不正というエラーです。値の形式が間違っている可能性が高いです。`int` の部分は対象となるオプションによって変わります。この README に載っているオプションの「設定例」が役立つかもしれません。
|
||||
* `voluptuous.error.MultipleInvalid: extra keys not allowed @ ...`: 対応していないオプション名が存在している場合に発生するエラーです。オプション名を間違って記述しているか、誤って紛れ込んでいる可能性が高いです。
|
||||
|
||||
|
||||
@ -36,6 +36,7 @@ from library.dreambooth_folder_creation_gui import (
|
||||
gradio_dreambooth_folder_creation_tab,
|
||||
)
|
||||
from library.utilities import utilities_tab
|
||||
from library.sampler_gui import sample_gradio_config, run_cmd_sample
|
||||
from easygui import msgbox
|
||||
|
||||
folder_symbol = '\U0001f4c2' # 📂
|
||||
@ -100,6 +101,10 @@ def save_configuration(
|
||||
optimizer,
|
||||
optimizer_args,
|
||||
noise_offset,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
):
|
||||
# Get list of function parameters and values
|
||||
parameters = list(locals().items())
|
||||
@ -201,6 +206,10 @@ def open_configuration(
|
||||
optimizer,
|
||||
optimizer_args,
|
||||
noise_offset,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
):
|
||||
# Get list of function parameters and values
|
||||
parameters = list(locals().items())
|
||||
@ -281,6 +290,10 @@ def train_model(
|
||||
optimizer,
|
||||
optimizer_args,
|
||||
noise_offset,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
):
|
||||
if pretrained_model_name_or_path == '':
|
||||
msgbox('Source model information is missing')
|
||||
@ -454,6 +467,13 @@ def train_model(
|
||||
noise_offset=noise_offset,
|
||||
)
|
||||
|
||||
run_cmd += run_cmd_sample(
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
)
|
||||
|
||||
print(run_cmd)
|
||||
|
||||
# Run the command
|
||||
@ -654,11 +674,14 @@ def dreambooth_tab(
|
||||
inputs=[color_aug],
|
||||
outputs=[cache_latents],
|
||||
)
|
||||
# optimizer.change(
|
||||
# set_legacy_8bitadam,
|
||||
# inputs=[optimizer, use_8bit_adam],
|
||||
# outputs=[optimizer, use_8bit_adam],
|
||||
# )
|
||||
|
||||
(
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
) = sample_gradio_config()
|
||||
|
||||
with gr.Tab('Tools'):
|
||||
gr.Markdown(
|
||||
'This section provide Dreambooth tools to help setup your dataset...'
|
||||
@ -740,6 +763,10 @@ def dreambooth_tab(
|
||||
optimizer,
|
||||
optimizer_args,
|
||||
noise_offset,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
]
|
||||
|
||||
button_open_config.click(
|
||||
|
||||
60
fine_tune.py
60
fine_tune.py
@ -13,7 +13,11 @@ import diffusers
|
||||
from diffusers import DDPMScheduler
|
||||
|
||||
import library.train_util as train_util
|
||||
|
||||
import library.config_util as config_util
|
||||
from library.config_util import (
|
||||
ConfigSanitizer,
|
||||
BlueprintGenerator,
|
||||
)
|
||||
|
||||
def collate_fn(examples):
|
||||
return examples[0]
|
||||
@ -30,25 +34,36 @@ def train(args):
|
||||
|
||||
tokenizer = train_util.load_tokenizer(args)
|
||||
|
||||
train_dataset = train_util.FineTuningDataset(args.in_json, args.train_batch_size, args.train_data_dir,
|
||||
tokenizer, args.max_token_length, args.shuffle_caption, args.keep_tokens,
|
||||
args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
|
||||
args.bucket_reso_steps, args.bucket_no_upscale,
|
||||
args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop,
|
||||
args.dataset_repeats, args.debug_dataset)
|
||||
blueprint_generator = BlueprintGenerator(ConfigSanitizer(False, True, True))
|
||||
if args.dataset_config is not None:
|
||||
print(f"Load dataset config from {args.dataset_config}")
|
||||
user_config = config_util.load_user_config(args.dataset_config)
|
||||
ignored = ["train_data_dir", "in_json"]
|
||||
if any(getattr(args, attr) is not None for attr in ignored):
|
||||
print("ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(', '.join(ignored)))
|
||||
else:
|
||||
user_config = {
|
||||
"datasets": [{
|
||||
"subsets": [{
|
||||
"image_dir": args.train_data_dir,
|
||||
"metadata_file": args.in_json,
|
||||
}]
|
||||
}]
|
||||
}
|
||||
|
||||
# 学習データのdropout率を設定する
|
||||
train_dataset.set_caption_dropout(args.caption_dropout_rate, args.caption_dropout_every_n_epochs, args.caption_tag_dropout_rate)
|
||||
|
||||
train_dataset.make_buckets()
|
||||
blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
|
||||
train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
|
||||
|
||||
if args.debug_dataset:
|
||||
train_util.debug_dataset(train_dataset)
|
||||
train_util.debug_dataset(train_dataset_group)
|
||||
return
|
||||
if len(train_dataset) == 0:
|
||||
if len(train_dataset_group) == 0:
|
||||
print("No data found. Please verify the metadata file and train_data_dir option. / 画像がありません。メタデータおよびtrain_data_dirオプションを確認してください。")
|
||||
return
|
||||
|
||||
if cache_latents:
|
||||
assert train_dataset_group.is_latent_cacheable(), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
|
||||
|
||||
# acceleratorを準備する
|
||||
print("prepare accelerator")
|
||||
accelerator, unwrap_model = train_util.prepare_accelerator(args)
|
||||
@ -109,7 +124,7 @@ def train(args):
|
||||
vae.requires_grad_(False)
|
||||
vae.eval()
|
||||
with torch.no_grad():
|
||||
train_dataset.cache_latents(vae)
|
||||
train_dataset_group.cache_latents(vae)
|
||||
vae.to("cpu")
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
@ -155,7 +170,7 @@ def train(args):
|
||||
# DataLoaderのプロセス数：0はメインプロセスになる
|
||||
n_workers = min(args.max_data_loader_n_workers, os.cpu_count() - 1) # cpu_count-1 ただし最大で指定された数まで
|
||||
train_dataloader = torch.utils.data.DataLoader(
|
||||
train_dataset, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
|
||||
train_dataset_group, batch_size=1, shuffle=True, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
|
||||
|
||||
# 学習ステップ数を計算する
|
||||
if args.max_train_epochs is not None:
|
||||
@ -199,7 +214,7 @@ def train(args):
|
||||
# 学習する
|
||||
total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
|
||||
print("running training / 学習開始")
|
||||
print(f" num examples / サンプル数: {train_dataset.num_train_images}")
|
||||
print(f" num examples / サンプル数: {train_dataset_group.num_train_images}")
|
||||
print(f" num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
|
||||
print(f" num epochs / epoch数: {num_train_epochs}")
|
||||
print(f" batch size per device / バッチサイズ: {args.train_batch_size}")
|
||||
@ -218,7 +233,7 @@ def train(args):
|
||||
|
||||
for epoch in range(num_train_epochs):
|
||||
print(f"epoch {epoch+1}/{num_train_epochs}")
|
||||
train_dataset.set_current_epoch(epoch + 1)
|
||||
train_dataset_group.set_current_epoch(epoch + 1)
|
||||
|
||||
for m in training_models:
|
||||
m.train()
|
||||
@ -282,17 +297,13 @@ def train(args):
|
||||
progress_bar.update(1)
|
||||
global_step += 1
|
||||
|
||||
train_util.sample_images(accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
|
||||
|
||||
current_loss = loss.detach().item() # 平均なのでbatch sizeは関係ないはず
|
||||
if args.logging_dir is not None:
|
||||
logs = {"loss": current_loss, "lr": float(lr_scheduler.get_last_lr()[0])}
|
||||
logs = {"avr_loss": loss_total / (step+1)}
|
||||
if args.optimizer_type.lower() == "DAdaptation".lower(): # tracking d*lr value
|
||||
# print(lr_scheduler.optimizers)
|
||||
logs["lr/d*lr"] = lr_scheduler.optimizers[0].param_groups[0]['d']*lr_scheduler.optimizers[0].param_groups[0]['lr']
|
||||
logs["d"] = lr_scheduler.optimizers[0].param_groups[0]['d']
|
||||
logs["lrD"] = lr_scheduler.optimizers[0].param_groups[0]['lr']
|
||||
logs["gsq_weighted"] = lr_scheduler.optimizers[0].param_groups[0]['gsq_weighted']
|
||||
|
||||
accelerator.log(logs, step=global_step)
|
||||
|
||||
# TODO moving averageにする
|
||||
@ -315,6 +326,8 @@ def train(args):
|
||||
train_util.save_sd_model_on_epoch_end(args, accelerator, src_path, save_stable_diffusion_format, use_safetensors,
|
||||
save_dtype, epoch, num_train_epochs, global_step, unwrap_model(text_encoder), unwrap_model(unet), vae)
|
||||
|
||||
train_util.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
|
||||
|
||||
is_main_process = accelerator.is_main_process
|
||||
if is_main_process:
|
||||
unet = unwrap_model(unet)
|
||||
@ -342,6 +355,7 @@ if __name__ == '__main__':
|
||||
train_util.add_training_arguments(parser, False)
|
||||
train_util.add_sd_saving_arguments(parser)
|
||||
train_util.add_optimizer_arguments(parser)
|
||||
config_util.add_config_arguments(parser)
|
||||
|
||||
parser.add_argument("--diffusers_xformers", action='store_true',
|
||||
help='use xformers by diffusers / Diffusersでxformersを使用する')
|
||||
|
||||
@ -324,7 +324,7 @@ __※引数を都度書き換えて、別のメタデータファイルに書き
|
||||
## 学習の実行
|
||||
たとえば以下のように実行します。以下は省メモリ化のための設定です。
|
||||
```
|
||||
accelerate launch --num_cpu_threads_per_process 8 fine_tune.py
|
||||
accelerate launch --num_cpu_threads_per_process 1 fine_tune.py
|
||||
--pretrained_model_name_or_path=model.ckpt
|
||||
--in_json meta_lat.json
|
||||
--train_data_dir=train_data
|
||||
@ -336,7 +336,7 @@ accelerate launch --num_cpu_threads_per_process 8 fine_tune.py
|
||||
--save_every_n_epochs=4
|
||||
```
|
||||
|
||||
accelerateのnum_cpu_threads_per_processにはCPUのコア数を指定するとよいようです。
|
||||
accelerateのnum_cpu_threads_per_processには通常は1を指定するとよいようです。
|
||||
|
||||
pretrained_model_name_or_pathに学習対象のモデルを指定します（Stable DiffusionのcheckpointかDiffusersのモデル）。Stable Diffusionのcheckpointは.ckptと.safetensorsに対応しています（拡張子で自動判定）。
|
||||
|
||||
|
||||
@ -27,6 +27,7 @@ from library.tensorboard_gui import (
|
||||
stop_tensorboard,
|
||||
)
|
||||
from library.utilities import utilities_tab
|
||||
from library.sampler_gui import sample_gradio_config, run_cmd_sample
|
||||
|
||||
folder_symbol = '\U0001f4c2' # 📂
|
||||
refresh_symbol = '\U0001f504' # 🔄
|
||||
@ -98,6 +99,10 @@ def save_configuration(
|
||||
optimizer,
|
||||
optimizer_args,
|
||||
noise_offset,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
):
|
||||
# Get list of function parameters and values
|
||||
parameters = list(locals().items())
|
||||
@ -205,6 +210,10 @@ def open_config_file(
|
||||
optimizer,
|
||||
optimizer_args,
|
||||
noise_offset,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
):
|
||||
# Get list of function parameters and values
|
||||
parameters = list(locals().items())
|
||||
@ -291,6 +300,10 @@ def train_model(
|
||||
optimizer,
|
||||
optimizer_args,
|
||||
noise_offset,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
):
|
||||
# create caption json file
|
||||
if generate_caption_database:
|
||||
@ -446,6 +459,13 @@ def train_model(
|
||||
noise_offset=noise_offset,
|
||||
)
|
||||
|
||||
run_cmd += run_cmd_sample(
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
)
|
||||
|
||||
print(run_cmd)
|
||||
|
||||
# Run the command
|
||||
@ -656,11 +676,13 @@ def finetune_tab():
|
||||
inputs=[color_aug],
|
||||
outputs=[cache_latents], # Not applicable to fine_tune.py
|
||||
)
|
||||
# optimizer.change(
|
||||
# set_legacy_8bitadam,
|
||||
# inputs=[optimizer, use_8bit_adam],
|
||||
# outputs=[optimizer, use_8bit_adam],
|
||||
# )
|
||||
|
||||
(
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
) = sample_gradio_config()
|
||||
|
||||
button_run = gr.Button('Train model', variant='primary')
|
||||
|
||||
@ -737,6 +759,10 @@ def finetune_tab():
|
||||
optimizer,
|
||||
optimizer_args,
|
||||
noise_offset,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
]
|
||||
|
||||
button_run.click(train_model, inputs=settings_list)
|
||||
|
||||
@ -80,6 +80,7 @@ from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
import library.model_util as model_util
|
||||
import library.train_util as train_util
|
||||
import tools.original_control_net as original_control_net
|
||||
from tools.original_control_net import ControlNetInfo
|
||||
|
||||
@ -589,6 +590,8 @@ class PipelineLike():
|
||||
latents: Optional[torch.FloatTensor] = None,
|
||||
max_embeddings_multiples: Optional[int] = 3,
|
||||
output_type: Optional[str] = "pil",
|
||||
vae_batch_size: float = None,
|
||||
return_latents: bool = False,
|
||||
# return_dict: bool = True,
|
||||
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
|
||||
is_cancelled_callback: Optional[Callable[[], bool]] = None,
|
||||
@ -680,6 +683,9 @@ class PipelineLike():
|
||||
else:
|
||||
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
|
||||
|
||||
vae_batch_size = batch_size if vae_batch_size is None else (
|
||||
int(vae_batch_size) if vae_batch_size >= 1 else max(1, int(batch_size * vae_batch_size)))
|
||||
|
||||
if strength < 0 or strength > 1:
|
||||
raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
|
||||
|
||||
@ -793,7 +799,6 @@ class PipelineLike():
|
||||
latents_dtype = text_embeddings.dtype
|
||||
init_latents_orig = None
|
||||
mask = None
|
||||
noise = None
|
||||
|
||||
if init_image is None:
|
||||
# get the initial random noise unless the user supplied it
|
||||
@ -825,6 +830,8 @@ class PipelineLike():
|
||||
if isinstance(init_image[0], PIL.Image.Image):
|
||||
init_image = [preprocess_image(im) for im in init_image]
|
||||
init_image = torch.cat(init_image)
|
||||
if isinstance(init_image, list):
|
||||
init_image = torch.stack(init_image)
|
||||
|
||||
# mask image to tensor
|
||||
if mask_image is not None:
|
||||
@ -835,9 +842,24 @@ class PipelineLike():
|
||||
|
||||
# encode the init image into latents and scale the latents
|
||||
init_image = init_image.to(device=self.device, dtype=latents_dtype)
|
||||
if init_image.size()[2:] == (height // 8, width // 8):
|
||||
init_latents = init_image
|
||||
else:
|
||||
if vae_batch_size >= batch_size:
|
||||
init_latent_dist = self.vae.encode(init_image).latent_dist
|
||||
init_latents = init_latent_dist.sample(generator=generator)
|
||||
else:
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
init_latents = []
|
||||
for i in tqdm(range(0, batch_size, vae_batch_size)):
|
||||
init_latent_dist = self.vae.encode(init_image[i:i + vae_batch_size]
|
||||
if vae_batch_size > 1 else init_image[i].unsqueeze(0)).latent_dist
|
||||
init_latents.append(init_latent_dist.sample(generator=generator))
|
||||
init_latents = torch.cat(init_latents)
|
||||
|
||||
init_latents = 0.18215 * init_latents
|
||||
|
||||
if len(init_latents) == 1:
|
||||
init_latents = init_latents.repeat((batch_size, 1, 1, 1))
|
||||
init_latents_orig = init_latents
|
||||
@ -932,8 +954,19 @@ class PipelineLike():
|
||||
if is_cancelled_callback is not None and is_cancelled_callback():
|
||||
return None
|
||||
|
||||
if return_latents:
|
||||
return (latents, False)
|
||||
|
||||
latents = 1 / 0.18215 * latents
|
||||
if vae_batch_size >= batch_size:
|
||||
image = self.vae.decode(latents).sample
|
||||
else:
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
images = []
|
||||
for i in tqdm(range(0, batch_size, vae_batch_size)):
|
||||
images.append(self.vae.decode(latents[i:i + vae_batch_size] if vae_batch_size > 1 else latents[i].unsqueeze(0)).sample)
|
||||
image = torch.cat(images)
|
||||
|
||||
image = (image / 2 + 0.5).clamp(0, 1)
|
||||
|
||||
@ -1820,7 +1853,7 @@ def preprocess_mask(mask):
|
||||
mask = mask.convert("L")
|
||||
w, h = mask.size
|
||||
w, h = map(lambda x: x - x % 32, (w, h)) # resize to integer multiple of 32
|
||||
mask = mask.resize((w // 8, h // 8), resample=PIL.Image.LANCZOS)
|
||||
mask = mask.resize((w // 8, h // 8), resample=PIL.Image.BILINEAR) # LANCZOS)
|
||||
mask = np.array(mask).astype(np.float32) / 255.0
|
||||
mask = np.tile(mask, (4, 1, 1))
|
||||
mask = mask[None].transpose(0, 1, 2, 3) # what does this step do?
|
||||
@ -1862,6 +1895,7 @@ class BatchDataExt(NamedTuple):
|
||||
|
||||
|
||||
class BatchData(NamedTuple):
|
||||
return_latents: bool
|
||||
base: BatchDataBase
|
||||
ext: BatchDataExt
|
||||
|
||||
@ -1930,10 +1964,7 @@ def main(args):
|
||||
# tokenizerを読み込む
|
||||
print("loading tokenizer")
|
||||
if use_stable_diffusion_format:
|
||||
if args.v2:
|
||||
tokenizer = CLIPTokenizer.from_pretrained(V2_STABLE_DIFFUSION_PATH, subfolder="tokenizer")
|
||||
else:
|
||||
tokenizer = CLIPTokenizer.from_pretrained(TOKENIZER_PATH) # , model_max_length=max_token_length + 2)
|
||||
tokenizer = train_util.load_tokenizer(args)
|
||||
|
||||
# schedulerを用意する
|
||||
sched_init_args = {}
|
||||
@ -2296,9 +2327,9 @@ def main(args):
|
||||
# highres_fixの処理
|
||||
if highres_fix and not highres_1st:
|
||||
# 1st stageのバッチを作成して呼び出す：サイズを小さくして呼び出す
|
||||
print("process 1st stage1")
|
||||
print("process 1st stage")
|
||||
batch_1st = []
|
||||
for base, ext in batch:
|
||||
for _, base, ext in batch:
|
||||
width_1st = int(ext.width * args.highres_fix_scale + .5)
|
||||
height_1st = int(ext.height * args.highres_fix_scale + .5)
|
||||
width_1st = width_1st - width_1st % 32
|
||||
@ -2306,20 +2337,29 @@ def main(args):
|
||||
|
||||
ext_1st = BatchDataExt(width_1st, height_1st, args.highres_fix_steps, ext.scale,
|
||||
ext.negative_scale, ext.strength, ext.network_muls)
|
||||
batch_1st.append(BatchData(base, ext_1st))
|
||||
batch_1st.append(BatchData(args.highres_fix_latents_upscaling, base, ext_1st))
|
||||
images_1st = process_batch(batch_1st, True, True)
|
||||
|
||||
# 2nd stageのバッチを作成して以下処理する
|
||||
print("process 2nd stage1")
|
||||
print("process 2nd stage")
|
||||
if args.highres_fix_latents_upscaling:
|
||||
org_dtype = images_1st.dtype
|
||||
if images_1st.dtype == torch.bfloat16:
|
||||
images_1st = images_1st.to(torch.float) # interpolateがbf16をサポートしていない
|
||||
images_1st = torch.nn.functional.interpolate(
|
||||
images_1st, (batch[0].ext.height // 8, batch[0].ext.width // 8), mode='bilinear') # , antialias=True)
|
||||
images_1st = images_1st.to(org_dtype)
|
||||
|
||||
batch_2nd = []
|
||||
for i, (bd, image) in enumerate(zip(batch, images_1st)):
|
||||
if not args.highres_fix_latents_upscaling:
|
||||
image = image.resize((bd.ext.width, bd.ext.height), resample=PIL.Image.LANCZOS) # img2imgとして設定
|
||||
bd_2nd = BatchData(BatchDataBase(*bd.base[0:3], bd.base.seed+1, image, None, *bd.base[6:]), bd.ext)
|
||||
bd_2nd = BatchData(False, BatchDataBase(*bd.base[0:3], bd.base.seed+1, image, None, *bd.base[6:]), bd.ext)
|
||||
batch_2nd.append(bd_2nd)
|
||||
batch = batch_2nd
|
||||
|
||||
# このバッチの情報を取り出す
|
||||
(step_first, _, _, _, init_image, mask_image, _, guide_image), \
|
||||
return_latents, (step_first, _, _, _, init_image, mask_image, _, guide_image), \
|
||||
(width, height, steps, scale, negative_scale, strength, network_muls) = batch[0]
|
||||
noise_shape = (LATENT_CHANNELS, height // DOWNSAMPLING_FACTOR, width // DOWNSAMPLING_FACTOR)
|
||||
|
||||
@ -2353,7 +2393,7 @@ def main(args):
|
||||
all_images_are_same = True
|
||||
all_masks_are_same = True
|
||||
all_guide_images_are_same = True
|
||||
for i, ((_, prompt, negative_prompt, seed, init_image, mask_image, clip_prompt, guide_image), _) in enumerate(batch):
|
||||
for i, (_, (_, prompt, negative_prompt, seed, init_image, mask_image, clip_prompt, guide_image), _) in enumerate(batch):
|
||||
prompts.append(prompt)
|
||||
negative_prompts.append(negative_prompt)
|
||||
seeds.append(seed)
|
||||
@ -2413,8 +2453,10 @@ def main(args):
|
||||
n.set_multiplier(m)
|
||||
|
||||
images = pipe(prompts, negative_prompts, init_images, mask_images, height, width, steps, scale, negative_scale, strength, latents=start_code,
|
||||
output_type='pil', max_embeddings_multiples=max_embeddings_multiples, img2img_noise=i2i_noises, clip_prompts=clip_prompts, clip_guide_images=guide_images)[0]
|
||||
if highres_1st and not args.highres_fix_save_1st:
|
||||
output_type='pil', max_embeddings_multiples=max_embeddings_multiples, img2img_noise=i2i_noises,
|
||||
vae_batch_size=args.vae_batch_size, return_latents=return_latents,
|
||||
clip_prompts=clip_prompts, clip_guide_images=guide_images)[0]
|
||||
if highres_1st and not args.highres_fix_save_1st: # return images or latents
|
||||
return images
|
||||
|
||||
# save image
|
||||
@ -2612,9 +2654,9 @@ def main(args):
|
||||
print("Use previous image as guide image.")
|
||||
guide_image = prev_image
|
||||
|
||||
b1 = BatchData(BatchDataBase(global_step, prompt, negative_prompt, seed, init_image, mask_image, clip_prompt, guide_image),
|
||||
b1 = BatchData(False, BatchDataBase(global_step, prompt, negative_prompt, seed, init_image, mask_image, clip_prompt, guide_image),
|
||||
BatchDataExt(width, height, steps, scale, negative_scale, strength, tuple(network_muls) if network_muls else None))
|
||||
if len(batch_data) > 0 and batch_data[-1][1] != b1[1]: # バッチ分割必要？
|
||||
if len(batch_data) > 0 and batch_data[-1].ext != b1.ext: # バッチ分割必要？
|
||||
process_batch(batch_data, highres_fix)
|
||||
batch_data.clear()
|
||||
|
||||
@ -2658,6 +2700,8 @@ if __name__ == '__main__':
|
||||
parser.add_argument("--H", type=int, default=None, help="image height, in pixel space / 生成画像高さ")
|
||||
parser.add_argument("--W", type=int, default=None, help="image width, in pixel space / 生成画像幅")
|
||||
parser.add_argument("--batch_size", type=int, default=1, help="batch size / バッチサイズ")
|
||||
parser.add_argument("--vae_batch_size", type=float, default=None,
|
||||
help="batch size for VAE, < 1.0 for ratio / VAE処理時のバッチサイズ、1未満の値の場合は通常バッチサイズの比率")
|
||||
parser.add_argument("--steps", type=int, default=50, help="number of ddim sampling steps / サンプリングステップ数")
|
||||
parser.add_argument('--sampler', type=str, default='ddim',
|
||||
choices=['ddim', 'pndm', 'lms', 'euler', 'euler_a', 'heun', 'dpm_2', 'dpm_2_a', 'dpmsolver',
|
||||
@ -2669,6 +2713,8 @@ if __name__ == '__main__':
|
||||
parser.add_argument("--ckpt", type=str, default=None, help="path to checkpoint of model / モデルのcheckpointファイルまたはディレクトリ")
|
||||
parser.add_argument("--vae", type=str, default=None,
|
||||
help="path to checkpoint of vae to replace / VAEを入れ替える場合、VAEのcheckpointファイルまたはディレクトリ")
|
||||
parser.add_argument("--tokenizer_cache_dir", type=str, default=None,
|
||||
help="directory for caching Tokenizer (for offline training) / Tokenizerをキャッシュするディレクトリ（ネット接続なしでの学習のため）")
|
||||
# parser.add_argument("--replace_clip_l14_336", action='store_true',
|
||||
# help="Replace CLIP (Text Encoder) to l/14@336 / CLIP(Text Encoder)をl/14@336に入れ替える")
|
||||
parser.add_argument("--seed", type=int, default=None,
|
||||
@ -2713,6 +2759,8 @@ if __name__ == '__main__':
|
||||
help="1st stage steps for highres fix / highres fixの最初のステージのステップ数")
|
||||
parser.add_argument("--highres_fix_save_1st", action='store_true',
|
||||
help="save 1st stage images for highres fix / highres fixの最初のステージの画像を保存する")
|
||||
parser.add_argument("--highres_fix_latents_upscaling", action='store_true',
|
||||
help="use latents upscaling for highres fix / highres fixでlatentで拡大する")
|
||||
parser.add_argument("--negative_scale", type=float, default=None,
|
||||
help="set another guidance scale for negative prompt / ネガティブプロンプトのscaleを指定する")
|
||||
|
||||
|
||||
@ -70,7 +70,10 @@ if __name__ == '__main__':
|
||||
# torch.cuda.set_per_process_memory_fraction(0.48)
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
'--listen', type=str, default='127.0.0.1', help='IP to listen on for connections to Gradio'
|
||||
'--listen',
|
||||
type=str,
|
||||
default='127.0.0.1',
|
||||
help='IP to listen on for connections to Gradio',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--username', type=str, default='', help='Username for authentication'
|
||||
|
||||
527
library/config_util.py
Normal file
527
library/config_util.py
Normal file
@ -0,0 +1,527 @@
|
||||
import argparse
|
||||
from dataclasses import (
|
||||
asdict,
|
||||
dataclass,
|
||||
)
|
||||
import functools
|
||||
from textwrap import dedent, indent
|
||||
import json
|
||||
from pathlib import Path
|
||||
# from toolz import curry
|
||||
from typing import (
|
||||
List,
|
||||
Optional,
|
||||
Sequence,
|
||||
Tuple,
|
||||
Union,
|
||||
)
|
||||
|
||||
import toml
|
||||
import voluptuous
|
||||
from voluptuous import (
|
||||
Any,
|
||||
ExactSequence,
|
||||
MultipleInvalid,
|
||||
Object,
|
||||
Required,
|
||||
Schema,
|
||||
)
|
||||
from transformers import CLIPTokenizer
|
||||
|
||||
from . import train_util
|
||||
from .train_util import (
|
||||
DreamBoothSubset,
|
||||
FineTuningSubset,
|
||||
DreamBoothDataset,
|
||||
FineTuningDataset,
|
||||
DatasetGroup,
|
||||
)
|
||||
|
||||
|
||||
def add_config_arguments(parser: argparse.ArgumentParser):
|
||||
parser.add_argument("--dataset_config", type=Path, default=None, help="config file for detail settings / 詳細な設定用の設定ファイル")
|
||||
|
||||
# TODO: inherit Params class in Subset, Dataset
|
||||
|
||||
@dataclass
|
||||
class BaseSubsetParams:
|
||||
image_dir: Optional[str] = None
|
||||
num_repeats: int = 1
|
||||
shuffle_caption: bool = False
|
||||
keep_tokens: int = 0
|
||||
color_aug: bool = False
|
||||
flip_aug: bool = False
|
||||
face_crop_aug_range: Optional[Tuple[float, float]] = None
|
||||
random_crop: bool = False
|
||||
caption_dropout_rate: float = 0.0
|
||||
caption_dropout_every_n_epochs: int = 0
|
||||
caption_tag_dropout_rate: float = 0.0
|
||||
|
||||
@dataclass
|
||||
class DreamBoothSubsetParams(BaseSubsetParams):
|
||||
is_reg: bool = False
|
||||
class_tokens: Optional[str] = None
|
||||
caption_extension: str = ".caption"
|
||||
|
||||
@dataclass
|
||||
class FineTuningSubsetParams(BaseSubsetParams):
|
||||
metadata_file: Optional[str] = None
|
||||
|
||||
@dataclass
|
||||
class BaseDatasetParams:
|
||||
tokenizer: CLIPTokenizer = None
|
||||
max_token_length: int = None
|
||||
resolution: Optional[Tuple[int, int]] = None
|
||||
debug_dataset: bool = False
|
||||
|
||||
@dataclass
|
||||
class DreamBoothDatasetParams(BaseDatasetParams):
|
||||
batch_size: int = 1
|
||||
enable_bucket: bool = False
|
||||
min_bucket_reso: int = 256
|
||||
max_bucket_reso: int = 1024
|
||||
bucket_reso_steps: int = 64
|
||||
bucket_no_upscale: bool = False
|
||||
prior_loss_weight: float = 1.0
|
||||
|
||||
@dataclass
|
||||
class FineTuningDatasetParams(BaseDatasetParams):
|
||||
batch_size: int = 1
|
||||
enable_bucket: bool = False
|
||||
min_bucket_reso: int = 256
|
||||
max_bucket_reso: int = 1024
|
||||
bucket_reso_steps: int = 64
|
||||
bucket_no_upscale: bool = False
|
||||
|
||||
@dataclass
|
||||
class SubsetBlueprint:
|
||||
params: Union[DreamBoothSubsetParams, FineTuningSubsetParams]
|
||||
|
||||
@dataclass
|
||||
class DatasetBlueprint:
|
||||
is_dreambooth: bool
|
||||
params: Union[DreamBoothDatasetParams, FineTuningDatasetParams]
|
||||
subsets: Sequence[SubsetBlueprint]
|
||||
|
||||
@dataclass
|
||||
class DatasetGroupBlueprint:
|
||||
datasets: Sequence[DatasetBlueprint]
|
||||
@dataclass
|
||||
class Blueprint:
|
||||
dataset_group: DatasetGroupBlueprint
|
||||
|
||||
|
||||
class ConfigSanitizer:
|
||||
# @curry
|
||||
@staticmethod
|
||||
def __validate_and_convert_twodim(klass, value: Sequence) -> Tuple:
|
||||
Schema(ExactSequence([klass, klass]))(value)
|
||||
return tuple(value)
|
||||
|
||||
# @curry
|
||||
@staticmethod
|
||||
def __validate_and_convert_scalar_or_twodim(klass, value: Union[float, Sequence]) -> Tuple:
|
||||
Schema(Any(klass, ExactSequence([klass, klass])))(value)
|
||||
try:
|
||||
Schema(klass)(value)
|
||||
return (value, value)
|
||||
except:
|
||||
return ConfigSanitizer.__validate_and_convert_twodim(klass, value)
|
||||
|
||||
# subset schema
|
||||
SUBSET_ASCENDABLE_SCHEMA = {
|
||||
"color_aug": bool,
|
||||
"face_crop_aug_range": functools.partial(__validate_and_convert_twodim.__func__, float),
|
||||
"flip_aug": bool,
|
||||
"num_repeats": int,
|
||||
"random_crop": bool,
|
||||
"shuffle_caption": bool,
|
||||
"keep_tokens": int,
|
||||
}
|
||||
# DO means DropOut
|
||||
DO_SUBSET_ASCENDABLE_SCHEMA = {
|
||||
"caption_dropout_every_n_epochs": int,
|
||||
"caption_dropout_rate": Any(float, int),
|
||||
"caption_tag_dropout_rate": Any(float, int),
|
||||
}
|
||||
# DB means DreamBooth
|
||||
DB_SUBSET_ASCENDABLE_SCHEMA = {
|
||||
"caption_extension": str,
|
||||
"class_tokens": str,
|
||||
}
|
||||
DB_SUBSET_DISTINCT_SCHEMA = {
|
||||
Required("image_dir"): str,
|
||||
"is_reg": bool,
|
||||
}
|
||||
# FT means FineTuning
|
||||
FT_SUBSET_DISTINCT_SCHEMA = {
|
||||
Required("metadata_file"): str,
|
||||
"image_dir": str,
|
||||
}
|
||||
|
||||
# datasets schema
|
||||
DATASET_ASCENDABLE_SCHEMA = {
|
||||
"batch_size": int,
|
||||
"bucket_no_upscale": bool,
|
||||
"bucket_reso_steps": int,
|
||||
"enable_bucket": bool,
|
||||
"max_bucket_reso": int,
|
||||
"min_bucket_reso": int,
|
||||
"resolution": functools.partial(__validate_and_convert_scalar_or_twodim.__func__, int),
|
||||
}
|
||||
|
||||
# options handled by argparse but not handled by user config
|
||||
ARGPARSE_SPECIFIC_SCHEMA = {
|
||||
"debug_dataset": bool,
|
||||
"max_token_length": Any(None, int),
|
||||
"prior_loss_weight": Any(float, int),
|
||||
}
|
||||
# for handling default None value of argparse
|
||||
ARGPARSE_NULLABLE_OPTNAMES = [
|
||||
"face_crop_aug_range",
|
||||
"resolution",
|
||||
]
|
||||
# prepare map because option name may differ among argparse and user config
|
||||
ARGPARSE_OPTNAME_TO_CONFIG_OPTNAME = {
|
||||
"train_batch_size": "batch_size",
|
||||
"dataset_repeats": "num_repeats",
|
||||
}
|
||||
|
||||
def __init__(self, support_dreambooth: bool, support_finetuning: bool, support_dropout: bool) -> None:
|
||||
assert support_dreambooth or support_finetuning, "Neither DreamBooth mode nor fine tuning mode specified. Please specify one mode or more. / DreamBooth モードか fine tuning モードのどちらも指定されていません。1つ以上指定してください。"
|
||||
|
||||
self.db_subset_schema = self.__merge_dict(
|
||||
self.SUBSET_ASCENDABLE_SCHEMA,
|
||||
self.DB_SUBSET_DISTINCT_SCHEMA,
|
||||
self.DB_SUBSET_ASCENDABLE_SCHEMA,
|
||||
self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {},
|
||||
)
|
||||
|
||||
self.ft_subset_schema = self.__merge_dict(
|
||||
self.SUBSET_ASCENDABLE_SCHEMA,
|
||||
self.FT_SUBSET_DISTINCT_SCHEMA,
|
||||
self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {},
|
||||
)
|
||||
|
||||
self.db_dataset_schema = self.__merge_dict(
|
||||
self.DATASET_ASCENDABLE_SCHEMA,
|
||||
self.SUBSET_ASCENDABLE_SCHEMA,
|
||||
self.DB_SUBSET_ASCENDABLE_SCHEMA,
|
||||
self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {},
|
||||
{"subsets": [self.db_subset_schema]},
|
||||
)
|
||||
|
||||
self.ft_dataset_schema = self.__merge_dict(
|
||||
self.DATASET_ASCENDABLE_SCHEMA,
|
||||
self.SUBSET_ASCENDABLE_SCHEMA,
|
||||
self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {},
|
||||
{"subsets": [self.ft_subset_schema]},
|
||||
)
|
||||
|
||||
if support_dreambooth and support_finetuning:
|
||||
def validate_flex_dataset(dataset_config: dict):
|
||||
subsets_config = dataset_config.get("subsets", [])
|
||||
|
||||
# check dataset meets FT style
|
||||
# NOTE: all FT subsets should have "metadata_file"
|
||||
if all(["metadata_file" in subset for subset in subsets_config]):
|
||||
return Schema(self.ft_dataset_schema)(dataset_config)
|
||||
# check dataset meets DB style
|
||||
# NOTE: all DB subsets should have no "metadata_file"
|
||||
elif all(["metadata_file" not in subset for subset in subsets_config]):
|
||||
return Schema(self.db_dataset_schema)(dataset_config)
|
||||
else:
|
||||
raise voluptuous.Invalid("DreamBooth subset and fine tuning subset cannot be mixed in the same dataset. Please split them into separate datasets. / DreamBoothのサブセットとfine tuninのサブセットを同一のデータセットに混在させることはできません。別々のデータセットに分割してください。")
|
||||
|
||||
self.dataset_schema = validate_flex_dataset
|
||||
elif support_dreambooth:
|
||||
self.dataset_schema = self.db_dataset_schema
|
||||
else:
|
||||
self.dataset_schema = self.ft_dataset_schema
|
||||
|
||||
self.general_schema = self.__merge_dict(
|
||||
self.DATASET_ASCENDABLE_SCHEMA,
|
||||
self.SUBSET_ASCENDABLE_SCHEMA,
|
||||
self.DB_SUBSET_ASCENDABLE_SCHEMA if support_dreambooth else {},
|
||||
self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {},
|
||||
)
|
||||
|
||||
self.user_config_validator = Schema({
|
||||
"general": self.general_schema,
|
||||
"datasets": [self.dataset_schema],
|
||||
})
|
||||
|
||||
self.argparse_schema = self.__merge_dict(
|
||||
self.general_schema,
|
||||
self.ARGPARSE_SPECIFIC_SCHEMA,
|
||||
{optname: Any(None, self.general_schema[optname]) for optname in self.ARGPARSE_NULLABLE_OPTNAMES},
|
||||
{a_name: self.general_schema[c_name] for a_name, c_name in self.ARGPARSE_OPTNAME_TO_CONFIG_OPTNAME.items()},
|
||||
)
|
||||
|
||||
self.argparse_config_validator = Schema(Object(self.argparse_schema), extra=voluptuous.ALLOW_EXTRA)
|
||||
|
||||
def sanitize_user_config(self, user_config: dict) -> dict:
|
||||
try:
|
||||
return self.user_config_validator(user_config)
|
||||
except MultipleInvalid:
|
||||
# TODO: エラー発生時のメッセージをわかりやすくする
|
||||
print("Invalid user config / ユーザ設定の形式が正しくないようです")
|
||||
raise
|
||||
|
||||
# NOTE: In nature, argument parser result is not needed to be sanitize
|
||||
# However this will help us to detect program bug
|
||||
def sanitize_argparse_namespace(self, argparse_namespace: argparse.Namespace) -> argparse.Namespace:
|
||||
try:
|
||||
return self.argparse_config_validator(argparse_namespace)
|
||||
except MultipleInvalid:
|
||||
# XXX: this should be a bug
|
||||
print("Invalid cmdline parsed arguments. This should be a bug. / コマンドラインのパース結果が正しくないようです。プログラムのバグの可能性が高いです。")
|
||||
raise
|
||||
|
||||
# NOTE: value would be overwritten by latter dict if there is already the same key
|
||||
@staticmethod
|
||||
def __merge_dict(*dict_list: dict) -> dict:
|
||||
merged = {}
|
||||
for schema in dict_list:
|
||||
# merged |= schema
|
||||
for k, v in schema.items():
|
||||
merged[k] = v
|
||||
return merged
|
||||
|
||||
|
||||
class BlueprintGenerator:
|
||||
BLUEPRINT_PARAM_NAME_TO_CONFIG_OPTNAME = {
|
||||
}
|
||||
|
||||
def __init__(self, sanitizer: ConfigSanitizer):
|
||||
self.sanitizer = sanitizer
|
||||
|
||||
# runtime_params is for parameters which is only configurable on runtime, such as tokenizer
|
||||
def generate(self, user_config: dict, argparse_namespace: argparse.Namespace, **runtime_params) -> Blueprint:
|
||||
sanitized_user_config = self.sanitizer.sanitize_user_config(user_config)
|
||||
sanitized_argparse_namespace = self.sanitizer.sanitize_argparse_namespace(argparse_namespace)
|
||||
|
||||
# convert argparse namespace to dict like config
|
||||
# NOTE: it is ok to have extra entries in dict
|
||||
optname_map = self.sanitizer.ARGPARSE_OPTNAME_TO_CONFIG_OPTNAME
|
||||
argparse_config = {optname_map.get(optname, optname): value for optname, value in vars(sanitized_argparse_namespace).items()}
|
||||
|
||||
general_config = sanitized_user_config.get("general", {})
|
||||
|
||||
dataset_blueprints = []
|
||||
for dataset_config in sanitized_user_config.get("datasets", []):
|
||||
# NOTE: if subsets have no "metadata_file", these are DreamBooth datasets/subsets
|
||||
subsets = dataset_config.get("subsets", [])
|
||||
is_dreambooth = all(["metadata_file" not in subset for subset in subsets])
|
||||
if is_dreambooth:
|
||||
subset_params_klass = DreamBoothSubsetParams
|
||||
dataset_params_klass = DreamBoothDatasetParams
|
||||
else:
|
||||
subset_params_klass = FineTuningSubsetParams
|
||||
dataset_params_klass = FineTuningDatasetParams
|
||||
|
||||
subset_blueprints = []
|
||||
for subset_config in subsets:
|
||||
params = self.generate_params_by_fallbacks(subset_params_klass,
|
||||
[subset_config, dataset_config, general_config, argparse_config, runtime_params])
|
||||
subset_blueprints.append(SubsetBlueprint(params))
|
||||
|
||||
params = self.generate_params_by_fallbacks(dataset_params_klass,
|
||||
[dataset_config, general_config, argparse_config, runtime_params])
|
||||
dataset_blueprints.append(DatasetBlueprint(is_dreambooth, params, subset_blueprints))
|
||||
|
||||
dataset_group_blueprint = DatasetGroupBlueprint(dataset_blueprints)
|
||||
|
||||
return Blueprint(dataset_group_blueprint)
|
||||
|
||||
@staticmethod
|
||||
def generate_params_by_fallbacks(param_klass, fallbacks: Sequence[dict]):
|
||||
name_map = BlueprintGenerator.BLUEPRINT_PARAM_NAME_TO_CONFIG_OPTNAME
|
||||
search_value = BlueprintGenerator.search_value
|
||||
default_params = asdict(param_klass())
|
||||
param_names = default_params.keys()
|
||||
|
||||
params = {name: search_value(name_map.get(name, name), fallbacks, default_params.get(name)) for name in param_names}
|
||||
|
||||
return param_klass(**params)
|
||||
|
||||
@staticmethod
|
||||
def search_value(key: str, fallbacks: Sequence[dict], default_value = None):
|
||||
for cand in fallbacks:
|
||||
value = cand.get(key)
|
||||
if value is not None:
|
||||
return value
|
||||
|
||||
return default_value
|
||||
|
||||
|
||||
def generate_dataset_group_by_blueprint(dataset_group_blueprint: DatasetGroupBlueprint):
|
||||
datasets: List[Union[DreamBoothDataset, FineTuningDataset]] = []
|
||||
|
||||
for dataset_blueprint in dataset_group_blueprint.datasets:
|
||||
if dataset_blueprint.is_dreambooth:
|
||||
subset_klass = DreamBoothSubset
|
||||
dataset_klass = DreamBoothDataset
|
||||
else:
|
||||
subset_klass = FineTuningSubset
|
||||
dataset_klass = FineTuningDataset
|
||||
|
||||
subsets = [subset_klass(**asdict(subset_blueprint.params)) for subset_blueprint in dataset_blueprint.subsets]
|
||||
dataset = dataset_klass(subsets=subsets, **asdict(dataset_blueprint.params))
|
||||
datasets.append(dataset)
|
||||
|
||||
# print info
|
||||
info = ""
|
||||
for i, dataset in enumerate(datasets):
|
||||
is_dreambooth = isinstance(dataset, DreamBoothDataset)
|
||||
info += dedent(f"""\
|
||||
[Dataset {i}]
|
||||
batch_size: {dataset.batch_size}
|
||||
resolution: {(dataset.width, dataset.height)}
|
||||
enable_bucket: {dataset.enable_bucket}
|
||||
""")
|
||||
|
||||
if dataset.enable_bucket:
|
||||
info += indent(dedent(f"""\
|
||||
min_bucket_reso: {dataset.min_bucket_reso}
|
||||
max_bucket_reso: {dataset.max_bucket_reso}
|
||||
bucket_reso_steps: {dataset.bucket_reso_steps}
|
||||
bucket_no_upscale: {dataset.bucket_no_upscale}
|
||||
\n"""), " ")
|
||||
else:
|
||||
info += "\n"
|
||||
|
||||
for j, subset in enumerate(dataset.subsets):
|
||||
info += indent(dedent(f"""\
|
||||
[Subset {j} of Dataset {i}]
|
||||
image_dir: "{subset.image_dir}"
|
||||
image_count: {subset.img_count}
|
||||
num_repeats: {subset.num_repeats}
|
||||
shuffle_caption: {subset.shuffle_caption}
|
||||
keep_tokens: {subset.keep_tokens}
|
||||
caption_dropout_rate: {subset.caption_dropout_rate}
|
||||
caption_dropout_every_n_epoches: {subset.caption_dropout_every_n_epochs}
|
||||
caption_tag_dropout_rate: {subset.caption_tag_dropout_rate}
|
||||
color_aug: {subset.color_aug}
|
||||
flip_aug: {subset.flip_aug}
|
||||
face_crop_aug_range: {subset.face_crop_aug_range}
|
||||
random_crop: {subset.random_crop}
|
||||
"""), " ")
|
||||
|
||||
if is_dreambooth:
|
||||
info += indent(dedent(f"""\
|
||||
is_reg: {subset.is_reg}
|
||||
class_tokens: {subset.class_tokens}
|
||||
caption_extension: {subset.caption_extension}
|
||||
\n"""), " ")
|
||||
else:
|
||||
info += indent(dedent(f"""\
|
||||
metadata_file: {subset.metadata_file}
|
||||
\n"""), " ")
|
||||
|
||||
print(info)
|
||||
|
||||
# make buckets first because it determines the length of dataset
|
||||
for i, dataset in enumerate(datasets):
|
||||
print(f"[Dataset {i}]")
|
||||
dataset.make_buckets()
|
||||
|
||||
return DatasetGroup(datasets)
|
||||
|
||||
|
||||
def generate_dreambooth_subsets_config_by_subdirs(train_data_dir: Optional[str] = None, reg_data_dir: Optional[str] = None):
|
||||
def extract_dreambooth_params(name: str) -> Tuple[int, str]:
|
||||
tokens = name.split('_')
|
||||
try:
|
||||
n_repeats = int(tokens[0])
|
||||
except ValueError as e:
|
||||
print(f"ignore directory without repeats / 繰り返し回数のないディレクトリを無視します: {dir}")
|
||||
return 0, ""
|
||||
caption_by_folder = '_'.join(tokens[1:])
|
||||
return n_repeats, caption_by_folder
|
||||
|
||||
def generate(base_dir: Optional[str], is_reg: bool):
|
||||
if base_dir is None:
|
||||
return []
|
||||
|
||||
base_dir: Path = Path(base_dir)
|
||||
if not base_dir.is_dir():
|
||||
return []
|
||||
|
||||
subsets_config = []
|
||||
for subdir in base_dir.iterdir():
|
||||
if not subdir.is_dir():
|
||||
continue
|
||||
|
||||
num_repeats, class_tokens = extract_dreambooth_params(subdir.name)
|
||||
if num_repeats < 1:
|
||||
continue
|
||||
|
||||
subset_config = {"image_dir": str(subdir), "num_repeats": num_repeats, "is_reg": is_reg, "class_tokens": class_tokens}
|
||||
subsets_config.append(subset_config)
|
||||
|
||||
return subsets_config
|
||||
|
||||
subsets_config = []
|
||||
subsets_config += generate(train_data_dir, False)
|
||||
subsets_config += generate(reg_data_dir, True)
|
||||
|
||||
return subsets_config
|
||||
|
||||
|
||||
def load_user_config(file: str) -> dict:
|
||||
file: Path = Path(file)
|
||||
if not file.is_file():
|
||||
raise ValueError(f"file not found / ファイルが見つかりません: {file}")
|
||||
|
||||
if file.name.lower().endswith('.json'):
|
||||
try:
|
||||
config = json.load(file)
|
||||
except Exception:
|
||||
print(f"Error on parsing JSON config file. Please check the format. / JSON 形式の設定ファイルの読み込みに失敗しました。文法が正しいか確認してください。: {file}")
|
||||
raise
|
||||
elif file.name.lower().endswith('.toml'):
|
||||
try:
|
||||
config = toml.load(file)
|
||||
except Exception:
|
||||
print(f"Error on parsing TOML config file. Please check the format. / TOML 形式の設定ファイルの読み込みに失敗しました。文法が正しいか確認してください。: {file}")
|
||||
raise
|
||||
else:
|
||||
raise ValueError(f"not supported config file format / 対応していない設定ファイルの形式です: {file}")
|
||||
|
||||
return config
|
||||
|
||||
|
||||
# for config test
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--support_dreambooth", action="store_true")
|
||||
parser.add_argument("--support_finetuning", action="store_true")
|
||||
parser.add_argument("--support_dropout", action="store_true")
|
||||
parser.add_argument("dataset_config")
|
||||
config_args, remain = parser.parse_known_args()
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
train_util.add_dataset_arguments(parser, config_args.support_dreambooth, config_args.support_finetuning, config_args.support_dropout)
|
||||
train_util.add_training_arguments(parser, config_args.support_dreambooth)
|
||||
argparse_namespace = parser.parse_args(remain)
|
||||
train_util.prepare_dataset_args(argparse_namespace, config_args.support_finetuning)
|
||||
|
||||
print("[argparse_namespace]")
|
||||
print(vars(argparse_namespace))
|
||||
|
||||
user_config = load_user_config(config_args.dataset_config)
|
||||
|
||||
print("\n[user_config]")
|
||||
print(user_config)
|
||||
|
||||
sanitizer = ConfigSanitizer(config_args.support_dreambooth, config_args.support_finetuning, config_args.support_dropout)
|
||||
sanitized_user_config = sanitizer.sanitize_user_config(user_config)
|
||||
|
||||
print("\n[sanitized_user_config]")
|
||||
print(sanitized_user_config)
|
||||
|
||||
blueprint = BlueprintGenerator(sanitizer).generate(user_config, argparse_namespace)
|
||||
|
||||
print("\n[blueprint]")
|
||||
print(blueprint)
|
||||
96
library/sampler_gui.py
Normal file
96
library/sampler_gui.py
Normal file
@ -0,0 +1,96 @@
|
||||
import tempfile
|
||||
import gradio as gr
|
||||
from easygui import msgbox
|
||||
|
||||
folder_symbol = '\U0001f4c2' # 📂
|
||||
refresh_symbol = '\U0001f504' # 🔄
|
||||
save_style_symbol = '\U0001f4be' # 💾
|
||||
document_symbol = '\U0001F4C4' # 📄
|
||||
|
||||
|
||||
###
|
||||
### Gradio common sampler GUI section
|
||||
###
|
||||
|
||||
|
||||
def sample_gradio_config():
|
||||
with gr.Accordion('Sample images config', open=False):
|
||||
with gr.Row():
|
||||
sample_every_n_steps = gr.Number(
|
||||
label='Sample every n steps',
|
||||
value=0,
|
||||
precision=0,
|
||||
interactive=True,
|
||||
)
|
||||
sample_every_n_epochs = gr.Number(
|
||||
label='Sample every n epochs',
|
||||
value=0,
|
||||
precision=0,
|
||||
interactive=True,
|
||||
)
|
||||
sample_sampler = gr.Dropdown(
|
||||
label='Sample sampler',
|
||||
choices=[
|
||||
'ddim',
|
||||
'pndm',
|
||||
'lms',
|
||||
'euler',
|
||||
'euler_a',
|
||||
'heun',
|
||||
'dpm_2',
|
||||
'dpm_2_a',
|
||||
'dpmsolver',
|
||||
'dpmsolver++',
|
||||
'dpmsingle',
|
||||
'k_lms',
|
||||
'k_euler',
|
||||
'k_euler_a',
|
||||
'k_dpm_2',
|
||||
'k_dpm_2_a',
|
||||
],
|
||||
value='euler_a',
|
||||
interactive=True,
|
||||
)
|
||||
with gr.Row():
|
||||
sample_prompts = gr.Textbox(
|
||||
lines=5,
|
||||
label='Sample prompts',
|
||||
interactive=True,
|
||||
)
|
||||
return (
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
)
|
||||
|
||||
|
||||
def run_cmd_sample(
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
):
|
||||
run_cmd = ''
|
||||
|
||||
if sample_every_n_epochs == 0 and sample_every_n_steps == 0:
|
||||
return run_cmd
|
||||
|
||||
# Create a temporary file and get its path
|
||||
with tempfile.NamedTemporaryFile(mode='w', delete=False) as temp_file:
|
||||
# Write the contents of the variable to the file
|
||||
temp_file.write(sample_prompts)
|
||||
|
||||
# Get the path of the temporary file
|
||||
sample_prompts_path = temp_file.name
|
||||
|
||||
run_cmd += f' --sample_sampler={sample_sampler}'
|
||||
run_cmd += f' --sample_prompts="{sample_prompts_path}"'
|
||||
|
||||
if not sample_every_n_epochs == 0:
|
||||
run_cmd += f' --sample_every_n_epochs="{sample_every_n_epochs}"'
|
||||
|
||||
if not sample_every_n_steps == 0:
|
||||
run_cmd += f' --sample_every_n_steps="{sample_every_n_steps}"'
|
||||
|
||||
return run_cmd
|
||||
File diff suppressed because it is too large
Load Diff
47
lora_gui.py
47
lora_gui.py
@ -40,6 +40,7 @@ from library.utilities import utilities_tab
|
||||
from library.merge_lora_gui import gradio_merge_lora_tab
|
||||
from library.verify_lora_gui import gradio_verify_lora_tab
|
||||
from library.resize_lora_gui import gradio_resize_lora_tab
|
||||
from library.sampler_gui import sample_gradio_config, run_cmd_sample
|
||||
from easygui import msgbox
|
||||
|
||||
folder_symbol = '\U0001f4c2' # 📂
|
||||
@ -112,9 +113,13 @@ def save_configuration(
|
||||
optimizer,
|
||||
optimizer_args,
|
||||
noise_offset,
|
||||
LoRA_type='Standard',
|
||||
conv_dim=0,
|
||||
conv_alpha=0,
|
||||
LoRA_type,
|
||||
conv_dim,
|
||||
conv_alpha,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
):
|
||||
# Get list of function parameters and values
|
||||
parameters = list(locals().items())
|
||||
@ -223,9 +228,13 @@ def open_configuration(
|
||||
optimizer,
|
||||
optimizer_args,
|
||||
noise_offset,
|
||||
LoRA_type='Standard',
|
||||
conv_dim=0,
|
||||
conv_alpha=0,
|
||||
LoRA_type,
|
||||
conv_dim,
|
||||
conv_alpha,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
):
|
||||
# Get list of function parameters and values
|
||||
parameters = list(locals().items())
|
||||
@ -323,6 +332,10 @@ def train_model(
|
||||
LoRA_type,
|
||||
conv_dim,
|
||||
conv_alpha,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
):
|
||||
if pretrained_model_name_or_path == '':
|
||||
msgbox('Source model information is missing')
|
||||
@ -544,6 +557,13 @@ def train_model(
|
||||
noise_offset=noise_offset,
|
||||
)
|
||||
|
||||
run_cmd += run_cmd_sample(
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
)
|
||||
|
||||
print(run_cmd)
|
||||
|
||||
# Run the command
|
||||
@ -826,11 +846,12 @@ def lora_tab(
|
||||
outputs=[cache_latents],
|
||||
)
|
||||
|
||||
# optimizer.change(
|
||||
# set_legacy_8bitadam,
|
||||
# inputs=[optimizer, use_8bit_adam],
|
||||
# outputs=[optimizer, use_8bit_adam],
|
||||
# )
|
||||
(
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
) = sample_gradio_config()
|
||||
|
||||
with gr.Tab('Tools'):
|
||||
gr.Markdown(
|
||||
@ -927,6 +948,10 @@ def lora_tab(
|
||||
LoRA_type,
|
||||
conv_dim,
|
||||
conv_alpha,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
]
|
||||
|
||||
button_open_config.click(
|
||||
|
||||
@ -14,7 +14,9 @@ pytorch-lightning==1.9.0
|
||||
safetensors==0.2.6
|
||||
tensorboard==2.10.1
|
||||
tk==0.1.0
|
||||
toml==0.10.2
|
||||
transformers==4.26.0
|
||||
voluptuous==0.13.1
|
||||
# for BLIP captioning
|
||||
fairscale==0.4.13
|
||||
requests==2.28.2
|
||||
|
||||
@ -36,6 +36,7 @@ from library.dreambooth_folder_creation_gui import (
|
||||
gradio_dreambooth_folder_creation_tab,
|
||||
)
|
||||
from library.utilities import utilities_tab
|
||||
from library.sampler_gui import sample_gradio_config, run_cmd_sample
|
||||
from easygui import msgbox
|
||||
|
||||
folder_symbol = '\U0001f4c2' # 📂
|
||||
@ -106,6 +107,10 @@ def save_configuration(
|
||||
optimizer,
|
||||
optimizer_args,
|
||||
noise_offset,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
):
|
||||
# Get list of function parameters and values
|
||||
parameters = list(locals().items())
|
||||
@ -213,6 +218,10 @@ def open_configuration(
|
||||
optimizer,
|
||||
optimizer_args,
|
||||
noise_offset,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
):
|
||||
# Get list of function parameters and values
|
||||
parameters = list(locals().items())
|
||||
@ -299,6 +308,10 @@ def train_model(
|
||||
optimizer,
|
||||
optimizer_args,
|
||||
noise_offset,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
):
|
||||
if pretrained_model_name_or_path == '':
|
||||
msgbox('Source model information is missing')
|
||||
@ -496,6 +509,13 @@ def train_model(
|
||||
elif template == 'style template':
|
||||
run_cmd += f' --use_style_template'
|
||||
|
||||
run_cmd += run_cmd_sample(
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
)
|
||||
|
||||
print(run_cmd)
|
||||
|
||||
# Run the command
|
||||
@ -740,11 +760,14 @@ def ti_tab(
|
||||
inputs=[color_aug],
|
||||
outputs=[cache_latents],
|
||||
)
|
||||
# optimizer.change(
|
||||
# set_legacy_8bitadam,
|
||||
# inputs=[optimizer, use_8bit_adam],
|
||||
# outputs=[optimizer, use_8bit_adam],
|
||||
# )
|
||||
|
||||
(
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
) = sample_gradio_config()
|
||||
|
||||
with gr.Tab('Tools'):
|
||||
gr.Markdown(
|
||||
'This section provide Dreambooth tools to help setup your dataset...'
|
||||
@ -832,6 +855,10 @@ def ti_tab(
|
||||
optimizer,
|
||||
optimizer_args,
|
||||
noise_offset,
|
||||
sample_every_n_steps,
|
||||
sample_every_n_epochs,
|
||||
sample_sampler,
|
||||
sample_prompts,
|
||||
]
|
||||
|
||||
button_open_config.click(
|
||||
|
||||
619
train_README-ja.md
Normal file
619
train_README-ja.md
Normal file
@ -0,0 +1,619 @@
|
||||
当リポジトリではモデルのfine tuning、DreamBooth、およびLoRAとTextual Inversionの学習をサポートします。この文書ではそれらに共通する、学習データの準備方法やスクリプトオプションについて説明します。
|
||||
|
||||
# 概要
|
||||
|
||||
あらかじめこのリポジトリのREADMEを参照し、環境整備を行ってください。
|
||||
|
||||
|
||||
以下について説明します。
|
||||
|
||||
1. 学習データの準備について（設定ファイルを用いる新形式）
|
||||
1. Aspect Ratio Bucketingについて
|
||||
1. 以前の指定形式（設定ファイルを用いずコマンドラインから指定）
|
||||
1. fine tuning 方式のメタデータ準備：キャプションニングなど
|
||||
|
||||
1.だけ実行すればとりあえず学習は可能です（学習については各スクリプトのドキュメントを参照）。2.以降は必要に応じて参照してください。
|
||||
|
||||
<!--
|
||||
1. 各スクリプトで共通のオプション
|
||||
-->
|
||||
|
||||
# 学習データの準備について
|
||||
|
||||
任意のフォルダ（複数でも可）に学習データの画像ファイルを用意しておきます。`.png`, `.jpg`, `.jpeg`, `.webp`, `.bmp` をサポートします。リサイズなどの前処理は基本的に必要ありません。
|
||||
|
||||
ただし学習解像度（後述）よりも極端に小さい画像は使わないか、あらかじめ超解像AIなどで拡大しておくことをお勧めします。また極端に大きな画像（3000x3000ピクセル程度？）よりも大きな画像はエラーになる場合があるようですので事前に縮小してください。
|
||||
|
||||
学習時には、モデルに学ばせる画像データを整理し、スクリプトに対して指定する必要があります。学習データの数、学習対象、キャプション（画像の説明）が用意できるか否かなどにより、いくつかの方法で学習データを指定できます。以下の方式があります（それぞれの名前は一般的なものではなく、当リポジトリ独自の定義です）。正則化画像については後述します。
|
||||
|
||||
1. DreamBooth、class+identifier方式（正則化画像使用可）
|
||||
|
||||
特定の単語 (identifier) に学習対象を紐づけるように学習します。キャプションを用意する必要はありません。たとえば特定のキャラを学ばせる場合に使うとキャプションを用意する必要がない分、手軽ですが、髪型や服装、背景など学習データの全要素が identifier に紐づけられて学習されるため、生成時のプロンプトで服が変えられない、といった事態も起こりえます。
|
||||
|
||||
1. DreamBooth、キャプション方式（正則化画像使用可）
|
||||
|
||||
画像ごとにキャプションが記録されたテキストファイルを用意して学習します。たとえば特定のキャラを学ばせると、画像の詳細をキャプションに記述することで（白い服を着たキャラA、赤い服を着たキャラA、など）キャラとそれ以外の要素が分離され、より厳密にモデルがキャラだけを学ぶことが期待できます。
|
||||
|
||||
1. fine tuning方式（正則化画像使用不可）
|
||||
|
||||
あらかじめキャプションをメタデータファイルにまとめます。タグとキャプションを分けて管理したり、学習を高速化するためlatentsを事前キャッシュしたりなどの機能をサポートします（いずれも別文書で説明しています）。
|
||||
|
||||
学習したいものと使用できる指定方法の組み合わせは以下の通りです。
|
||||
|
||||
| 学習対象または方法 | スクリプト | DB / class+identifier | DB / キャプション | fine tuning |
|
||||
| ----- | ----- | ----- | ----- | ----- |
|
||||
| モデルをfine tuning | `fine_tune.py`| x | x | o |
|
||||
| モデルをDreamBooth | `train_db.py`| o | o | x |
|
||||
| LoRA | `train_network.py`| o | o | o |
|
||||
| Textual Invesion | `train_textual_inversion.py`| o | o | o |
|
||||
|
||||
## どれを選ぶか
|
||||
|
||||
LoRA、Textual Inversionについては、手軽にキャプションファイルを用意せずに学習したい場合はDreamBooth class+identifier、用意できるならDreamBooth キャプション方式がよいでしょう。学習データの枚数が多く、かつ正則化画像を使用しない場合はfine tuning方式も検討してください。
|
||||
|
||||
DreamBoothについても同様ですが、fine tuning方式は使えません。fine tuningの場合はfine tuning方式のみです。
|
||||
|
||||
# 各方式の指定方法について
|
||||
|
||||
ここではそれぞれの指定方法で典型的なパターンについてだけ説明します。より詳細な指定方法については [データセット設定](./config_README-ja.md) をご覧ください。
|
||||
|
||||
# DreamBooth、class+identifier方式（正則化画像使用可）
|
||||
|
||||
この方式では、各画像は `class identifier` というキャプションで学習されたのと同じことになります（`shs dog` など）。
|
||||
|
||||
## step 1. 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`` などです。本当は Danbooru Tag に含まれないやつがより望ましいです。）
|
||||
|
||||
## step 2. 正則化画像を使うか否かを決め、使う場合には正則化画像を生成する
|
||||
|
||||
正則化画像とは、前述のclass全体が、学習対象に引っ張られることを防ぐための画像です（language drift）。正則化画像を使わないと、たとえば `shs 1girl` で特定のキャラクタを学ばせると、単なる `1girl` というプロンプトで生成してもそのキャラに似てきます。これは `1girl` が学習時のキャプションに含まれているためです。
|
||||
|
||||
学習対象の画像と正則化画像を同時に学ばせることで、class は class のままで留まり、identifier をプロンプトにつけた時だけ学習対象が生成されるようになります。
|
||||
|
||||
LoRAやDreamBoothで特定のキャラだけ出てくればよい場合は、正則化画像を用いなくても良いといえます。
|
||||
|
||||
Textual Inversionでは用いなくてよいでしょう（学ばせる token string がキャプションに含まれない場合はなにも学習されないため）。
|
||||
|
||||
正則化画像としては、学習対象のモデルで、class 名だけで生成した画像を用いるのが一般的です（たとえば `1girl`）。ただし生成画像の品質が悪い場合には、プロンプトを工夫したり、ネットから別途ダウンロードした画像を用いることもできます。
|
||||
|
||||
（正則化画像も学習されるため、その品質はモデルに影響します。）
|
||||
|
||||
一般的には数百枚程度、用意するのが望ましいようです（枚数が少ないと class 画像が一般化されずそれらの特徴を学んでしまいます）。
|
||||
|
||||
生成画像を使う場合、通常、生成画像のサイズは学習解像度（より正確にはbucketの解像度、後述）にあわせてください。
|
||||
|
||||
## step 2. 設定ファイルの記述
|
||||
|
||||
テキストファイルを作成し、拡張子を `.toml` にします。たとえば以下のように記述します。
|
||||
|
||||
（`#` で始まっている部分はコメントですので、このままコピペしてそのままでもよいですし、削除しても問題ありません。）
|
||||
|
||||
```toml
|
||||
[general]
|
||||
enable_bucket = true # Aspect Ratio Bucketingを使うか否か
|
||||
|
||||
[[datasets]]
|
||||
resolution = 512 # 学習解像度
|
||||
batch_size = 4 # バッチサイズ
|
||||
|
||||
[[datasets.subsets]]
|
||||
image_dir = 'C:\hoge' # 学習用画像を入れたフォルダを指定
|
||||
class_tokens = 'hoge girl' # identifier class を指定
|
||||
num_repeats = 10 # 学習用画像の繰り返し回数
|
||||
|
||||
# 以下は正則化画像を用いる場合のみ記述する。用いない場合は削除する
|
||||
[[datasets.subsets]]
|
||||
is_reg = true
|
||||
image_dir = 'C:\reg' # 正則化画像を入れたフォルダを指定
|
||||
class_tokens = 'girl' # class を指定
|
||||
num_repeats = 1 # 正則化画像の繰り返し回数、基本的には1でよい
|
||||
```
|
||||
|
||||
基本的には以下を場所のみ書き換えれば学習できます。
|
||||
|
||||
1. 学習解像度
|
||||
|
||||
数値1つを指定すると正方形（`512`なら512x512）、鍵カッコカンマ区切りで2つ指定すると横×縦（`[512,768]`なら512x768）になります。SD1.x系ではもともとの学習解像度は512です。`[512,768]` 等の大きめの解像度を指定すると縦長、横長画像生成時の破綻を小さくできるかもしれません。SD2.x 768系では `768` です。
|
||||
|
||||
1. バッチサイズ
|
||||
|
||||
同時に何件のデータを学習するかを指定します。GPUのVRAMサイズ、学習解像度によって変わってきます。またfine tuning/DreamBooth/LoRA等でも変わってきますので、詳しくは各スクリプトの説明をご覧ください。
|
||||
|
||||
1. フォルダ指定
|
||||
|
||||
学習用画像、正則化画像（使用する場合のみ）のフォルダを指定します。画像データが含まれているフォルダそのものを指定します。
|
||||
|
||||
1. identifier と class の指定
|
||||
|
||||
前述のサンプルの通りです。
|
||||
|
||||
1. 繰り返し回数
|
||||
|
||||
後述します。
|
||||
|
||||
### 繰り返し回数について
|
||||
|
||||
繰り返し回数は、正則化画像の枚数と学習用画像の枚数を調整するために用いられます。正則化画像の枚数は学習用画像よりも多いため、学習用画像を繰り返して枚数を合わせ、1対1の比率で学習できるようにします。
|
||||
|
||||
繰り返し回数は「 __学習用画像の繰り返し回数×学習用画像の枚数≧正則化画像の繰り返し回数×正則化画像の枚数__ 」となるように指定してください。
|
||||
|
||||
（1 epoch（データが一周すると1 epoch）のデータ数が「学習用画像の繰り返し回数×学習用画像の枚数」となります。正則化画像の枚数がそれより多いと、余った部分の正則化画像は使用されません。）
|
||||
|
||||
## step 3. 学習
|
||||
|
||||
それぞれのドキュメントを参考に学習を行ってください。
|
||||
|
||||
# DreamBooth、キャプション方式（正則化画像使用可）
|
||||
|
||||
この方式では各画像はキャプションで学習されます。
|
||||
|
||||
## step 1. キャプションファイルを準備する
|
||||
|
||||
学習用画像のフォルダに、画像と同じファイル名で、拡張子 `.caption`（設定で変えられます）のファイルを置いてください。それぞれのファイルは1行のみとしてください。エンコーディングは `UTF-8` です。
|
||||
|
||||
## step 2. 正則化画像を使うか否かを決め、使う場合には正則化画像を生成する
|
||||
|
||||
class+identifier形式と同様です。なお正則化画像にもキャプションを付けることができますが、通常は不要でしょう。
|
||||
|
||||
## step 2. 設定ファイルの記述
|
||||
|
||||
テキストファイルを作成し、拡張子を `.toml` にします。たとえば以下のように記述します。
|
||||
|
||||
```toml
|
||||
[general]
|
||||
enable_bucket = true # Aspect Ratio Bucketingを使うか否か
|
||||
|
||||
[[datasets]]
|
||||
resolution = 512 # 学習解像度
|
||||
batch_size = 4 # バッチサイズ
|
||||
|
||||
[[datasets.subsets]]
|
||||
image_dir = 'C:\hoge' # 学習用画像を入れたフォルダを指定
|
||||
caption_extension = '.caption' # キャプションファイルの拡張子 .txt を使う場合には書き換える
|
||||
num_repeats = 10 # 学習用画像の繰り返し回数
|
||||
|
||||
# 以下は正則化画像を用いる場合のみ記述する。用いない場合は削除する
|
||||
[[datasets.subsets]]
|
||||
is_reg = true
|
||||
image_dir = 'C:\reg' # 正則化画像を入れたフォルダを指定
|
||||
class_tokens = 'girl' # class を指定
|
||||
num_repeats = 1 # 正則化画像の繰り返し回数、基本的には1でよい
|
||||
```
|
||||
|
||||
基本的には以下を場所のみ書き換えれば学習できます。特に記述がない部分は class+identifier 方式と同じです。
|
||||
|
||||
1. 学習解像度
|
||||
1. バッチサイズ
|
||||
1. フォルダ指定
|
||||
1. キャプションファイルの拡張子
|
||||
|
||||
任意の拡張子を指定できます。
|
||||
1. 繰り返し回数
|
||||
|
||||
## step 3. 学習
|
||||
|
||||
それぞれのドキュメントを参考に学習を行ってください。
|
||||
|
||||
# fine tuning 方式
|
||||
|
||||
## step 1. メタデータを準備する
|
||||
|
||||
キャプションやタグをまとめた管理用ファイルをメタデータと呼びます。json形式で拡張子は `.json`
|
||||
です。作成方法は長くなりますのでこの文書の末尾に書きました。
|
||||
|
||||
## step 2. 設定ファイルの記述
|
||||
|
||||
テキストファイルを作成し、拡張子を `.toml` にします。たとえば以下のように記述します。
|
||||
|
||||
```toml
|
||||
[general]
|
||||
shuffle_caption = true
|
||||
keep_tokens = 1
|
||||
|
||||
[[datasets]]
|
||||
resolution = 512 # 学習解像度
|
||||
batch_size = 4 # バッチサイズ
|
||||
|
||||
[[datasets.subsets]]
|
||||
image_dir = 'C:\piyo' # 学習用画像を入れたフォルダを指定
|
||||
metadata_file = 'C:\piyo\piyo_md.json' # メタデータファイル名
|
||||
```
|
||||
|
||||
基本的には以下を場所のみ書き換えれば学習できます。特に記述がない部分は DreamBooth, class+identifier 方式と同じです。
|
||||
|
||||
1. 学習解像度
|
||||
1. バッチサイズ
|
||||
1. フォルダ指定
|
||||
1. メタデータファイル名
|
||||
|
||||
後述の方法で作成したメタデータファイルを指定します。
|
||||
|
||||
|
||||
## step 3. 学習
|
||||
|
||||
それぞれのドキュメントを参考に学習を行ってください。
|
||||
|
||||
# Aspect Ratio Bucketing について
|
||||
|
||||
Stable Diffusion のv1は512\*512で学習されていますが、それに加えて256\*1024や384\*640といった解像度でも学習します。これによりトリミングされる部分が減り、より正しくキャプションと画像の関係が学習されることが期待されます。
|
||||
|
||||
また任意の解像度で学習するため、事前に画像データの縦横比を統一しておく必要がなくなります。
|
||||
|
||||
設定で有効、向こうが切り替えられますが、ここまでの設定ファイルの記述例では有効になっています（`true` が設定されています）。
|
||||
|
||||
学習解像度はパラメータとして与えられた解像度の面積（＝メモリ使用量）を超えない範囲で、64ピクセル単位（デフォルト、変更可）で縦横に調整、作成されます。
|
||||
|
||||
機械学習では入力サイズをすべて統一するのが一般的ですが、特に制約があるわけではなく、実際は同一のバッチ内で統一されていれば大丈夫です。NovelAIの言うbucketingは、あらかじめ教師データを、アスペクト比に応じた学習解像度ごとに分類しておくことを指しているようです。そしてバッチを各bucket内の画像で作成することで、バッチの画像サイズを統一します。
|
||||
|
||||
# 以前のデータ指定方法
|
||||
|
||||
フォルダ名で繰り返し回数を指定する方法です。
|
||||
|
||||
## step 1. 学習用画像の準備
|
||||
|
||||
学習用画像を格納するフォルダを作成します。 __さらにその中に__ 、以下の名前でディレクトリを作成します。
|
||||
|
||||
```
|
||||
<繰り返し回数>_<identifier> <class>
|
||||
```
|
||||
|
||||
間の``_``を忘れないでください。
|
||||
|
||||
たとえば「sls frog」というプロンプトで、データを20回繰り返す場合、「20_sls frog」となります。以下のようになります。
|
||||
|
||||
|
||||
|
||||
### 複数class、複数対象（identifier）の学習
|
||||
|
||||
方法は単純で、学習用画像のフォルダ内に ``繰り返し回数_<identifier> <class>`` のフォルダを複数、正則化画像フォルダにも同様に ``繰り返し回数_<class>`` のフォルダを複数、用意してください。
|
||||
|
||||
たとえば「sls frog」と「cpc rabbit」を同時に学習する場合、以下のようになります。
|
||||
|
||||
|
||||
|
||||
classがひとつで対象が複数の場合、正則化画像フォルダはひとつで構いません。たとえば1girlにキャラAとキャラBがいる場合は次のようにします。
|
||||
|
||||
- train_girls
|
||||
- 10_sls 1girl
|
||||
- 10_cpc 1girl
|
||||
- reg_girls
|
||||
- 1_1girl
|
||||
|
||||
### DreamBoothでキャプションを使う
|
||||
|
||||
学習用画像、正則化画像のフォルダに、画像と同じファイル名で、拡張子.caption（オプションで変えられます）のファイルを置くと、そのファイルからキャプションを読み込みプロンプトとして学習します。
|
||||
|
||||
※それらの画像の学習に、フォルダ名（identifier class）は使用されなくなります。
|
||||
|
||||
キャプションファイルの拡張子はデフォルトで.captionです。学習スクリプトの `--caption_extension` オプションで変更できます。`--shuffle_caption` オプションで学習時のキャプションについて、カンマ区切りの各部分をシャッフルしながら学習します。
|
||||
|
||||
## step 2. 正則化画像の準備
|
||||
|
||||
正則化画像を使う場合の手順です。
|
||||
|
||||
正則化画像を格納するフォルダを作成します。 __さらにその中に__ ``<繰り返し回数>_<class>`` という名前でディレクトリを作成します。
|
||||
|
||||
たとえば「frog」というプロンプトで、データを繰り返さない（1回だけ）場合、以下のようになります。
|
||||
|
||||
|
||||
|
||||
|
||||
## step 3. 学習の実行
|
||||
|
||||
各学習スクリプトを実行します。 `--train_data_dir` オプションで前述の学習用データのフォルダを（__画像を含むフォルダではなく、その親フォルダ__）、`--reg_data_dir` オプションで正則化画像のフォルダ（__画像を含むフォルダではなく、その親フォルダ__）を指定してください。
|
||||
|
||||
<!--
|
||||
# 学習スクリプト共通のオプション
|
||||
|
||||
スクリプトの更新後、ドキュメントの更新が追い付いていない場合があります。その場合は `--help` オプションで使用できるオプションを確認してください。
|
||||
|
||||
## TODO 書きます
|
||||
-->
|
||||
|
||||
# メタデータファイルの作成
|
||||
|
||||
## 教師データの用意
|
||||
|
||||
前述のように学習させたい画像データを用意し、任意のフォルダに入れてください。
|
||||
|
||||
たとえば以下のように画像を格納します。
|
||||
|
||||
|
||||
|
||||
## 自動キャプショニング
|
||||
|
||||
キャプションを使わずタグだけで学習する場合はスキップしてください。
|
||||
|
||||
また手動でキャプションを用意する場合、キャプションは教師データ画像と同じディレクトリに、同じファイル名、拡張子.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でキャプションファイルが生成されます。
|
||||
|
||||
|
||||
|
||||
たとえば以下のようなキャプションが付きます。
|
||||
|
||||
|
||||
|
||||
## 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をクリックして開き、そこからダウンロードします。
|
||||
|
||||
|
||||
|
||||
以下のようなこういうディレクトリ構造にしてください
|
||||
|
||||
|
||||
|
||||
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件ずつ処理されるためわりと遅いです。
|
||||
|
||||
複数の教師データフォルダがある場合には、それぞれのフォルダに対して実行してください。
|
||||
|
||||
以下のように生成されます。
|
||||
|
||||
|
||||
|
||||
こんな感じにタグが付きます（すごい情報量……）。
|
||||
|
||||
|
||||
|
||||
## 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フォルダに自動的にダウンロードされます（フォルダはオプションで変えられます）。以下のようになります。
|
||||
|
||||
|
||||
|
||||
タグファイルが教師データ画像と同じディレクトリに、同じファイル名、拡張子.txtで作成されます。
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
threshオプションで、判定されたタグのconfidence（確信度）がいくつ以上でタグをつけるかが指定できます。デフォルトはWD14Taggerのサンプルと同じ0.35です。値を下げるとより多くのタグが付与されますが、精度は下がります。
|
||||
|
||||
batch_sizeはGPUのVRAM容量に応じて増減してください。大きいほうが速くなります（VRAM 12GBでももう少し増やせると思います）。caption_extensionオプションでタグファイルの拡張子を変更できます。デフォルトは.txtです。
|
||||
|
||||
model_dirオプションでモデルの保存先フォルダを指定できます。
|
||||
|
||||
またforce_downloadオプションを指定すると保存先フォルダがあってもモデルを再ダウンロードします。
|
||||
|
||||
複数の教師データフォルダがある場合には、それぞれのフォルダに対して実行してください。
|
||||
|
||||
## キャプションとタグ情報の前処理
|
||||
|
||||
スクリプトから処理しやすいようにキャプションとタグをメタデータとしてひとつのファイルにまとめます。
|
||||
|
||||
### キャプションの前処理
|
||||
|
||||
キャプションをメタデータに入れるには、作業フォルダ内で以下を実行してください（キャプションを学習に使わない場合は実行不要です）（実際は1行で記述します、以下同様）。`--full_path` オプションを指定してメタデータに画像ファイルの場所をフルパスで格納します。このオプションを省略すると相対パスで記録されますが、フォルダ指定が `.toml` ファイル内で別途必要になります。
|
||||
|
||||
```
|
||||
python merge_captions_to_metadata.py --full_apth <教師データフォルダ>
|
||||
--in_json <読み込むメタデータファイル名> <メタデータファイル名>
|
||||
```
|
||||
|
||||
メタデータファイル名は任意の名前です。
|
||||
教師データがtrain_data、読み込むメタデータファイルなし、メタデータファイルがmeta_cap.jsonの場合、以下のようになります。
|
||||
|
||||
```
|
||||
python merge_captions_to_metadata.py --full_path 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 --full_path <教師データフォルダ>
|
||||
--in_json <読み込むメタデータファイル名> <書き込むメタデータファイル名>
|
||||
```
|
||||
|
||||
先と同じディレクトリ構成で、meta_cap.jsonを読み、meta_cap_dd.jsonに書きだす場合、以下となります。
|
||||
```
|
||||
python merge_dd_tags_to_metadata.py --full_path 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の事前取得
|
||||
|
||||
※ このステップは必須ではありません。省略しても学習時にlatentsを取得しながら学習できます。
|
||||
また学習時に `random_crop` や `color_aug` などを行う場合にはlatentsの事前取得はできません（画像を毎回変えながら学習するため）。事前取得をしない場合、ここまでのメタデータで学習できます。
|
||||
|
||||
あらかじめ画像の潜在表現を取得しディスクに保存しておきます。それにより、学習を高速に進めることができます。あわせてbucketing（教師データをアスペクト比に応じて分類する）を行います。
|
||||
|
||||
作業フォルダで以下のように入力してください。
|
||||
```
|
||||
python prepare_buckets_latents.py --full_path <教師データフォルダ>
|
||||
<読み込むメタデータファイル名> <書き込むメタデータファイル名>
|
||||
<fine tuningするモデル名またはcheckpoint>
|
||||
--batch_size <バッチサイズ>
|
||||
--max_resolution <解像度 幅,高さ>
|
||||
--mixed_precision <精度>
|
||||
```
|
||||
|
||||
モデルがmodel.ckpt、バッチサイズ4、学習解像度は512\*512、精度no（float32）で、meta_clean.jsonからメタデータを読み込み、meta_lat.jsonに書き込む場合、以下のようになります。
|
||||
|
||||
```
|
||||
python prepare_buckets_latents.py --full_path
|
||||
train_data meta_clean.json meta_lat.json model.ckpt
|
||||
--batch_size 4 --max_resolution 512,512 --mixed_precision no
|
||||
```
|
||||
|
||||
教師データフォルダにnumpyのnpz形式でlatentsが保存されます。
|
||||
|
||||
解像度の最小サイズを--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の結果が表示されます。
|
||||
|
||||
|
||||
|
||||
複数の教師データフォルダがある場合には、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
|
||||
|
||||
```
|
||||
読み込み元と書き込み先を同じにすることも可能ですが別々の方が安全です。
|
||||
|
||||
__※引数を都度書き換えて、別のメタデータファイルに書き込むと安全です。__
|
||||
|
||||
54
train_db.py
54
train_db.py
@ -15,7 +15,11 @@ import diffusers
|
||||
from diffusers import DDPMScheduler
|
||||
|
||||
import library.train_util as train_util
|
||||
from library.train_util import DreamBoothDataset
|
||||
import library.config_util as config_util
|
||||
from library.config_util import (
|
||||
ConfigSanitizer,
|
||||
BlueprintGenerator,
|
||||
)
|
||||
|
||||
|
||||
def collate_fn(examples):
|
||||
@ -33,24 +37,33 @@ def train(args):
|
||||
|
||||
tokenizer = train_util.load_tokenizer(args)
|
||||
|
||||
train_dataset = DreamBoothDataset(args.train_batch_size, args.train_data_dir, args.reg_data_dir,
|
||||
tokenizer, args.max_token_length, args.caption_extension, args.shuffle_caption, args.keep_tokens,
|
||||
args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
|
||||
args.bucket_reso_steps, args.bucket_no_upscale,
|
||||
args.prior_loss_weight, args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop, args.debug_dataset)
|
||||
blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, False, True))
|
||||
if args.dataset_config is not None:
|
||||
print(f"Load dataset config from {args.dataset_config}")
|
||||
user_config = config_util.load_user_config(args.dataset_config)
|
||||
ignored = ["train_data_dir", "reg_data_dir"]
|
||||
if any(getattr(args, attr) is not None for attr in ignored):
|
||||
print("ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(', '.join(ignored)))
|
||||
else:
|
||||
user_config = {
|
||||
"datasets": [{
|
||||
"subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)
|
||||
}]
|
||||
}
|
||||
|
||||
blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
|
||||
train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
|
||||
|
||||
if args.no_token_padding:
|
||||
train_dataset.disable_token_padding()
|
||||
|
||||
# 学習データのdropout率を設定する
|
||||
train_dataset.set_caption_dropout(args.caption_dropout_rate, args.caption_dropout_every_n_epochs, args.caption_tag_dropout_rate)
|
||||
|
||||
train_dataset.make_buckets()
|
||||
train_dataset_group.disable_token_padding()
|
||||
|
||||
if args.debug_dataset:
|
||||
train_util.debug_dataset(train_dataset)
|
||||
train_util.debug_dataset(train_dataset_group)
|
||||
return
|
||||
|
||||
if cache_latents:
|
||||
assert train_dataset_group.is_latent_cacheable(), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
|
||||
|
||||
# acceleratorを準備する
|
||||
print("prepare accelerator")
|
||||
|
||||
@ -91,7 +104,7 @@ def train(args):
|
||||
vae.requires_grad_(False)
|
||||
vae.eval()
|
||||
with torch.no_grad():
|
||||
train_dataset.cache_latents(vae)
|
||||
train_dataset_group.cache_latents(vae)
|
||||
vae.to("cpu")
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
@ -126,7 +139,7 @@ def train(args):
|
||||
# DataLoaderのプロセス数：0はメインプロセスになる
|
||||
n_workers = min(args.max_data_loader_n_workers, os.cpu_count() - 1) # cpu_count-1 ただし最大で指定された数まで
|
||||
train_dataloader = torch.utils.data.DataLoader(
|
||||
train_dataset, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
|
||||
train_dataset_group, batch_size=1, shuffle=True, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
|
||||
|
||||
# 学習ステップ数を計算する
|
||||
if args.max_train_epochs is not None:
|
||||
@ -176,8 +189,8 @@ def train(args):
|
||||
# 学習する
|
||||
total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
|
||||
print("running training / 学習開始")
|
||||
print(f" num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset.num_train_images}")
|
||||
print(f" num reg images / 正則化画像の数: {train_dataset.num_reg_images}")
|
||||
print(f" num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
|
||||
print(f" num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
|
||||
print(f" num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
|
||||
print(f" num epochs / epoch数: {num_train_epochs}")
|
||||
print(f" batch size per device / バッチサイズ: {args.train_batch_size}")
|
||||
@ -198,7 +211,7 @@ def train(args):
|
||||
loss_total = 0.0
|
||||
for epoch in range(num_train_epochs):
|
||||
print(f"epoch {epoch+1}/{num_train_epochs}")
|
||||
train_dataset.set_current_epoch(epoch + 1)
|
||||
train_dataset_group.set_current_epoch(epoch + 1)
|
||||
|
||||
# 指定したステップ数までText Encoderを学習する：epoch最初の状態
|
||||
unet.train()
|
||||
@ -278,6 +291,8 @@ def train(args):
|
||||
progress_bar.update(1)
|
||||
global_step += 1
|
||||
|
||||
train_util.sample_images(accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
|
||||
|
||||
current_loss = loss.detach().item()
|
||||
if args.logging_dir is not None:
|
||||
logs = {"loss": current_loss, "lr": float(lr_scheduler.get_last_lr()[0])}
|
||||
@ -309,6 +324,8 @@ def train(args):
|
||||
train_util.save_sd_model_on_epoch_end(args, accelerator, src_path, save_stable_diffusion_format, use_safetensors,
|
||||
save_dtype, epoch, num_train_epochs, global_step, unwrap_model(text_encoder), unwrap_model(unet), vae)
|
||||
|
||||
train_util.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
|
||||
|
||||
is_main_process = accelerator.is_main_process
|
||||
if is_main_process:
|
||||
unet = unwrap_model(unet)
|
||||
@ -336,6 +353,7 @@ if __name__ == '__main__':
|
||||
train_util.add_training_arguments(parser, True)
|
||||
train_util.add_sd_saving_arguments(parser)
|
||||
train_util.add_optimizer_arguments(parser)
|
||||
config_util.add_config_arguments(parser)
|
||||
|
||||
parser.add_argument("--no_token_padding", action="store_true",
|
||||
help="disable token padding (same as Diffuser's DreamBooth) / トークンのpaddingを無効にする（Diffusers版DreamBoothと同じ動作）")
|
||||
|
||||
227
train_network.py
227
train_network.py
@ -1,4 +1,3 @@
|
||||
from torch.cuda.amp import autocast
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
import importlib
|
||||
import argparse
|
||||
@ -12,11 +11,17 @@ import json
|
||||
from tqdm import tqdm
|
||||
import torch
|
||||
from accelerate.utils import set_seed
|
||||
import diffusers
|
||||
from diffusers import DDPMScheduler
|
||||
|
||||
import library.train_util as train_util
|
||||
from library.train_util import DreamBoothDataset, FineTuningDataset
|
||||
from library.train_util import (
|
||||
DreamBoothDataset,
|
||||
)
|
||||
import library.config_util as config_util
|
||||
from library.config_util import (
|
||||
ConfigSanitizer,
|
||||
BlueprintGenerator,
|
||||
)
|
||||
|
||||
|
||||
def collate_fn(examples):
|
||||
@ -49,6 +54,7 @@ def train(args):
|
||||
|
||||
cache_latents = args.cache_latents
|
||||
use_dreambooth_method = args.in_json is None
|
||||
use_user_config = args.dataset_config is not None
|
||||
|
||||
if args.seed is not None:
|
||||
set_seed(args.seed)
|
||||
@ -56,35 +62,47 @@ def train(args):
|
||||
tokenizer = train_util.load_tokenizer(args)
|
||||
|
||||
# データセットを準備する
|
||||
blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, True))
|
||||
if use_user_config:
|
||||
print(f"Load dataset config from {args.dataset_config}")
|
||||
user_config = config_util.load_user_config(args.dataset_config)
|
||||
ignored = ["train_data_dir", "reg_data_dir", "in_json"]
|
||||
if any(getattr(args, attr) is not None for attr in ignored):
|
||||
print(
|
||||
"ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(', '.join(ignored)))
|
||||
else:
|
||||
if use_dreambooth_method:
|
||||
print("Use DreamBooth method.")
|
||||
train_dataset = DreamBoothDataset(args.train_batch_size, args.train_data_dir, args.reg_data_dir,
|
||||
tokenizer, args.max_token_length, args.caption_extension, args.shuffle_caption, args.keep_tokens,
|
||||
args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
|
||||
args.bucket_reso_steps, args.bucket_no_upscale,
|
||||
args.prior_loss_weight, args.flip_aug, args.color_aug, args.face_crop_aug_range,
|
||||
args.random_crop, args.debug_dataset)
|
||||
user_config = {
|
||||
"datasets": [{
|
||||
"subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)
|
||||
}]
|
||||
}
|
||||
else:
|
||||
print("Train with captions.")
|
||||
train_dataset = FineTuningDataset(args.in_json, args.train_batch_size, args.train_data_dir,
|
||||
tokenizer, args.max_token_length, args.shuffle_caption, args.keep_tokens,
|
||||
args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
|
||||
args.bucket_reso_steps, args.bucket_no_upscale,
|
||||
args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop,
|
||||
args.dataset_repeats, args.debug_dataset)
|
||||
user_config = {
|
||||
"datasets": [{
|
||||
"subsets": [{
|
||||
"image_dir": args.train_data_dir,
|
||||
"metadata_file": args.in_json,
|
||||
}]
|
||||
}]
|
||||
}
|
||||
|
||||
# 学習データのdropout率を設定する
|
||||
train_dataset.set_caption_dropout(args.caption_dropout_rate, args.caption_dropout_every_n_epochs, args.caption_tag_dropout_rate)
|
||||
|
||||
train_dataset.make_buckets()
|
||||
blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
|
||||
train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
|
||||
|
||||
if args.debug_dataset:
|
||||
train_util.debug_dataset(train_dataset)
|
||||
train_util.debug_dataset(train_dataset_group)
|
||||
return
|
||||
if len(train_dataset) == 0:
|
||||
if len(train_dataset_group) == 0:
|
||||
print("No data found. Please verify arguments (train_data_dir must be the parent of folders with images) / 画像がありません。引数指定を確認してください（train_data_dirには画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）")
|
||||
return
|
||||
|
||||
if cache_latents:
|
||||
assert train_dataset_group.is_latent_cacheable(
|
||||
), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
|
||||
|
||||
# acceleratorを準備する
|
||||
print("prepare accelerator")
|
||||
accelerator, unwrap_model = train_util.prepare_accelerator(args)
|
||||
@ -109,7 +127,7 @@ def train(args):
|
||||
vae.requires_grad_(False)
|
||||
vae.eval()
|
||||
with torch.no_grad():
|
||||
train_dataset.cache_latents(vae)
|
||||
train_dataset_group.cache_latents(vae)
|
||||
vae.to("cpu")
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
@ -153,7 +171,7 @@ def train(args):
|
||||
# DataLoaderのプロセス数：0はメインプロセスになる
|
||||
n_workers = min(args.max_data_loader_n_workers, os.cpu_count() - 1) # cpu_count-1 ただし最大で指定された数まで
|
||||
train_dataloader = torch.utils.data.DataLoader(
|
||||
train_dataset, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
|
||||
train_dataset_group, batch_size=1, shuffle=True, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
|
||||
|
||||
# 学習ステップ数を計算する
|
||||
if args.max_train_epochs is not None:
|
||||
@ -231,17 +249,19 @@ def train(args):
|
||||
args.save_every_n_epochs = math.floor(num_train_epochs / args.save_n_epoch_ratio) or 1
|
||||
|
||||
# 学習する
|
||||
# TODO: find a way to handle total batch size when there are multiple datasets
|
||||
total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
|
||||
print("running training / 学習開始")
|
||||
print(f" num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset.num_train_images}")
|
||||
print(f" num reg images / 正則化画像の数: {train_dataset.num_reg_images}")
|
||||
print(f" num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
|
||||
print(f" num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
|
||||
print(f" num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
|
||||
print(f" num epochs / epoch数: {num_train_epochs}")
|
||||
print(f" batch size per device / バッチサイズ: {args.train_batch_size}")
|
||||
print(f" total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
|
||||
print(f" batch size per device / バッチサイズ: {', '.join([str(d.batch_size) for d in train_dataset_group.datasets])}")
|
||||
# print(f" total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
|
||||
print(f" gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
|
||||
print(f" total optimization steps / 学習ステップ数: {args.max_train_steps}")
|
||||
|
||||
# TODO refactor metadata creation and move to util
|
||||
metadata = {
|
||||
"ss_session_id": session_id, # random integer indicating which group of epochs the model came from
|
||||
"ss_training_started_at": training_started_at, # unix timestamp
|
||||
@ -249,12 +269,10 @@ def train(args):
|
||||
"ss_learning_rate": args.learning_rate,
|
||||
"ss_text_encoder_lr": args.text_encoder_lr,
|
||||
"ss_unet_lr": args.unet_lr,
|
||||
"ss_num_train_images": train_dataset.num_train_images, # includes repeating
|
||||
"ss_num_reg_images": train_dataset.num_reg_images,
|
||||
"ss_num_train_images": train_dataset_group.num_train_images,
|
||||
"ss_num_reg_images": train_dataset_group.num_reg_images,
|
||||
"ss_num_batches_per_epoch": len(train_dataloader),
|
||||
"ss_num_epochs": num_train_epochs,
|
||||
"ss_batch_size_per_device": args.train_batch_size,
|
||||
"ss_total_batch_size": total_batch_size,
|
||||
"ss_gradient_checkpointing": args.gradient_checkpointing,
|
||||
"ss_gradient_accumulation_steps": args.gradient_accumulation_steps,
|
||||
"ss_max_train_steps": args.max_train_steps,
|
||||
@ -266,26 +284,12 @@ def train(args):
|
||||
"ss_mixed_precision": args.mixed_precision,
|
||||
"ss_full_fp16": bool(args.full_fp16),
|
||||
"ss_v2": bool(args.v2),
|
||||
"ss_resolution": args.resolution,
|
||||
"ss_clip_skip": args.clip_skip,
|
||||
"ss_max_token_length": args.max_token_length,
|
||||
"ss_color_aug": bool(args.color_aug),
|
||||
"ss_flip_aug": bool(args.flip_aug),
|
||||
"ss_random_crop": bool(args.random_crop),
|
||||
"ss_shuffle_caption": bool(args.shuffle_caption),
|
||||
"ss_cache_latents": bool(args.cache_latents),
|
||||
"ss_enable_bucket": bool(train_dataset.enable_bucket),
|
||||
"ss_bucket_no_upscale": bool(train_dataset.bucket_no_upscale),
|
||||
"ss_min_bucket_reso": train_dataset.min_bucket_reso,
|
||||
"ss_max_bucket_reso": train_dataset.max_bucket_reso,
|
||||
"ss_seed": args.seed,
|
||||
"ss_lowram": args.lowram,
|
||||
"ss_keep_tokens": args.keep_tokens,
|
||||
"ss_noise_offset": args.noise_offset,
|
||||
"ss_dataset_dirs": json.dumps(train_dataset.dataset_dirs_info),
|
||||
"ss_reg_dataset_dirs": json.dumps(train_dataset.reg_dataset_dirs_info),
|
||||
"ss_tag_frequency": json.dumps(train_dataset.tag_frequency),
|
||||
"ss_bucket_info": json.dumps(train_dataset.bucket_info),
|
||||
"ss_training_comment": args.training_comment, # will not be updated after training
|
||||
"ss_sd_scripts_commit_hash": train_util.get_git_revision_hash(),
|
||||
"ss_optimizer": optimizer_name + (f"({optimizer_args})" if len(optimizer_args) > 0 else ""),
|
||||
@ -297,6 +301,132 @@ def train(args):
|
||||
"ss_prior_loss_weight": args.prior_loss_weight,
|
||||
}
|
||||
|
||||
if use_user_config:
|
||||
# save metadata of multiple datasets
|
||||
# NOTE: pack "ss_datasets" value as json one time
|
||||
# or should also pack nested collections as json?
|
||||
datasets_metadata = []
|
||||
tag_frequency = {} # merge tag frequency for metadata editor
|
||||
dataset_dirs_info = {} # merge subset dirs for metadata editor
|
||||
|
||||
for dataset in train_dataset_group.datasets:
|
||||
is_dreambooth_dataset = isinstance(dataset, DreamBoothDataset)
|
||||
dataset_metadata = {
|
||||
"is_dreambooth": is_dreambooth_dataset,
|
||||
"batch_size_per_device": dataset.batch_size,
|
||||
"num_train_images": dataset.num_train_images, # includes repeating
|
||||
"num_reg_images": dataset.num_reg_images,
|
||||
"resolution": (dataset.width, dataset.height),
|
||||
"enable_bucket": bool(dataset.enable_bucket),
|
||||
"min_bucket_reso": dataset.min_bucket_reso,
|
||||
"max_bucket_reso": dataset.max_bucket_reso,
|
||||
"tag_frequency": dataset.tag_frequency,
|
||||
"bucket_info": dataset.bucket_info,
|
||||
}
|
||||
|
||||
subsets_metadata = []
|
||||
for subset in dataset.subsets:
|
||||
subset_metadata = {
|
||||
"img_count": subset.img_count,
|
||||
"num_repeats": subset.num_repeats,
|
||||
"color_aug": bool(subset.color_aug),
|
||||
"flip_aug": bool(subset.flip_aug),
|
||||
"random_crop": bool(subset.random_crop),
|
||||
"shuffle_caption": bool(subset.shuffle_caption),
|
||||
"keep_tokens": subset.keep_tokens,
|
||||
}
|
||||
|
||||
image_dir_or_metadata_file = None
|
||||
if subset.image_dir:
|
||||
image_dir = os.path.basename(subset.image_dir)
|
||||
subset_metadata["image_dir"] = image_dir
|
||||
image_dir_or_metadata_file = image_dir
|
||||
|
||||
if is_dreambooth_dataset:
|
||||
subset_metadata["class_tokens"] = subset.class_tokens
|
||||
subset_metadata["is_reg"] = subset.is_reg
|
||||
if subset.is_reg:
|
||||
image_dir_or_metadata_file = None # not merging reg dataset
|
||||
else:
|
||||
metadata_file = os.path.basename(subset.metadata_file)
|
||||
subset_metadata["metadata_file"] = metadata_file
|
||||
image_dir_or_metadata_file = metadata_file # may overwrite
|
||||
|
||||
subsets_metadata.append(subset_metadata)
|
||||
|
||||
# merge dataset dir: not reg subset only
|
||||
# TODO update additional-network extension to show detailed dataset config from metadata
|
||||
if image_dir_or_metadata_file is not None:
|
||||
# datasets may have a certain dir multiple times
|
||||
v = image_dir_or_metadata_file
|
||||
i = 2
|
||||
while v in dataset_dirs_info:
|
||||
v = image_dir_or_metadata_file + f" ({i})"
|
||||
i += 1
|
||||
image_dir_or_metadata_file = v
|
||||
|
||||
dataset_dirs_info[image_dir_or_metadata_file] = {
|
||||
"n_repeats": subset.num_repeats,
|
||||
"img_count": subset.img_count
|
||||
}
|
||||
|
||||
dataset_metadata["subsets"] = subsets_metadata
|
||||
datasets_metadata.append(dataset_metadata)
|
||||
|
||||
# merge tag frequency:
|
||||
for ds_dir_name, ds_freq_for_dir in dataset.tag_frequency.items():
|
||||
# あるディレクトリが複数のdatasetで使用されている場合、一度だけ数える
|
||||
# もともと繰り返し回数を指定しているので、キャプション内でのタグの出現回数と、それが学習で何度使われるかは一致しない
|
||||
# なので、ここで複数datasetの回数を合算してもあまり意味はない
|
||||
if ds_dir_name in tag_frequency:
|
||||
continue
|
||||
tag_frequency[ds_dir_name] = ds_freq_for_dir
|
||||
|
||||
metadata["ss_datasets"] = json.dumps(datasets_metadata)
|
||||
metadata["ss_tag_frequency"] = json.dumps(tag_frequency)
|
||||
metadata["ss_dataset_dirs"] = json.dumps(dataset_dirs_info)
|
||||
else:
|
||||
# conserving backward compatibility when using train_dataset_dir and reg_dataset_dir
|
||||
assert len(
|
||||
train_dataset_group.datasets) == 1, f"There should be a single dataset but {len(train_dataset_group.datasets)} found. This seems to be a bug. / データセットは1個だけ存在するはずですが、実際には{len(train_dataset_group.datasets)}個でした。プログラムのバグかもしれません。"
|
||||
|
||||
dataset = train_dataset_group.datasets[0]
|
||||
|
||||
dataset_dirs_info = {}
|
||||
reg_dataset_dirs_info = {}
|
||||
if use_dreambooth_method:
|
||||
for subset in dataset.subsets:
|
||||
info = reg_dataset_dirs_info if subset.is_reg else dataset_dirs_info
|
||||
info[os.path.basename(subset.image_dir)] = {
|
||||
"n_repeats": subset.num_repeats,
|
||||
"img_count": subset.img_count
|
||||
}
|
||||
else:
|
||||
for subset in dataset.subsets:
|
||||
dataset_dirs_info[os.path.basename(subset.metadata_file)] = {
|
||||
"n_repeats": subset.num_repeats,
|
||||
"img_count": subset.img_count
|
||||
}
|
||||
|
||||
metadata.update({
|
||||
"ss_batch_size_per_device": args.train_batch_size,
|
||||
"ss_total_batch_size": total_batch_size,
|
||||
"ss_resolution": args.resolution,
|
||||
"ss_color_aug": bool(args.color_aug),
|
||||
"ss_flip_aug": bool(args.flip_aug),
|
||||
"ss_random_crop": bool(args.random_crop),
|
||||
"ss_shuffle_caption": bool(args.shuffle_caption),
|
||||
"ss_enable_bucket": bool(dataset.enable_bucket),
|
||||
"ss_bucket_no_upscale": bool(dataset.bucket_no_upscale),
|
||||
"ss_min_bucket_reso": dataset.min_bucket_reso,
|
||||
"ss_max_bucket_reso": dataset.max_bucket_reso,
|
||||
"ss_keep_tokens": args.keep_tokens,
|
||||
"ss_dataset_dirs": json.dumps(dataset_dirs_info),
|
||||
"ss_reg_dataset_dirs": json.dumps(reg_dataset_dirs_info),
|
||||
"ss_tag_frequency": json.dumps(dataset.tag_frequency),
|
||||
"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
|
||||
@ -332,7 +462,7 @@ def train(args):
|
||||
loss_total = 0.0
|
||||
for epoch in range(num_train_epochs):
|
||||
print(f"epoch {epoch+1}/{num_train_epochs}")
|
||||
train_dataset.set_current_epoch(epoch + 1)
|
||||
train_dataset_group.set_current_epoch(epoch + 1)
|
||||
|
||||
metadata["ss_epoch"] = str(epoch+1)
|
||||
|
||||
@ -400,6 +530,8 @@ def train(args):
|
||||
progress_bar.update(1)
|
||||
global_step += 1
|
||||
|
||||
train_util.sample_images(accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
|
||||
|
||||
current_loss = loss.detach().item()
|
||||
if epoch == 0:
|
||||
loss_list.append(current_loss)
|
||||
@ -445,6 +577,8 @@ def train(args):
|
||||
if saving and args.save_state:
|
||||
train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)
|
||||
|
||||
train_util.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
|
||||
|
||||
# end of epoch
|
||||
|
||||
metadata["ss_epoch"] = str(num_train_epochs)
|
||||
@ -480,6 +614,7 @@ if __name__ == '__main__':
|
||||
train_util.add_dataset_arguments(parser, True, True, True)
|
||||
train_util.add_training_arguments(parser, True)
|
||||
train_util.add_optimizer_arguments(parser)
|
||||
config_util.add_config_arguments(parser)
|
||||
|
||||
parser.add_argument("--no_metadata", action='store_true', help="do not save metadata in output model / メタデータを出力先モデルに保存しない")
|
||||
parser.add_argument("--save_model_as", type=str, default="safetensors", choices=[None, "ckpt", "pt", "safetensors"],
|
||||
|
||||
@ -11,7 +11,11 @@ import diffusers
|
||||
from diffusers import DDPMScheduler
|
||||
|
||||
import library.train_util as train_util
|
||||
from library.train_util import DreamBoothDataset, FineTuningDataset
|
||||
import library.config_util as config_util
|
||||
from library.config_util import (
|
||||
ConfigSanitizer,
|
||||
BlueprintGenerator,
|
||||
)
|
||||
|
||||
imagenet_templates_small = [
|
||||
"a photo of a {}",
|
||||
@ -79,7 +83,6 @@ def train(args):
|
||||
train_util.prepare_dataset_args(args, True)
|
||||
|
||||
cache_latents = args.cache_latents
|
||||
use_dreambooth_method = args.in_json is None
|
||||
|
||||
if args.seed is not None:
|
||||
set_seed(args.seed)
|
||||
@ -139,21 +142,35 @@ def train(args):
|
||||
print(f"create embeddings for {args.num_vectors_per_token} tokens, for {args.token_string}")
|
||||
|
||||
# データセットを準備する
|
||||
blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False))
|
||||
if args.dataset_config is not None:
|
||||
print(f"Load dataset config from {args.dataset_config}")
|
||||
user_config = config_util.load_user_config(args.dataset_config)
|
||||
ignored = ["train_data_dir", "reg_data_dir", "in_json"]
|
||||
if any(getattr(args, attr) is not None for attr in ignored):
|
||||
print("ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(', '.join(ignored)))
|
||||
else:
|
||||
use_dreambooth_method = args.in_json is None
|
||||
if use_dreambooth_method:
|
||||
print("Use DreamBooth method.")
|
||||
train_dataset = DreamBoothDataset(args.train_batch_size, args.train_data_dir, args.reg_data_dir,
|
||||
tokenizer, args.max_token_length, args.caption_extension, args.shuffle_caption, args.keep_tokens,
|
||||
args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
|
||||
args.bucket_reso_steps, args.bucket_no_upscale,
|
||||
args.prior_loss_weight, args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop, args.debug_dataset)
|
||||
user_config = {
|
||||
"datasets": [{
|
||||
"subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)
|
||||
}]
|
||||
}
|
||||
else:
|
||||
print("Train with captions.")
|
||||
train_dataset = FineTuningDataset(args.in_json, args.train_batch_size, args.train_data_dir,
|
||||
tokenizer, args.max_token_length, args.shuffle_caption, args.keep_tokens,
|
||||
args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
|
||||
args.bucket_reso_steps, args.bucket_no_upscale,
|
||||
args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop,
|
||||
args.dataset_repeats, args.debug_dataset)
|
||||
user_config = {
|
||||
"datasets": [{
|
||||
"subsets": [{
|
||||
"image_dir": args.train_data_dir,
|
||||
"metadata_file": args.in_json,
|
||||
}]
|
||||
}]
|
||||
}
|
||||
|
||||
blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
|
||||
train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
|
||||
|
||||
# make captions: tokenstring tokenstring1 tokenstring2 ...tokenstringn という文字列に書き換える超乱暴な実装
|
||||
if use_template:
|
||||
@ -163,20 +180,25 @@ def train(args):
|
||||
captions = []
|
||||
for tmpl in templates:
|
||||
captions.append(tmpl.format(replace_to))
|
||||
train_dataset.add_replacement("", captions)
|
||||
elif args.num_vectors_per_token > 1:
|
||||
train_dataset_group.add_replacement("", captions)
|
||||
else:
|
||||
if args.num_vectors_per_token > 1:
|
||||
replace_to = " ".join(token_strings)
|
||||
train_dataset.add_replacement(args.token_string, replace_to)
|
||||
|
||||
train_dataset.make_buckets()
|
||||
train_dataset_group.add_replacement(args.token_string, replace_to)
|
||||
prompt_replacement = (args.token_string, replace_to)
|
||||
else:
|
||||
prompt_replacement = None
|
||||
|
||||
if args.debug_dataset:
|
||||
train_util.debug_dataset(train_dataset, show_input_ids=True)
|
||||
train_util.debug_dataset(train_dataset_group, show_input_ids=True)
|
||||
return
|
||||
if len(train_dataset) == 0:
|
||||
if len(train_dataset_group) == 0:
|
||||
print("No data found. Please verify arguments / 画像がありません。引数指定を確認してください")
|
||||
return
|
||||
|
||||
if cache_latents:
|
||||
assert train_dataset_group.is_latent_cacheable(), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
|
||||
|
||||
# モデルに xformers とか memory efficient attention を組み込む
|
||||
train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
|
||||
|
||||
@ -186,7 +208,7 @@ def train(args):
|
||||
vae.requires_grad_(False)
|
||||
vae.eval()
|
||||
with torch.no_grad():
|
||||
train_dataset.cache_latents(vae)
|
||||
train_dataset_group.cache_latents(vae)
|
||||
vae.to("cpu")
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
@ -205,7 +227,7 @@ def train(args):
|
||||
# DataLoaderのプロセス数：0はメインプロセスになる
|
||||
n_workers = min(args.max_data_loader_n_workers, os.cpu_count() - 1) # cpu_count-1 ただし最大で指定された数まで
|
||||
train_dataloader = torch.utils.data.DataLoader(
|
||||
train_dataset, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
|
||||
train_dataset_group, batch_size=1, shuffle=True, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
|
||||
|
||||
# 学習ステップ数を計算する
|
||||
if args.max_train_epochs is not None:
|
||||
@ -263,8 +285,8 @@ def train(args):
|
||||
# 学習する
|
||||
total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
|
||||
print("running training / 学習開始")
|
||||
print(f" num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset.num_train_images}")
|
||||
print(f" num reg images / 正則化画像の数: {train_dataset.num_reg_images}")
|
||||
print(f" num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
|
||||
print(f" num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
|
||||
print(f" num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
|
||||
print(f" num epochs / epoch数: {num_train_epochs}")
|
||||
print(f" batch size per device / バッチサイズ: {args.train_batch_size}")
|
||||
@ -283,12 +305,11 @@ def train(args):
|
||||
|
||||
for epoch in range(num_train_epochs):
|
||||
print(f"epoch {epoch+1}/{num_train_epochs}")
|
||||
train_dataset.set_current_epoch(epoch + 1)
|
||||
train_dataset_group.set_current_epoch(epoch + 1)
|
||||
|
||||
text_encoder.train()
|
||||
|
||||
loss_total = 0
|
||||
bef_epo_embs = unwrap_model(text_encoder).get_input_embeddings().weight[token_ids].data.detach().clone()
|
||||
for step, batch in enumerate(train_dataloader):
|
||||
with accelerator.accumulate(text_encoder):
|
||||
with torch.no_grad():
|
||||
@ -354,6 +375,9 @@ def train(args):
|
||||
progress_bar.update(1)
|
||||
global_step += 1
|
||||
|
||||
train_util.sample_images(accelerator, args, None, global_step, accelerator.device,
|
||||
vae, tokenizer, text_encoder, unet, prompt_replacement)
|
||||
|
||||
current_loss = loss.detach().item()
|
||||
if args.logging_dir is not None:
|
||||
logs = {"loss": current_loss, "lr": float(lr_scheduler.get_last_lr()[0])}
|
||||
@ -376,8 +400,6 @@ def train(args):
|
||||
accelerator.wait_for_everyone()
|
||||
|
||||
updated_embs = unwrap_model(text_encoder).get_input_embeddings().weight[token_ids].data.detach().clone()
|
||||
# d = updated_embs - bef_epo_embs
|
||||
# print(bef_epo_embs.size(), updated_embs.size(), d.mean(), d.min())
|
||||
|
||||
if args.save_every_n_epochs is not None:
|
||||
model_name = train_util.DEFAULT_EPOCH_NAME if args.output_name is None else args.output_name
|
||||
@ -399,6 +421,9 @@ def train(args):
|
||||
if saving and args.save_state:
|
||||
train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)
|
||||
|
||||
train_util.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device,
|
||||
vae, tokenizer, text_encoder, unet, prompt_replacement)
|
||||
|
||||
# end of epoch
|
||||
|
||||
is_main_process = accelerator.is_main_process
|
||||
@ -474,6 +499,7 @@ if __name__ == '__main__':
|
||||
train_util.add_dataset_arguments(parser, True, True, False)
|
||||
train_util.add_training_arguments(parser, True)
|
||||
train_util.add_optimizer_arguments(parser)
|
||||
config_util.add_config_arguments(parser)
|
||||
|
||||
parser.add_argument("--save_model_as", type=str, default="pt", choices=[None, "ckpt", "pt", "safetensors"],
|
||||
help="format to save the model (default is .pt) / モデル保存時の形式（デフォルトはpt）")
|
||||
|
||||
Loading…
Reference in New Issue
Block a user