Merge branch 'dev' into dev

README.md

@@ -127,7 +127,33 @@ The majority of scripts is licensed under ASL 2.0 (including codes from Diffuser
 
 ## Change History
 
-- 9 Mar. 2023, 2023/3/9:
+- 11 Mar. 2023, 2023/3/11:
+    - Fix `svd_merge_lora.py` causes an error about the device.
+    - `svd_merge_lora.py` でデバイス関連のエラーが発生する不具合を修正しました。
+- 10 Mar. 2023, 2023/3/10: release v0.5.1
+  - Fix to LoRA modules in the model are same to the previous (before 0.5.0) if Conv2d-3x3 is disabled (no `conv_dim` arg, default).
+    - Conv2D with kernel size 1x1 in ResNet modules were accidentally included in v0.5.0.
+    - Trained models with v0.5.0 will work with Web UI's built-in LoRA and Additional Networks extension.
+  - Fix an issue that dim (rank) of LoRA module is limited to the in/out dimensions of the target Linear/Conv2d (in case of the dim > 320).
+  - `resize_lora.py` now have a feature to `dynamic resizing` which means each LoRA module can have different ranks (dims). Thanks to mgz-dev for this great work!
+    - The appropriate rank is selected based on the complexity of each module with an algorithm specified in the command line arguments. For details: https://github.com/kohya-ss/sd-scripts/pull/243
+  - Multiple GPUs training is finally supported in `train_network.py`. Thanks to ddPn08 to solve this long running issue!
+  - Dataset with fine-tuning method (with metadata json) now works without images if `.npz` files exist. Thanks to rvhfxb!
+  - `train_network.py` can work if the current directory is not the directory where the script is in. Thanks to mio2333!
+  - Fix `extract_lora_from_models.py` and `svd_merge_lora.py` doesn't work with higher rank (>320).
+
+  - LoRAのConv2d-3x3拡張を行わない場合（`conv_dim` を指定しない場合）、以前（v0.5.0）と同じ構成になるよう修正しました。
+    - ResNetのカーネルサイズ1x1のConv2dが誤って対象になっていました。
+    - ただv0.5.0で学習したモデルは Additional Networks 拡張、およびWeb UIのLoRA機能で問題なく使えると思われます。
+  - LoRAモジュールの dim (rank) が、対象モジュールの次元数以下に制限される不具合を修正しました（320より大きい dim を指定した場合）。
+  - `resize_lora.py` に `dynamic resizing` （リサイズ後の各LoRAモジュールが異なるrank (dim) を持てる機能）を追加しました。mgz-dev 氏の貢献に感謝します。
+    - 適切なランクがコマンドライン引数で指定したアルゴリズムにより自動的に選択されます。詳細はこちらをご覧ください: https://github.com/kohya-ss/sd-scripts/pull/243
+  - `train_network.py` でマルチGPU学習をサポートしました。長年の懸案を解決された ddPn08 氏に感謝します。
+  - fine-tuning方式のデータセット（メタデータ.jsonファイルを使うデータセット）で `.npz` が存在するときには画像がなくても動作するようになりました。rvhfxb 氏に感謝します。
+  - 他のディレクトリから `train_network.py` を呼び出しても動作するよう変更しました。 mio2333 氏に感謝します。
+  - `extract_lora_from_models.py` および `svd_merge_lora.py` が320より大きいrankを指定すると動かない不具合を修正しました。
+  
+- 9 Mar. 2023, 2023/3/9: release v0.5.0
   - There may be problems due to major changes. If you cannot revert back to the previous version when problems occur, please do not update for a while.
   - Minimum metadata (module name, dim, alpha and network_args) is recorded even with `--no_metadata`, issue https://github.com/kohya-ss/sd-scripts/issues/254
   - `train_network.py` supports LoRA for Conv2d-3x3 (extended to conv2d with a kernel size not 1x1).

fine_tune.py

@@ -5,6 +5,7 @@ import argparse
 import gc
 import math
 import os
+import toml
 
 from tqdm import tqdm
 import torch
@@ -19,6 +20,7 @@ from library.config_util import (
     BlueprintGenerator,
 )
 
+
 def collate_fn(examples):
     return examples[0]
 
@@ -40,15 +42,23 @@ def train(args):
         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)))
+            print(
+                "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
+                    ", ".join(ignored)
+                )
+            )
     else:
         user_config = {
-      "datasets": [{
+            "datasets": [
-        "subsets": [{
+                {
+                    "subsets": [
+                        {
                             "image_dir": args.train_data_dir,
                             "metadata_file": args.in_json,
-        }]
+                        }
-      }]
+                    ]
+                }
+            ]
         }
 
     blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
@@ -58,11 +68,15 @@ def train(args):
         train_util.debug_dataset(train_dataset_group)
         return
     if len(train_dataset_group) == 0:
-    print("No data found. Please verify the metadata file and train_data_dir option. / 画像がありません。メタデータおよびtrain_data_dirオプションを確認してください。")
+        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は使えません"
+        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")
@@ -86,7 +100,7 @@ def train(args):
         save_stable_diffusion_format = load_stable_diffusion_format
         use_safetensors = args.use_safetensors
     else:
-    save_stable_diffusion_format = args.save_model_as.lower() == 'ckpt' or args.save_model_as.lower() == 'safetensors'
+        save_stable_diffusion_format = args.save_model_as.lower() == "ckpt" or args.save_model_as.lower() == "safetensors"
         use_safetensors = args.use_safetensors or ("safetensors" in args.save_model_as.lower())
 
     # Diffusers版のxformers使用フラグを設定する関数
@@ -170,7 +184,13 @@ 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_group, batch_size=1, shuffle=True, 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:
@@ -182,7 +202,9 @@ def train(args):
 
     # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
     if args.full_fp16:
-    assert args.mixed_precision == "fp16", "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
+        assert (
+            args.mixed_precision == "fp16"
+        ), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
         print("enable full fp16 training.")
         unet.to(weight_dtype)
         text_encoder.to(weight_dtype)
@@ -190,7 +212,8 @@ def train(args):
     # acceleratorがなんかよろしくやってくれるらしい
     if args.train_text_encoder:
         unet, text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        unet, text_encoder, optimizer, train_dataloader, lr_scheduler)
+            unet, text_encoder, optimizer, train_dataloader, lr_scheduler
+        )
     else:
         unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, train_dataloader, lr_scheduler)
 
@@ -223,8 +246,9 @@ def train(args):
     progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
     global_step = 0
 
