format with black

fine_tune.py

@@ -20,6 +20,7 @@ from library.config_util import (
     BlueprintGenerator,
 )
 
+
 def collate_fn(examples):
     return examples[0]
 
@@ -41,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)
@@ -59,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")
@@ -87,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使用フラグを設定する関数
@@ -171,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:
@@ -179,13 +198,20 @@ def train(args):
         print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
 
     # lr schedulerを用意する
-  lr_scheduler = train_util.get_scheduler_fix(args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps,
+    lr_scheduler = train_util.get_scheduler_fix(
+        args.lr_scheduler,
+        optimizer,
+        num_warmup_steps=args.lr_warmup_steps,
         num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
-                                              num_cycles=args.lr_scheduler_num_cycles, power=args.lr_scheduler_power)
+        num_cycles=args.lr_scheduler_num_cycles,
+        power=args.lr_scheduler_power,
+    )
 
     # 実験的機能：勾配も含めた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)
@@ -193,7 +219,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)
 
@@ -226,8 +253,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")
@@ -255,7 +283,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)
@@ -298,18 +327,22 @@ 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にする
             loss_total += current_loss
-      avr_loss = loss_total / (step+1)
+            avr_loss = loss_total / (step + 1)
             logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
             progress_bar.set_postfix(**logs)
 
@@ -318,14 +351,26 @@ def train(args):
 
         if args.logging_dir is not None:
             logs = {"loss/epoch": loss_total / len(train_dataloader)}
-      accelerator.log(logs, step=epoch+1)
+            accelerator.log(logs, step=epoch + 1)
 
         accelerator.wait_for_everyone()
 
         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)
 
@@ -343,12 +388,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)
@@ -358,8 +404,7 @@ 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()

train_db.py

@@ -44,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)
@@ -63,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)
 
@@ -93,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 を組み込む
@@ -130,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()
 
@@ -140,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:
@@ -151,13 +166,20 @@ def train(args):
         args.stop_text_encoder_training = args.max_train_steps + 1  # do not stop until end
 
     # lr schedulerを用意する TODO gradient_accumulation_stepsの扱いが何かおかしいかもしれない。後で確認する
-  lr_scheduler = train_util.get_scheduler_fix(args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps,
+    lr_scheduler = train_util.get_scheduler_fix(
+        args.lr_scheduler,
+        optimizer,
+        num_warmup_steps=args.lr_warmup_steps,
         num_training_steps=args.max_train_steps,
-                                              num_cycles=args.lr_scheduler_num_cycles, power=args.lr_scheduler_power)
+        num_cycles=args.lr_scheduler_num_cycles,
+        power=args.lr_scheduler_power,
+    )
 
     # 実験的機能：勾配も含めた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)
@@ -165,7 +187,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)
 
@@ -202,8 +225,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")
@@ -248,7 +272,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)
@@ -278,7 +303,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)
@@ -292,13 +317,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:
@@ -316,14 +345,26 @@ def train(args):
 
         if args.logging_dir is not None:
             logs = {"loss/epoch": loss_total / len(loss_list)}
-      accelerator.log(logs, step=epoch+1)
+            accelerator.log(logs, step=epoch + 1)
 
         accelerator.wait_for_everyone()
 
         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)
 
@@ -341,12 +382,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)
@@ -356,10 +398,17 @@ 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()
 

train_network.py

@@ -42,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
 
@@ -70,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)
@@ -97,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を準備する
@@ -137,6 +147,7 @@ def train(args):
 
     # 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)
@@ -144,7 +155,7 @@ def train(args):
     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 (;'∀')
@@ -175,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:
@@ -184,29 +201,38 @@ def train(args):
             print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
 
     # lr schedulerを用意する
-  lr_scheduler = train_util.get_scheduler_fix(args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps,
+    lr_scheduler = train_util.get_scheduler_fix(
+        args.lr_scheduler,
+        optimizer,
+        num_warmup_steps=args.lr_warmup_steps,
         num_training_steps=args.max_train_steps * accelerator.num_processes * args.gradient_accumulation_steps,
-                                              num_cycles=args.lr_scheduler_num_cycles, power=args.lr_scheduler_power)
+        num_cycles=args.lr_scheduler_num_cycles,
+        power=args.lr_scheduler_power,
+    )
 
     # 実験的機能：勾配も含めた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)
@@ -372,10 +398,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)
@@ -394,8 +417,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]
 
@@ -404,18 +428,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,
@@ -432,7 +454,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:
@@ -469,8 +492,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")
@@ -482,7 +506,7 @@ def train(args):
             print(f"epoch {epoch+1}/{num_train_epochs}")
         train_dataset_group.set_current_epoch(epoch + 1)
 
-    metadata["ss_epoch"] = str(epoch+1)
+        metadata["ss_epoch"] = str(epoch + 1)
 
         network.on_epoch_start(text_encoder, unet)
 
@@ -548,7 +572,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:
@@ -570,7 +596,7 @@ def train(args):
 
         if args.logging_dir is not None:
             logs = {"loss/epoch": loss_total / len(loss_list)}
-      accelerator.log(logs, step=epoch+1)
+            accelerator.log(logs, step=epoch + 1)
 
         accelerator.wait_for_everyone()
 
@@ -578,14 +604,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}")
@@ -617,7 +643,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}")
@@ -625,7 +651,7 @@ def train(args):
         print("model saved.")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
 
     train_util.add_sd_models_arguments(parser)
@@ -634,27 +660,39 @@ 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()
 

train_textual_inversion.py

@@ -105,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}")
@@ -133,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
@@ -149,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)
@@ -203,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)
@@ -233,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:
@@ -241,13 +261,19 @@ def train(args):
         print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
 
     # lr schedulerを用意する
-  lr_scheduler = train_util.get_scheduler_fix(args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps,
+    lr_scheduler = train_util.get_scheduler_fix(
+        args.lr_scheduler,
+        optimizer,
+        num_warmup_steps=args.lr_warmup_steps,
         num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
-                                              num_cycles=args.lr_scheduler_num_cycles, power=args.lr_scheduler_power)
+        num_cycles=args.lr_scheduler_num_cycles,
+        power=args.lr_scheduler_power,
+    )
 
     # 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))
@@ -303,8 +329,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")
@@ -374,25 +401,30 @@ 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
-      avr_loss = loss_total / (step+1)
+            avr_loss = loss_total / (step + 1)
             logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
             progress_bar.set_postfix(**logs)
 
@@ -401,7 +433,7 @@ def train(args):
 
         if args.logging_dir is not None:
             logs = {"loss/epoch": loss_total / len(train_dataloader)}
-      accelerator.log(logs, step=epoch+1)
+            accelerator.log(logs, step=epoch + 1)
 
         accelerator.wait_for_everyone()
 
@@ -411,13 +443,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}")
@@ -427,8 +459,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
 
@@ -449,7 +482,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}")
@@ -466,27 +499,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:
@@ -498,7 +533,7 @@ def load_weights(file):
     return emb
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
 
     train_util.add_sd_models_arguments(parser)
@@ -507,21 +542,35 @@ 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()