-  noise_scheduler = DDPMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
+    noise_scheduler = DDPMScheduler(
-                                  num_train_timesteps=1000, clip_sample=False)
+        beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
+    )
 
     if accelerator.is_main_process:
         accelerator.init_trackers("finetuning")
@@ -252,7 +276,8 @@ def train(args):
                     # Get the text embedding for conditioning
                     input_ids = batch["input_ids"].to(accelerator.device)
                     encoder_hidden_states = train_util.get_hidden_states(
-              args, input_ids, tokenizer, text_encoder, None if not args.full_fp16 else weight_dtype)
+                        args, input_ids, tokenizer, text_encoder, None if not args.full_fp16 else weight_dtype
+                    )
 
                 # Sample noise that we'll add to the latents
                 noise = torch.randn_like(latents, device=latents.device)
@@ -295,13 +320,17 @@ 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)
+                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])}
                 if args.optimizer_type.lower() == "DAdaptation".lower():  # tracking d*lr value
-          logs["lr/d*lr"] = lr_scheduler.optimizers[0].param_groups[0]['d']*lr_scheduler.optimizers[0].param_groups[0]['lr']
+                    logs["lr/d*lr"] = (
+                        lr_scheduler.optimizers[0].param_groups[0]["d"] * lr_scheduler.optimizers[0].param_groups[0]["lr"]
+                    )
                 accelerator.log(logs, step=global_step)
 
             # TODO moving averageにする
@@ -321,8 +350,20 @@ def train(args):
 
         if args.save_every_n_epochs is not None:
             src_path = src_stable_diffusion_ckpt if save_stable_diffusion_format else src_diffusers_model_path
-      train_util.save_sd_model_on_epoch_end(args, accelerator, src_path, save_stable_diffusion_format, use_safetensors,
+            train_util.save_sd_model_on_epoch_end(
-                                            save_dtype, epoch, num_train_epochs, global_step,  unwrap_model(text_encoder), unwrap_model(unet), vae)
+                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)
 
@@ -340,12 +381,13 @@ def train(args):
 
     if is_main_process:
         src_path = src_stable_diffusion_ckpt if save_stable_diffusion_format else src_diffusers_model_path
-    train_util.save_sd_model_on_train_end(args, src_path, save_stable_diffusion_format, use_safetensors,
+        train_util.save_sd_model_on_train_end(
-                                          save_dtype, epoch, global_step,  text_encoder, unet, vae)
+            args, src_path, save_stable_diffusion_format, use_safetensors, save_dtype, epoch, global_step, text_encoder, unet, vae
+        )
         print("model saved.")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
 
     train_util.add_sd_models_arguments(parser)
@@ -355,9 +397,10 @@ if __name__ == '__main__':
     train_util.add_optimizer_arguments(parser)
     config_util.add_config_arguments(parser)
 
-  parser.add_argument("--diffusers_xformers", action='store_true',
+    parser.add_argument("--diffusers_xformers", action="store_true", help="use xformers by diffusers / Diffusersでxformersを使用する")
-                      help='use xformers by diffusers / Diffusersでxformersを使用する')
     parser.add_argument("--train_text_encoder", action="store_true", help="train text encoder / text encoderも学習する")
 
     args = parser.parse_args()
+    args = train_util.read_config_from_file(args, parser)
+
     train(args)

finetune/merge_captions_to_metadata.py

@@ -4,7 +4,7 @@ from pathlib import Path
 from typing import List
 from tqdm import tqdm
 import library.train_util as train_util
-
+import os
 
 def main(args):
   assert not args.recursive or (args.recursive and args.full_path), "recursive requires full_path / recursiveはfull_pathと同時に指定してください"
@@ -29,6 +29,9 @@ def main(args):
     caption_path = image_path.with_suffix(args.caption_extension)
     caption = caption_path.read_text(encoding='utf-8').strip()
 
+    if not os.path.exists(caption_path):
+      caption_path = os.path.join(image_path, args.caption_extension)
+
     image_key = str(image_path) if args.full_path else image_path.stem
     if image_key not in metadata:
       metadata[image_key] = {}

finetune/merge_dd_tags_to_metadata.py

@@ -4,7 +4,7 @@ from pathlib import Path
 from typing import List
 from tqdm import tqdm
 import library.train_util as train_util
-
+import os
 
 def main(args):
   assert not args.recursive or (args.recursive and args.full_path), "recursive requires full_path / recursiveはfull_pathと同時に指定してください"
@@ -29,6 +29,9 @@ def main(args):
     tags_path = image_path.with_suffix(args.caption_extension)
     tags = tags_path.read_text(encoding='utf-8').strip()
 
+    if not os.path.exists(tags_path):
+      tags_path = os.path.join(image_path, args.caption_extension)
+
     image_key = str(image_path) if args.full_path else image_path.stem
     if image_key not in metadata:
       metadata[image_key] = {}

networks/extract_lora_from_models.py

@@ -103,7 +103,8 @@ def svd(args):
 
       if args.device:
         mat = mat.to(args.device)
-      # print(mat.size(), mat.device, rank, in_dim, out_dim)
+
+      # print(lora_name, mat.size(), mat.device, rank, in_dim, out_dim)
       rank = min(rank, in_dim, out_dim)                           # LoRA rank cannot exceed the original dim
 
       if conv2d:
@@ -137,27 +138,17 @@ def svd(args):
       lora_weights[lora_name] = (U, Vh)
 
   # make state dict for LoRA
-  lora_network_o.apply_to(text_encoder_o, unet_o, text_encoder_different, True)   # to make state dict
+  lora_sd = {}
-  lora_sd = lora_network_o.state_dict()
+  for lora_name, (up_weight, down_weight) in lora_weights.items():
-  print(f"LoRA has {len(lora_sd)} weights.")
+    lora_sd[lora_name + '.lora_up.weight'] = up_weight
-
+    lora_sd[lora_name + '.lora_down.weight'] = down_weight
-  for key in list(lora_sd.keys()):
+    lora_sd[lora_name + '.alpha'] = torch.tensor(down_weight.size()[0])
-    if "alpha" in key:
-      continue
-
-    lora_name = key.split('.')[0]
-    i = 0 if "lora_up" in key else 1
-
-    weights = lora_weights[lora_name][i]
-    # print(key, i, weights.size(), lora_sd[key].size())
-    # if len(lora_sd[key].size()) == 4:
-    #   weights = weights.unsqueeze(2).unsqueeze(3)
-
-    assert weights.size() == lora_sd[key].size(), f"size unmatch: {key}"
-    lora_sd[key] = weights
 
   # load state dict to LoRA and save it
-  info = lora_network_o.load_state_dict(lora_sd)
+  lora_network_save = lora.create_network_from_weights(1.0, None, None, text_encoder_o, unet_o, weights_sd=lora_sd)
+  lora_network_save.apply_to(text_encoder_o, unet_o)        # create internal module references for state_dict
+
+  info = lora_network_save.load_state_dict(lora_sd)
   print(f"Loading extracted LoRA weights: {info}")
 
   dir_name = os.path.dirname(args.save_to)
@@ -167,7 +158,7 @@ def svd(args):
   # minimum metadata
   metadata = {"ss_network_module": "networks.lora", "ss_network_dim": str(args.dim), "ss_network_alpha": str(args.dim)}
 
-  lora_network_o.save_weights(args.save_to, save_dtype, metadata)
+  lora_network_save.save_weights(args.save_to, save_dtype, metadata)
   print(f"LoRA weights are saved to: {args.save_to}")
 
 

networks/lora.py

@@ -21,30 +21,34 @@ class LoRAModule(torch.nn.Module):
     """ if alpha == 0 or None, alpha is rank (no scaling). """
     super().__init__()
     self.lora_name = lora_name
-    self.lora_dim = lora_dim
 
     if org_module.__class__.__name__ == 'Conv2d':
       in_dim = org_module.in_channels
       out_dim = org_module.out_channels
+    else:
+      in_dim = org_module.in_features
+      out_dim = org_module.out_features
 
-      self.lora_dim = min(self.lora_dim, in_dim, out_dim)
+    # if limit_rank:
-      if self.lora_dim != lora_dim:
+    #   self.lora_dim = min(lora_dim, in_dim, out_dim)
-        print(f"{lora_name} dim (rank) is changed to: {self.lora_dim}")
+    #   if self.lora_dim != lora_dim:
+    #     print(f"{lora_name} dim (rank) is changed to: {self.lora_dim}")
+    # else:
+    self.lora_dim = lora_dim
 
+    if org_module.__class__.__name__ == 'Conv2d':
       kernel_size = org_module.kernel_size
       stride = org_module.stride
       padding = org_module.padding
       self.lora_down = torch.nn.Conv2d(in_dim, self.lora_dim, kernel_size, stride, padding, bias=False)
       self.lora_up = torch.nn.Conv2d(self.lora_dim, out_dim, (1, 1), (1, 1), bias=False)
     else:
-      in_dim = org_module.in_features
+      self.lora_down = torch.nn.Linear(in_dim, self.lora_dim, bias=False)
-      out_dim = org_module.out_features
+      self.lora_up = torch.nn.Linear(self.lora_dim, out_dim, bias=False)
-      self.lora_down = torch.nn.Linear(in_dim, lora_dim, bias=False)
-      self.lora_up = torch.nn.Linear(lora_dim, out_dim, bias=False)
 
     if type(alpha) == torch.Tensor:
       alpha = alpha.detach().float().numpy()                              # without casting, bf16 causes error
-    alpha = lora_dim if alpha is None or alpha == 0 else alpha
+    alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
     self.scale = alpha / self.lora_dim
     self.register_buffer('alpha', torch.tensor(alpha))                    # 定数として扱える
 
@@ -149,7 +153,8 @@ def create_network(multiplier, network_dim, network_alpha, vae, text_encoder, un
   return network
 
 
-def create_network_from_weights(multiplier, file, vae, text_encoder, unet, **kwargs):
+def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weights_sd=None, **kwargs):
+  if weights_sd is None:
     if os.path.splitext(file)[1] == '.safetensors':
       from safetensors.torch import load_file, safe_open
       weights_sd = load_file(file)
@@ -183,7 +188,8 @@ def create_network_from_weights(multiplier, file, vae, text_encoder, unet, **kwa
 
 class LoRANetwork(torch.nn.Module):
   # is it possible to apply conv_in and conv_out?
-  UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention", "ResnetBlock2D", "Downsample2D", "Upsample2D"]
+  UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]
+  UNET_TARGET_REPLACE_MODULE_CONV2D_3X3 = ["ResnetBlock2D", "Downsample2D", "Upsample2D"]
   TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
   LORA_PREFIX_UNET = 'lora_unet'
   LORA_PREFIX_TEXT_ENCODER = 'lora_te'
@@ -245,7 +251,12 @@ class LoRANetwork(torch.nn.Module):
                                              text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
     print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
 
-    self.unet_loras = create_modules(LoRANetwork.LORA_PREFIX_UNET, unet, LoRANetwork.UNET_TARGET_REPLACE_MODULE)
+    # extend U-Net target modules if conv2d 3x3 is enabled, or load from weights
+    target_modules = LoRANetwork.UNET_TARGET_REPLACE_MODULE
+    if modules_dim is not None or self.conv_lora_dim is not None:
+      target_modules += LoRANetwork.UNET_TARGET_REPLACE_MODULE_CONV2D_3X3
+
+    self.unet_loras = create_modules(LoRANetwork.LORA_PREFIX_UNET, unet, target_modules)
     print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
 
     self.weights_sd = None

networks/resize_lora.py

@@ -1,6 +1,6 @@
 # Convert LoRA to different rank approximation (should only be used to go to lower rank)
 # This code is based off the extract_lora_from_models.py file which is based on https://github.com/cloneofsimo/lora/blob/develop/lora_diffusion/cli_svd.py
-# Thanks to cloneofsimo and kohya
+# Thanks to cloneofsimo
 
 import argparse
 import torch

networks/svd_merge_lora.py

@@ -23,16 +23,16 @@ def load_state_dict(file_name, dtype):
   return sd
 
 
-def save_to_file(file_name, model, state_dict, dtype):
+def save_to_file(file_name, state_dict, dtype):
   if dtype is not None:
     for key in list(state_dict.keys()):
       if type(state_dict[key]) == torch.Tensor:
         state_dict[key] = state_dict[key].to(dtype)
 
   if os.path.splitext(file_name)[1] == '.safetensors':
-    save_file(model, file_name)
+    save_file(state_dict, file_name)
   else:
-    torch.save(model, file_name)
+    torch.save(state_dict, file_name)
 
 
 def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dtype):
@@ -77,6 +77,10 @@ def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dty
 
       # W <- W + U * D
       scale = (alpha / network_dim)
+
+      if device:                      # and isinstance(scale, torch.Tensor):
+        scale = scale.to(device)
+
       if not conv2d:        # linear
         weight = weight + ratio * (up_weight @ down_weight) * scale
       elif kernel_size == (1, 1):
@@ -105,6 +109,7 @@ def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dty
           mat = mat.squeeze()
 
       module_new_rank = new_conv_rank if conv2d_3x3 else new_rank
+      module_new_rank = min(module_new_rank, in_dim, out_dim)                           # LoRA rank cannot exceed the original dim
 
       U, S, Vh = torch.linalg.svd(mat)
 
@@ -156,7 +161,7 @@ def merge(args):
   state_dict = merge_lora_models(args.models, args.ratios, args.new_rank, new_conv_rank, args.device, merge_dtype)
 
   print(f"saving model to: {args.save_to}")
-  save_to_file(args.save_to, state_dict, state_dict, save_dtype)
+  save_to_file(args.save_to, state_dict, save_dtype)
 
 
 if __name__ == '__main__':

train_README-ja.md

@@ -502,6 +502,14 @@ masterpiece, best quality, 1boy, in business suit, standing at street, looking b
 
     clip_skipと同様に、モデルの学習状態と異なる長さで学習するには、ある程度の教師データ枚数、長めの学習時間が必要になると思われます。
 
+- `--persistent_data_loader_workers`
+
+    Windows環境で指定するとエポック間の待ち時間が大幅に短縮されます。
+
+- `--max_data_loader_n_workers`
+
+    データ読み込みのプロセス数を指定します。プロセス数が多いとデータ読み込みが速くなりGPUを効率的に利用できますが、メインメモリを消費します。デフォルトは「`8` または `CPU同時実行スレッド数-1` の小さいほう」なので、メインメモリに余裕がない場合や、GPU使用率が90%程度以上なら、それらの数値を見ながら `2` または `1` 程度まで下げてください。
+
 - `--logging_dir` / `--log_prefix`
 
     学習ログの保存に関するオプションです。logging_dirオプションにログ保存先フォルダを指定してください。TensorBoard形式のログが保存されます。

train_db.py

@@ -7,6 +7,7 @@ import argparse
 import itertools
 import math
 import os
+import toml
 
 from tqdm import tqdm
 import torch
@@ -43,12 +44,16 @@ def train(args):
         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)))
+            print(
+                "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
+                    ", ".join(ignored)
+                )
+            )
     else:
         user_config = {
-      "datasets": [{
+            "datasets": [
-        "subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)
+                {"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)
@@ -62,15 +67,20 @@ def train(args):
         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は使えません"
+        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")
 
     if args.gradient_accumulation_steps > 1:
-    print(f"gradient_accumulation_steps is {args.gradient_accumulation_steps}. accelerate does not support gradient_accumulation_steps when training multiple models (U-Net and Text Encoder), so something might be wrong")
         print(
-        f"gradient_accumulation_stepsが{args.gradient_accumulation_steps}に設定されています。accelerateは複数モデル（U-NetおよびText Encoder）の学習時にgradient_accumulation_stepsをサポートしていないため結果は未知数です")
+            f"gradient_accumulation_steps is {args.gradient_accumulation_steps}. accelerate does not support gradient_accumulation_steps when training multiple models (U-Net and Text Encoder), so something might be wrong"
+        )
+        print(
+            f"gradient_accumulation_stepsが{args.gradient_accumulation_steps}に設定されています。accelerateは複数モデル（U-NetおよびText Encoder）の学習時にgradient_accumulation_stepsをサポートしていないため結果は未知数です"
+        )
 
     accelerator, unwrap_model = train_util.prepare_accelerator(args)
 
@@ -92,7 +102,7 @@ def train(args):
         save_stable_diffusion_format = load_stable_diffusion_format
         use_safetensors = args.use_safetensors
     else:
-    save_stable_diffusion_format = args.save_model_as.lower() == 'ckpt' or args.save_model_as.lower() == 'safetensors'
+        save_stable_diffusion_format = args.save_model_as.lower() == "ckpt" or args.save_model_as.lower() == "safetensors"
         use_safetensors = args.use_safetensors or ("safetensors" in args.save_model_as.lower())
 
     # モデルに xformers とか memory efficient attention を組み込む
@@ -129,7 +139,7 @@ def train(args):
     # 学習に必要なクラスを準備する
     print("prepare optimizer, data loader etc.")
     if train_text_encoder:
-    trainable_params = (itertools.chain(unet.parameters(), text_encoder.parameters()))
+        trainable_params = itertools.chain(unet.parameters(), text_encoder.parameters())
     else:
         trainable_params = unet.parameters()
 
@@ -139,7 +149,13 @@ 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_group, batch_size=1, shuffle=True, 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:
@@ -154,7 +170,9 @@ def train(args):
 
     # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
     if args.full_fp16:
-    assert args.mixed_precision == "fp16", "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
+        assert (
+            args.mixed_precision == "fp16"
+        ), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
         print("enable full fp16 training.")
         unet.to(weight_dtype)
         text_encoder.to(weight_dtype)
@@ -162,7 +180,8 @@ def train(args):
     # acceleratorがなんかよろしくやってくれるらしい
     if train_text_encoder:
         unet, text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        unet, text_encoder, optimizer, train_dataloader, lr_scheduler)
+            unet, text_encoder, optimizer, train_dataloader, lr_scheduler
+        )
     else:
         unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, train_dataloader, lr_scheduler)
 
@@ -199,8 +218,9 @@ def train(args):
     progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
     global_step = 0
 
-  noise_scheduler = DDPMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
+    noise_scheduler = DDPMScheduler(
-                                  num_train_timesteps=1000, clip_sample=False)
+        beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
+    )
 
     if accelerator.is_main_process:
         accelerator.init_trackers("dreambooth")
@@ -245,7 +265,8 @@ def train(args):
                 with torch.set_grad_enabled(global_step < args.stop_text_encoder_training):
                     input_ids = batch["input_ids"].to(accelerator.device)
                     encoder_hidden_states = train_util.get_hidden_states(
-              args, input_ids, tokenizer, text_encoder, None if not args.full_fp16 else weight_dtype)
+                        args, input_ids, tokenizer, text_encoder, None if not args.full_fp16 else weight_dtype
+                    )
 
                 # Sample a random timestep for each image
                 timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (b_size,), device=latents.device)
@@ -275,7 +296,7 @@ def train(args):
                 accelerator.backward(loss)
                 if accelerator.sync_gradients and args.max_grad_norm != 0.0:
                     if train_text_encoder:
-            params_to_clip = (itertools.chain(unet.parameters(), text_encoder.parameters()))
+                        params_to_clip = itertools.chain(unet.parameters(), text_encoder.parameters())
                     else:
                         params_to_clip = unet.parameters()
                     accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
@@ -289,13 +310,17 @@ 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)
+                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])}
                 if args.optimizer_type.lower() == "DAdaptation".lower():  # tracking d*lr value
-          logs["lr/d*lr"] = lr_scheduler.optimizers[0].param_groups[0]['d']*lr_scheduler.optimizers[0].param_groups[0]['lr']
+                    logs["lr/d*lr"] = (
+                        lr_scheduler.optimizers[0].param_groups[0]["d"] * lr_scheduler.optimizers[0].param_groups[0]["lr"]
+                    )
                 accelerator.log(logs, step=global_step)
 
             if epoch == 0:
@@ -319,8 +344,20 @@ def train(args):
 
         if args.save_every_n_epochs is not None:
             src_path = src_stable_diffusion_ckpt if save_stable_diffusion_format else src_diffusers_model_path
-      train_util.save_sd_model_on_epoch_end(args, accelerator, src_path, save_stable_diffusion_format, use_safetensors,
+            train_util.save_sd_model_on_epoch_end(
-                                            save_dtype, epoch, num_train_epochs, global_step,  unwrap_model(text_encoder), unwrap_model(unet), vae)
+                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)
 
@@ -338,12 +375,13 @@ def train(args):
 
     if is_main_process:
         src_path = src_stable_diffusion_ckpt if save_stable_diffusion_format else src_diffusers_model_path
-    train_util.save_sd_model_on_train_end(args, src_path, save_stable_diffusion_format, use_safetensors,
+        train_util.save_sd_model_on_train_end(
-                                          save_dtype, epoch, global_step,  text_encoder, unet, vae)
+            args, src_path, save_stable_diffusion_format, use_safetensors, save_dtype, epoch, global_step, text_encoder, unet, vae
+        )
         print("model saved.")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
 
     train_util.add_sd_models_arguments(parser)
@@ -353,10 +391,19 @@ if __name__ == '__main__':
     train_util.add_optimizer_arguments(parser)
     config_util.add_config_arguments(parser)
 
-  parser.add_argument("--no_token_padding", action="store_true",
+    parser.add_argument(
-                      help="disable token padding (same as Diffuser's DreamBooth) / トークンのpaddingを無効にする（Diffusers版DreamBoothと同じ動作）")
+        "--no_token_padding",
-  parser.add_argument("--stop_text_encoder_training", type=int, default=None,
+        action="store_true",
-                      help="steps to stop text encoder training, -1 for no training / Text Encoderの学習を止めるステップ数、-1で最初から学習しない")
+        help="disable token padding (same as Diffuser's DreamBooth) / トークンのpaddingを無効にする（Diffusers版DreamBoothと同じ動作）",
+    )
+    parser.add_argument(
+        "--stop_text_encoder_training",
+        type=int,
+        default=None,
+        help="steps to stop text encoder training, -1 for no training / Text Encoderの学習を止めるステップ数、-1で最初から学習しない",
+    )
 
     args = parser.parse_args()
+    args = train_util.read_config_from_file(args, parser)
+
     train(args)

train_network.py

@@ -7,6 +7,7 @@ import os
 import random
 import time
 import json
+import toml
 
 from tqdm import tqdm
 import torch
@@ -41,7 +42,7 @@ def generate_step_logs(args: argparse.Namespace, current_loss, avr_loss, lr_sche
         logs["lr/unet"] = float(lr_scheduler.get_last_lr()[-1])  # may be same to textencoder
 
     if args.optimizer_type.lower() == "DAdaptation".lower():  # tracking d*lr value of unet.
-    logs["lr/d*lr"] = lr_scheduler.optimizers[-1].param_groups[0]['d']*lr_scheduler.optimizers[-1].param_groups[0]['lr']
+        logs["lr/d*lr"] = lr_scheduler.optimizers[-1].param_groups[0]["d"] * lr_scheduler.optimizers[-1].param_groups[0]["lr"]
 
     return logs
 
@@ -69,24 +70,31 @@ def train(args):
         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)))
+                "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
+                    ", ".join(ignored)
+                )
+            )
     else:
         if use_dreambooth_method:
             print("Use DreamBooth method.")
             user_config = {
-          "datasets": [{
+                "datasets": [
-              "subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)
+                    {"subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)}
-          }]
+                ]
             }
         else:
             print("Train with captions.")
             user_config = {
-          "datasets": [{
+                "datasets": [
-              "subsets": [{
+                    {
+                        "subsets": [
+                            {
                                 "image_dir": args.train_data_dir,
                                 "metadata_file": args.in_json,
-              }]
+                            }
-          }]
+                        ]
+                    }
+                ]
             }
 
     blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
@@ -96,11 +104,14 @@ def train(args):
         train_util.debug_dataset(train_dataset_group)
         return
     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には画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）")
+        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(
+        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を準備する
@@ -135,13 +146,16 @@ def train(args):
         gc.collect()
 
     # prepare network
+    import sys
+
+    sys.path.append(os.path.dirname(__file__))
     print("import network module:", args.network_module)
     network_module = importlib.import_module(args.network_module)
 
     net_kwargs = {}
     if args.network_args is not None:
         for net_arg in args.network_args:
-      key, value = net_arg.split('=')
+            key, value = net_arg.split("=")
             net_kwargs[key] = value
 
     # if a new network is added in future, add if ~ then blocks for each network (;'∀')
@@ -172,7 +186,13 @@ 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_group, batch_size=1, shuffle=True, 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:
@@ -185,23 +205,27 @@ def train(args):
 
     # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
     if args.full_fp16:
-    assert args.mixed_precision == "fp16", "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
+        assert (
+            args.mixed_precision == "fp16"
+        ), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
         print("enable full fp16 training.")
         network.to(weight_dtype)
 
     # acceleratorがなんかよろしくやってくれるらしい
     if train_unet and train_text_encoder:
         unet, text_encoder, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        unet, text_encoder, network, optimizer, train_dataloader, lr_scheduler)
+            unet, text_encoder, network, optimizer, train_dataloader, lr_scheduler
+        )
     elif train_unet:
         unet, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        unet,  network, optimizer, train_dataloader, lr_scheduler)
+            unet, network, optimizer, train_dataloader, lr_scheduler
+        )
     elif train_text_encoder:
         text_encoder, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        text_encoder, network, optimizer, train_dataloader, lr_scheduler)
+            text_encoder, network, optimizer, train_dataloader, lr_scheduler
+        )
     else:
-    network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+        network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(network, optimizer, train_dataloader, lr_scheduler)
-        network, optimizer, train_dataloader, lr_scheduler)
 
     unet.requires_grad_(False)
     unet.to(accelerator.device, dtype=weight_dtype)
@@ -367,10 +391,7 @@ def train(args):
                         i += 1
                     image_dir_or_metadata_file = v
 
-          dataset_dirs_info[image_dir_or_metadata_file] = {
+                    dataset_dirs_info[image_dir_or_metadata_file] = {"n_repeats": subset.num_repeats, "img_count": subset.img_count}
-              "n_repeats": subset.num_repeats,
-              "img_count": subset.img_count
-          }
 
             dataset_metadata["subsets"] = subsets_metadata
             datasets_metadata.append(dataset_metadata)
@@ -389,8 +410,9 @@ def train(args):
         metadata["ss_dataset_dirs"] = json.dumps(dataset_dirs_info)
     else:
         # conserving backward compatibility when using train_dataset_dir and reg_dataset_dir
-    assert len(
+        assert (
-        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)}個でした。プログラムのバグかもしれません。"
+            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]
 
@@ -399,18 +421,16 @@ def train(args):
         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)] = {
+                info[os.path.basename(subset.image_dir)] = {"n_repeats": subset.num_repeats, "img_count": subset.img_count}
-            "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
+                    "img_count": subset.img_count,
                 }
 
-    metadata.update({
+        metadata.update(
+            {
                 "ss_batch_size_per_device": args.train_batch_size,
                 "ss_total_batch_size": total_batch_size,
                 "ss_resolution": args.resolution,
@@ -427,7 +447,8 @@ def train(args):
                 "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),
-    })
+            }
+        )
 
     # add extra args
     if args.network_args:
@@ -464,8 +485,9 @@ def train(args):
     progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
     global_step = 0
 
-  noise_scheduler = DDPMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
+    noise_scheduler = DDPMScheduler(
-                                  num_train_timesteps=1000, clip_sample=False)
+        beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
+    )
 
     if accelerator.is_main_process:
         accelerator.init_trackers("network_train")
@@ -543,7 +565,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)
+                train_util.sample_images(
+                    accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet
+                )
 
             current_loss = loss.detach().item()
             if epoch == 0:
@@ -573,14 +597,14 @@ def train(args):
             model_name = train_util.DEFAULT_EPOCH_NAME if args.output_name is None else args.output_name
 
             def save_func():
-        ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, epoch + 1) + '.' + args.save_model_as
+                ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, epoch + 1) + "." + args.save_model_as
                 ckpt_file = os.path.join(args.output_dir, ckpt_name)
                 metadata["ss_training_finished_at"] = str(time.time())
                 print(f"saving checkpoint: {ckpt_file}")
                 unwrap_model(network).save_weights(ckpt_file, save_dtype, minimum_metadata if args.no_metadata else metadata)
 
             def remove_old_func(old_epoch_no):
-        old_ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, old_epoch_no) + '.' + args.save_model_as
+                old_ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, old_epoch_no) + "." + args.save_model_as
                 old_ckpt_file = os.path.join(args.output_dir, old_ckpt_name)
                 if os.path.exists(old_ckpt_file):
                     print(f"removing old checkpoint: {old_ckpt_file}")
@@ -612,7 +636,7 @@ def train(args):
         os.makedirs(args.output_dir, exist_ok=True)
 
         model_name = train_util.DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name
-    ckpt_name = model_name + '.' + args.save_model_as
+        ckpt_name = model_name + "." + args.save_model_as
         ckpt_file = os.path.join(args.output_dir, ckpt_name)
 
         print(f"save trained model to {ckpt_file}")
@@ -620,7 +644,7 @@ def train(args):
         print("model saved.")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
 
     train_util.add_sd_models_arguments(parser)
@@ -629,27 +653,41 @@ if __name__ == '__main__':
     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("--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"],
+    parser.add_argument(
-                      help="format to save the model (default is .safetensors) / モデル保存時の形式（デフォルトはsafetensors）")
+        "--save_model_as",
+        type=str,
+        default="safetensors",
+        choices=[None, "ckpt", "pt", "safetensors"],
+        help="format to save the model (default is .safetensors) / モデル保存時の形式（デフォルトはsafetensors）",
+    )
 
     parser.add_argument("--unet_lr", type=float, default=None, help="learning rate for U-Net / U-Netの学習率")
     parser.add_argument("--text_encoder_lr", type=float, default=None, help="learning rate for Text Encoder / Text Encoderの学習率")
 
-  parser.add_argument("--network_weights", type=str, default=None,
+    parser.add_argument("--network_weights", type=str, default=None, help="pretrained weights for network / 学習するネットワークの初期重み")
-                      help="pretrained weights for network / 学習するネットワークの初期重み")
+    parser.add_argument("--network_module", type=str, default=None, help="network module to train / 学習対象のネットワークのモジュール")
-  parser.add_argument("--network_module", type=str, default=None, help='network module to train / 学習対象のネットワークのモジュール')
+    parser.add_argument(
-  parser.add_argument("--network_dim", type=int, default=None,
+        "--network_dim", type=int, default=None, help="network dimensions (depends on each network) / モジュールの次元数（ネットワークにより定義は異なります）"
-                      help='network dimensions (depends on each network) / モジュールの次元数（ネットワークにより定義は異なります）')
+    )
-  parser.add_argument("--network_alpha", type=float, default=1,
+    parser.add_argument(
-                      help='alpha for LoRA weight scaling, default 1 (same as network_dim for same behavior as old version) / LoRaの重み調整のalpha値、デフォルト1（旧バージョンと同じ動作をするにはnetwork_dimと同じ値を指定）')
+        "--network_alpha",
-  parser.add_argument("--network_args", type=str, default=None, nargs='*',
+        type=float,
-                      help='additional argmuments for network (key=value) / ネットワークへの追加の引数')
+        default=1,
+        help="alpha for LoRA weight scaling, default 1 (same as network_dim for same behavior as old version) / LoRaの重み調整のalpha値、デフォルト1（旧バージョンと同じ動作をするにはnetwork_dimと同じ値を指定）",
+    )
+    parser.add_argument(
+        "--network_args", type=str, default=None, nargs="*", help="additional argmuments for network (key=value) / ネットワークへの追加の引数"
+    )
     parser.add_argument("--network_train_unet_only", action="store_true", help="only training U-Net part / U-Net関連部分のみ学習する")
-  parser.add_argument("--network_train_text_encoder_only", action="store_true",
+    parser.add_argument(
-                      help="only training Text Encoder part / Text Encoder関連部分のみ学習する")
+        "--network_train_text_encoder_only", action="store_true", help="only training Text Encoder part / Text Encoder関連部分のみ学習する"
-  parser.add_argument("--training_comment", type=str, default=None,
+    )
-                      help="arbitrary comment string stored in metadata / メタデータに記録する任意のコメント文字列")
+    parser.add_argument(
+        "--training_comment", type=str, default=None, help="arbitrary comment string stored in metadata / メタデータに記録する任意のコメント文字列"
+    )
 
     args = parser.parse_args()
+    args = train_util.read_config_from_file(args, parser)
+
     train(args)

train_network_README-ja.md

@@ -64,6 +64,10 @@ accelerate launch --num_cpu_threads_per_process 1 train_network.py
   * LoRAのRANKを指定します（``--networkdim=4``など）。省略時は4になります。数が多いほど表現力は増しますが、学習に必要なメモリ、時間は増えます。また闇雲に増やしても良くないようです。
 * `--network_alpha`
   *  アンダーフローを防ぎ安定して学習するための ``alpha`` 値を指定します。デフォルトは1です。``network_dim``と同じ値を指定すると以前のバージョンと同じ動作になります。
+* `--persistent_data_loader_workers`
+  * Windows環境で指定するとエポック間の待ち時間が大幅に短縮されます。
+* `--max_data_loader_n_workers`
+  * データ読み込みのプロセス数を指定します。プロセス数が多いとデータ読み込みが速くなりGPUを効率的に利用できますが、メインメモリを消費します。デフォルトは「`8` または `CPU同時実行スレッド数-1` の小さいほう」なので、メインメモリに余裕がない場合や、GPU使用率が90%程度以上なら、それらの数値を見ながら `2` または `1` 程度まで下げてください。
 * `--network_weights`
   * 学習前に学習済みのLoRAの重みを読み込み、そこから追加で学習します。
 * `--network_train_unet_only`

train_textual_inversion.py

@@ -3,6 +3,7 @@ import argparse
 import gc
 import math
 import os
+import toml
 
 from tqdm import tqdm
 import torch
@@ -104,14 +105,17 @@ def train(args):
         init_token_ids = tokenizer.encode(args.init_word, add_special_tokens=False)
         if len(init_token_ids) > 1 and len(init_token_ids) != args.num_vectors_per_token:
             print(
-          f"token length for init words is not same to num_vectors_per_token, init words is repeated or truncated / 初期化単語のトークン長がnum_vectors_per_tokenと合わないため、繰り返しまたは切り捨てが発生します: length {len(init_token_ids)}")
+                f"token length for init words is not same to num_vectors_per_token, init words is repeated or truncated / 初期化単語のトークン長がnum_vectors_per_tokenと合わないため、繰り返しまたは切り捨てが発生します: length {len(init_token_ids)}"
+            )
     else:
         init_token_ids = None
 
     # add new word to tokenizer, count is num_vectors_per_token
     token_strings = [args.token_string] + [f"{args.token_string}{i+1}" for i in range(args.num_vectors_per_token - 1)]
     num_added_tokens = tokenizer.add_tokens(token_strings)
-  assert num_added_tokens == args.num_vectors_per_token, f"tokenizer has same word to token string. please use another one / 指定したargs.token_stringは既に存在します。別の単語を使ってください: {args.token_string}"
+    assert (
+        num_added_tokens == args.num_vectors_per_token
+    ), f"tokenizer has same word to token string. please use another one / 指定したargs.token_stringは既に存在します。別の単語を使ってください: {args.token_string}"
 
     token_ids = tokenizer.convert_tokens_to_ids(token_strings)
     print(f"tokens are added: {token_ids}")
@@ -132,7 +136,8 @@ def train(args):
     if args.weights is not None:
         embeddings = load_weights(args.weights)
         assert len(token_ids) == len(
-        embeddings), f"num_vectors_per_token is mismatch for weights / 指定した重みとnum_vectors_per_tokenの値が異なります: {len(embeddings)}"
+            embeddings
+        ), f"num_vectors_per_token is mismatch for weights / 指定した重みとnum_vectors_per_tokenの値が異なります: {len(embeddings)}"
         # print(token_ids, embeddings.size())
         for token_id, embedding in zip(token_ids, embeddings):
             token_embeds[token_id] = embedding
@@ -148,25 +153,33 @@ def train(args):
         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)))
+            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.")
             user_config = {
-        "datasets": [{
+                "datasets": [
-          "subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)
+                    {"subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)}
-        }]
+                ]
             }
         else:
             print("Train with captions.")
             user_config = {
-        "datasets": [{
+                "datasets": [
-          "subsets": [{
+                    {
+                        "subsets": [
+                            {
                                 "image_dir": args.train_data_dir,
                                 "metadata_file": args.in_json,
-          }]
+                            }
-        }]
+                        ]
+                    }
+                ]
             }
 
     blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
@@ -202,7 +215,9 @@ def train(args):
         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は使えません"
+        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)
@@ -232,7 +247,13 @@ 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_group, batch_size=1, shuffle=True, 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:
@@ -244,7 +265,8 @@ def train(args):
 
     # acceleratorがなんかよろしくやってくれるらしい
     text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-      text_encoder, optimizer, train_dataloader, lr_scheduler)
+        text_encoder, optimizer, train_dataloader, lr_scheduler
+    )
 
     index_no_updates = torch.arange(len(tokenizer)) < token_ids[0]
     # print(len(index_no_updates), torch.sum(index_no_updates))
@@ -300,8 +322,9 @@ def train(args):
     progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
     global_step = 0
 
-  noise_scheduler = DDPMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
+    noise_scheduler = DDPMScheduler(
-                                  num_train_timesteps=1000, clip_sample=False)
+        beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
+    )
 
     if accelerator.is_main_process:
         accelerator.init_trackers("textual_inversion")
@@ -371,21 +394,26 @@ def train(args):
 
                 # Let's make sure we don't update any embedding weights besides the newly added token
                 with torch.no_grad():
-          unwrap_model(text_encoder).get_input_embeddings().weight[index_no_updates] = orig_embeds_params[index_no_updates]
+                    unwrap_model(text_encoder).get_input_embeddings().weight[index_no_updates] = orig_embeds_params[
+                        index_no_updates
+                    ]
 
             # Checks if the accelerator has performed an optimization step behind the scenes
             if accelerator.sync_gradients:
                 progress_bar.update(1)
                 global_step += 1
 
-        train_util.sample_images(accelerator, args, None, global_step, accelerator.device,
+                train_util.sample_images(
-                                 vae, tokenizer, text_encoder, unet, prompt_replacement)
+                    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])}
                 if args.optimizer_type.lower() == "DAdaptation".lower():  # tracking d*lr value
-          logs["lr/d*lr"] = lr_scheduler.optimizers[0].param_groups[0]['d']*lr_scheduler.optimizers[0].param_groups[0]['lr']
+                    logs["lr/d*lr"] = (
+                        lr_scheduler.optimizers[0].param_groups[0]["d"] * lr_scheduler.optimizers[0].param_groups[0]["lr"]
+                    )
                 accelerator.log(logs, step=global_step)
 
             loss_total += current_loss
@@ -408,13 +436,13 @@ def train(args):
             model_name = train_util.DEFAULT_EPOCH_NAME if args.output_name is None else args.output_name
 
             def save_func():
-        ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, epoch + 1) + '.' + args.save_model_as
+                ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, epoch + 1) + "." + args.save_model_as
                 ckpt_file = os.path.join(args.output_dir, ckpt_name)
                 print(f"saving checkpoint: {ckpt_file}")
                 save_weights(ckpt_file, updated_embs, save_dtype)
 
             def remove_old_func(old_epoch_no):
-        old_ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, old_epoch_no) + '.' + args.save_model_as
+                old_ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, old_epoch_no) + "." + args.save_model_as
                 old_ckpt_file = os.path.join(args.output_dir, old_ckpt_name)
                 if os.path.exists(old_ckpt_file):
                     print(f"removing old checkpoint: {old_ckpt_file}")
@@ -424,8 +452,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,
+        train_util.sample_images(
-                             vae, tokenizer, text_encoder, unet, prompt_replacement)
+            accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet, prompt_replacement
+        )
 
         # end of epoch
 
@@ -446,7 +475,7 @@ def train(args):
         os.makedirs(args.output_dir, exist_ok=True)
 
         model_name = train_util.DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name
-    ckpt_name = model_name + '.' + args.save_model_as
+        ckpt_name = model_name + "." + args.save_model_as
         ckpt_file = os.path.join(args.output_dir, ckpt_name)
 
         print(f"save trained model to {ckpt_file}")
@@ -463,27 +492,29 @@ def save_weights(file, updated_embs, save_dtype):
             v = v.detach().clone().to("cpu").to(save_dtype)
             state_dict[key] = v
 
-  if os.path.splitext(file)[1] == '.safetensors':
+    if os.path.splitext(file)[1] == ".safetensors":
         from safetensors.torch import save_file
+
         save_file(state_dict, file)
     else:
         torch.save(state_dict, file)  # can be loaded in Web UI
 
 
 def load_weights(file):
-  if os.path.splitext(file)[1] == '.safetensors':
+    if os.path.splitext(file)[1] == ".safetensors":
         from safetensors.torch import load_file
+
         data = load_file(file)
     else:
         # compatible to Web UI's file format
-    data = torch.load(file, map_location='cpu')
+        data = torch.load(file, map_location="cpu")
         if type(data) != dict:
             raise ValueError(f"weight file is not dict / 重みファイルがdict形式ではありません: {file}")
 
-    if 'string_to_param' in data:                           # textual inversion embeddings
+        if "string_to_param" in data:  # textual inversion embeddings
-      data = data['string_to_param']
+            data = data["string_to_param"]
-      if hasattr(data, '_parameters'):                      # support old PyTorch?
+            if hasattr(data, "_parameters"):  # support old PyTorch?
-        data = getattr(data, '_parameters')
+                data = getattr(data, "_parameters")
 
     emb = next(iter(data.values()))
     if type(emb) != torch.Tensor:
@@ -495,7 +526,7 @@ def load_weights(file):
     return emb
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
 
     train_util.add_sd_models_arguments(parser)
@@ -504,21 +535,37 @@ if __name__ == '__main__':
     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"],
+    parser.add_argument(
-                      help="format to save the model (default is .pt) / モデル保存時の形式（デフォルトはpt）")
+        "--save_model_as",
+        type=str,
+        default="pt",
+        choices=[None, "ckpt", "pt", "safetensors"],
+        help="format to save the model (default is .pt) / モデル保存時の形式（デフォルトはpt）",
+    )
 
-  parser.add_argument("--weights", type=str, default=None,
+    parser.add_argument("--weights", type=str, default=None, help="embedding weights to initialize / 学習するネットワークの初期重み")
-                      help="embedding weights to initialize / 学習するネットワークの初期重み")
+    parser.add_argument(
-  parser.add_argument("--num_vectors_per_token", type=int, default=1,
+        "--num_vectors_per_token", type=int, default=1, help="number of vectors per token / トークンに割り当てるembeddingsの要素数"
-                      help='number of vectors per token / トークンに割り当てるembeddingsの要素数')
+    )
-  parser.add_argument("--token_string", type=str, default=None,
+    parser.add_argument(
-                      help="token string used in training, must not exist in tokenizer / 学習時に使用されるトークン文字列、tokenizerに存在しない文字であること")
+        "--token_string",
-  parser.add_argument("--init_word", type=str, default=None,
+        type=str,
-                      help="words to initialize vector / ベクトルを初期化に使用する単語、複数可")
+        default=None,
-  parser.add_argument("--use_object_template", action='store_true',
+        help="token string used in training, must not exist in tokenizer / 学習時に使用されるトークン文字列、tokenizerに存在しない文字であること",
-                      help="ignore caption and use default templates for object / キャプションは使わずデフォルトの物体用テンプレートで学習する")
+    )
-  parser.add_argument("--use_style_template", action='store_true',
+    parser.add_argument("--init_word", type=str, default=None, help="words to initialize vector / ベクトルを初期化に使用する単語、複数可")
-                      help="ignore caption and use default templates for stype / キャプションは使わずデフォルトのスタイル用テンプレートで学習する")
+    parser.add_argument(
+        "--use_object_template",
+        action="store_true",
+        help="ignore caption and use default templates for object / キャプションは使わずデフォルトの物体用テンプレートで学習する",
+    )
+    parser.add_argument(
+        "--use_style_template",
+        action="store_true",
+        help="ignore caption and use default templates for stype / キャプションは使わずデフォルトのスタイル用テンプレートで学習する",
+    )
 
     args = parser.parse_args()
+    args = train_util.read_config_from_file(args, parser)
+
     train(args)