Upload 55 files

.gitattributes

@@ -1 +1,34 @@
-bitsandbytes_windows/libbitsandbytes_cuda116.dll filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

append_module.py

@@ -2,19 +2,7 @@ import argparse
 import json
 import shutil
 import time
-from typing import (
-  Dict,
-  List,
-  NamedTuple,
-  Optional,
-  Sequence,
-  Tuple,
-  Union,
-)
-from dataclasses import (
-  asdict,
-  dataclass,
-)
+from typing import Dict, List, NamedTuple, Tuple
 from accelerate import Accelerator
 from torch.autograd.function import Function
 import glob
@@ -40,7 +28,6 @@ import safetensors.torch
 
 import library.model_util as model_util
 import library.train_util as train_util
-import library.config_util as config_util
 
 #============================================================================================================
 #AdafactorScheduleに暫定的にinitial_lrを層別に適用できるようにしたもの
@@ -128,124 +115,6 @@ def make_bucket_resolutions_fix(max_reso, min_reso, min_size=256, max_size=1024,
   return area_size_resos_list, area_size_list
 
 #============================================================================================================
-#config_util 内より
-#============================================================================================================
-@dataclass
-class DreamBoothDatasetParams(config_util.DreamBoothDatasetParams):
-  min_resolution: Optional[Tuple[int, int]] = None
-  area_step : int = 2
-
-class ConfigSanitizer(config_util.ConfigSanitizer):
-  #@config_util.curry
-  @staticmethod
-  def __validate_and_convert_twodim(klass, value: Sequence) -> Tuple:
-    config_util.Schema(config_util.ExactSequence([klass, klass]))(value)
-    return tuple(value)
-
-  #@config_util.curry
-  @staticmethod
-  def __validate_and_convert_scalar_or_twodim(klass, value: Union[float, Sequence]) -> Tuple:
-    config_util.Schema(config_util.Any(klass, config_util.ExactSequence([klass, klass])))(value)
-    try:
-      config_util.Schema(klass)(value)
-      return (value, value)
-    except:
-      return ConfigSanitizer.__validate_and_convert_twodim(klass, value)
-  # datasets schema
-  DATASET_ASCENDABLE_SCHEMA = {
-    "batch_size": int,
-    "bucket_no_upscale": bool,
-    "bucket_reso_steps": int,
-    "enable_bucket": bool,
-    "max_bucket_reso": int,
-    "min_bucket_reso": int,
-    "resolution": config_util.functools.partial(__validate_and_convert_scalar_or_twodim.__func__, int),
-    "min_resolution": config_util.functools.partial(__validate_and_convert_scalar_or_twodim.__func__, int),
-    "area_step": int,
-  }
-  def __init__(self, support_dreambooth: bool, support_finetuning: bool, support_dropout: bool) -> None:
-    super().__init__(support_dreambooth, support_finetuning, support_dropout)
-  def _check(self):
-    print(self.db_dataset_schema)
-
-class BlueprintGenerator(config_util.BlueprintGenerator):
-  def __init__(self, sanitizer: ConfigSanitizer):
-    config_util.DreamBoothDatasetParams = DreamBoothDatasetParams
-    super().__init__(sanitizer)
-
-def generate_dataset_group_by_blueprint(dataset_group_blueprint: config_util.DatasetGroupBlueprint):
-  datasets: List[Union[DreamBoothDataset, train_util.FineTuningDataset]] = []
-
-  for dataset_blueprint in dataset_group_blueprint.datasets:
-    if dataset_blueprint.is_dreambooth:
-      subset_klass = train_util.DreamBoothSubset
-      dataset_klass = DreamBoothDataset
-    else:
-      subset_klass = train_util.FineTuningSubset
-      dataset_klass = train_util.FineTuningDataset
-
-    subsets = [subset_klass(**asdict(subset_blueprint.params)) for subset_blueprint in dataset_blueprint.subsets]
-    dataset = dataset_klass(subsets=subsets, **asdict(dataset_blueprint.params))
-    datasets.append(dataset)
-
-  # print info
-  info = ""
-  for i, dataset in enumerate(datasets):
-    is_dreambooth = isinstance(dataset, DreamBoothDataset)
-    info += config_util.dedent(f"""\
-      [Dataset {i}]
-        batch_size: {dataset.batch_size}
-        resolution: {(dataset.width, dataset.height)}
-        enable_bucket: {dataset.enable_bucket}
-    """)
-
-    if dataset.enable_bucket:
-      info += config_util.indent(config_util.dedent(f"""\
-        min_bucket_reso: {dataset.min_bucket_reso}
-        max_bucket_reso: {dataset.max_bucket_reso}
-        bucket_reso_steps: {dataset.bucket_reso_steps}
-        bucket_no_upscale: {dataset.bucket_no_upscale}
-      \n"""), "  ")
-    else:
-      info += "\n"
-
-    for j, subset in enumerate(dataset.subsets):
-      info += config_util.indent(config_util.dedent(f"""\
-        [Subset {j} of Dataset {i}]
-          image_dir: "{subset.image_dir}"
-          image_count: {subset.img_count}
-          num_repeats: {subset.num_repeats}
-          shuffle_caption: {subset.shuffle_caption}
-          keep_tokens: {subset.keep_tokens}
-          caption_dropout_rate: {subset.caption_dropout_rate}
-          caption_dropout_every_n_epoches: {subset.caption_dropout_every_n_epochs}
-          caption_tag_dropout_rate: {subset.caption_tag_dropout_rate}
-          color_aug: {subset.color_aug}
-          flip_aug: {subset.flip_aug}
-          face_crop_aug_range: {subset.face_crop_aug_range}
-          random_crop: {subset.random_crop}
-      """), "  ")
-
-      if is_dreambooth:
-        info += config_util.indent(config_util.dedent(f"""\
-          is_reg: {subset.is_reg}
-          class_tokens: {subset.class_tokens}
-          caption_extension: {subset.caption_extension}
-        \n"""), "    ")
-      else:
-        info += config_util.indent(config_util.dedent(f"""\
-          metadata_file: {subset.metadata_file}
-        \n"""), "    ")
-
-  print(info)
-
-  # make buckets first because it determines the length of dataset
-  for i, dataset in enumerate(datasets):
-    print(f"[Dataset {i}]")
-    dataset.make_buckets()
-
-  return train_util.DatasetGroup(datasets)
-#============================================================================================================
 #train_util 内より
 #============================================================================================================
 class BucketManager_append(train_util.BucketManager):
@@ -310,7 +179,7 @@ class BucketManager_append(train_util.BucketManager):
             bucket_size_id_list.append(bucket_size_id + i + 1)
         _min_error = 1000.
         _min_id = bucket_size_id
-        for now_size_id in bucket_size_id_list:
+        for now_size_id in bucket_size_id:
           self.predefined_aspect_ratios = self.predefined_aspect_ratios_list[now_size_id]
           ar_errors = self.predefined_aspect_ratios - aspect_ratio
           ar_error = np.abs(ar_errors).min()
@@ -384,13 +253,13 @@ class BucketManager_append(train_util.BucketManager):
     return reso, resized_size, ar_error
 
 class DreamBoothDataset(train_util.DreamBoothDataset):
-  def __init__(self, subsets: Sequence[train_util.DreamBoothSubset], batch_size: int, tokenizer, max_token_length, resolution, enable_bucket: bool, min_bucket_reso: int, max_bucket_reso: int, bucket_reso_steps: int, bucket_no_upscale: bool, prior_loss_weight: float, debug_dataset, min_resolution=None, area_step=None) -> None:
+  def __init__(self, batch_size, train_data_dir, reg_data_dir, tokenizer, max_token_length, caption_extension, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, prior_loss_weight, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset, min_resolution=None, area_step=None) -> None:
     print("use append DreamBoothDataset")
     self.min_resolution = min_resolution
     self.area_step = area_step
-    super().__init__(subsets, batch_size, tokenizer, max_token_length,
-                    resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale,
-                    prior_loss_weight, debug_dataset)
+    super().__init__(batch_size, train_data_dir, reg_data_dir, tokenizer, max_token_length, caption_extension, shuffle_caption, shuffle_keep_tokens,
+                      resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, prior_loss_weight,
+                      flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset)
   def make_buckets(self):
     '''
     bucketingを行わない場合も呼び出し必須（ひとつだけbucketを作る）
@@ -483,50 +352,40 @@ class DreamBoothDataset(train_util.DreamBoothDataset):
     self.shuffle_buckets()
     self._length = len(self.buckets_indices)
 
-import transformers
-from torch.optim import Optimizer
-from diffusers.optimization import SchedulerType
-from typing import Union
-def get_scheduler_Adafactor(
-    name: Union[str, SchedulerType],
-    optimizer: Optimizer,
-    scheduler_arg: Dict
-):
-  if name.startswith("adafactor"):
-    assert type(optimizer) == transformers.optimization.Adafactor, f"adafactor scheduler must be used with Adafactor optimizer / adafactor schedulerはAdafactorオプティマイザと同時に使ってください"
-    print(scheduler_arg)
-    return AdafactorSchedule_append(optimizer, **scheduler_arg)
+class FineTuningDataset(train_util.FineTuningDataset):
+  def __init__(self, json_file_name, batch_size, train_data_dir, tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, flip_aug, color_aug, face_crop_aug_range, random_crop, dataset_repeats, debug_dataset) -> None:
+    train_util.glob_images = glob_images
+    super().__init__( json_file_name, batch_size, train_data_dir, tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens,
+                      resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, flip_aug, color_aug, face_crop_aug_range,
+                      random_crop, dataset_repeats, debug_dataset)
+
+def glob_images(directory, base="*", npz_flag=True):
+  img_paths = []
+  dots = []
+  for ext in train_util.IMAGE_EXTENSIONS:
+    dots.append(ext)
+  if npz_flag:
+    dots.append(".npz")
+  for ext in dots:
+    if base == '*':
+      img_paths.extend(glob.glob(os.path.join(glob.escape(directory), base + ext)))
+    else:
+      img_paths.extend(glob.glob(glob.escape(os.path.join(directory, base + ext))))
+  return img_paths
+
 #============================================================================================================
 #networks.lora
 #============================================================================================================
-#from networks.lora import LoRANetwork
-def replace_prepare_optimizer_params(networks, network_module):
-  def prepare_optimizer_params(self, text_encoder_lr, unet_lr, loranames=None, lr_dic=None, block_args_dic=None):
-
+from networks.lora import LoRANetwork
+def replace_prepare_optimizer_params(networks):
+  def prepare_optimizer_params(self, text_encoder_lr, unet_lr, scheduler_lr=None, loranames=None):
+    
     def enumerate_params(loras, lora_name=None):
       params = []
       for lora in loras:
         if lora_name is not None:
-          get_param_flag = False
-          if "attentions" in lora_name or "lora_unet_up_blocks_0_resnets_2":
-            lora_names = [lora_name]
-            if "attentions" in lora_name:
-              lora_names.append(lora_name.replace("attentions", "resnets"))
-            elif "lora_unet_up_blocks_0_resnets_2" in lora_name:
-              lora_names.append("lora_unet_up_blocks_0_upsamplers_")
-            elif "lora_unet_up_blocks_1_attentions_2_" in lora_name:
-              lora_names.append("lora_unet_up_blocks_1_upsamplers_")
-            elif "lora_unet_up_blocks_2_attentions_2_" in lora_name:
-              lora_names.append("lora_unet_up_blocks_2_upsamplers_")
-            
-            for _name in lora_names:
-              if _name in lora.lora_name:
-                get_param_flag = True
-                break
-          else:
-            if lora_name in lora.lora_name:
-              get_param_flag = True
-          if get_param_flag: params.extend(lora.parameters())
+          if lora_name in lora.lora_name:
+            params.extend(lora.parameters())
         else:
           params.extend(lora.parameters())
       return params
@@ -534,7 +393,6 @@ def replace_prepare_optimizer_params(networks, network_module):
     self.requires_grad_(True)
     all_params = []
     ret_scheduler_lr = []
-    used_names = []
     
     if loranames is not None:
       textencoder_names = [None]
@@ -547,181 +405,37 @@ def replace_prepare_optimizer_params(networks, network_module):
     if self.text_encoder_loras:
       for textencoder_name in textencoder_names:
         param_data = {'params': enumerate_params(self.text_encoder_loras, lora_name=textencoder_name)}
-        used_names.append(textencoder_name)
         if text_encoder_lr is not None:
           param_data['lr'] = text_encoder_lr
-          if lr_dic is not None:
-            if textencoder_name in lr_dic:
-              param_data['lr'] = lr_dic[textencoder_name]
-              print(f"{textencoder_name} lr: {param_data['lr']}")
-
-        if block_args_dic is not None:
-          if "lora_te_" in block_args_dic:
-            for pname, value in block_args_dic["lora_te_"].items():
-              param_data[pname] = value
-          if textencoder_name in block_args_dic:
-            for pname, value in block_args_dic[textencoder_name].items():
-              param_data[pname] = value
-
-        if text_encoder_lr is not None:
-          ret_scheduler_lr.append(text_encoder_lr)
-        else:
-          ret_scheduler_lr.append(0.)
-        if lr_dic is not None:
-          if textencoder_name in lr_dic:
-            ret_scheduler_lr[-1] = lr_dic[textencoder_name]
+        if scheduler_lr is not None:
+          ret_scheduler_lr.append(scheduler_lr[0])
         all_params.append(param_data)
 
     if self.unet_loras:
       for unet_name in unet_names:
         param_data = {'params': enumerate_params(self.unet_loras, lora_name=unet_name)}
-        if len(param_data["params"])==0: continue
-        used_names.append(unet_name)
         if unet_lr is not None:
           param_data['lr'] = unet_lr
-          if lr_dic is not None:
-            if unet_name in lr_dic:
-              param_data['lr'] = lr_dic[unet_name]
-              print(f"{unet_name} lr: {param_data['lr']}")
-
-        if block_args_dic is not None:
-          if "lora_unet_" in block_args_dic:
-            for pname, value in block_args_dic["lora_unet_"].items():
-              param_data[pname] = value
-          if unet_name in block_args_dic:
-            for pname, value in block_args_dic[unet_name].items():
-              param_data[pname] = value
-        
-        if unet_lr is not None:
-          ret_scheduler_lr.append(unet_lr)
-        else:
-          ret_scheduler_lr.append(0.)
-        if lr_dic is not None:
-          if unet_name in lr_dic:
-            ret_scheduler_lr[-1] = lr_dic[unet_name]
+        if scheduler_lr is not None:
+          ret_scheduler_lr.append(scheduler_lr[1])
         all_params.append(param_data)
 
-    return all_params, {"initial_lr" : ret_scheduler_lr}, used_names
-  try:
-    network_module.LoRANetwork.prepare_optimizer_params = prepare_optimizer_params
-  except:
-    print("cant't replace prepare_optimizer_params")
+    return all_params, ret_scheduler_lr
+  
+  LoRANetwork.prepare_optimizer_params = prepare_optimizer_params
 
 #============================================================================================================
 #新規追加
 #============================================================================================================
 def add_append_arguments(parser: argparse.ArgumentParser):
   # for train_network_opt.py
-  #parser.add_argument("--optimizer", type=str, default="AdamW", choices=["AdamW", "RAdam", "AdaBound", "AdaBelief", "AggMo", "AdamP", "Adastand", "Adastand_belief", "Apollo", "Lamb", "Ranger", "RangerVA", "Lookahead_Adam", "Lookahead_DiffGrad", "Yogi", "NovoGrad", "QHAdam", "DiffGrad", "MADGRAD", "Adafactor"], help="使用するoptimizerを指定する")
-  #parser.add_argument("--optimizer_arg", type=str, default=None, nargs='*')
-  parser.add_argument("--use_lookahead", action="store_true")
-  parser.add_argument("--lookahead_arg", type=str, nargs="*", default=None)
+  parser.add_argument("--optimizer", type=str, default="AdamW", choices=["AdamW", "RAdam", "AdaBound", "AdaBelief", "AggMo", "AdamP", "Adastand", "Adastand_belief", "Apollo", "Lamb", "Ranger", "RangerVA", "Lookahead_Adam", "Lookahead_DiffGrad", "Yogi", "NovoGrad", "QHAdam", "DiffGrad", "MADGRAD", "Adafactor"], help="使用するoptimizerを指定する")
+  parser.add_argument("--optimizer_arg", type=str, default=None, nargs='*')
   parser.add_argument("--split_lora_networks", action="store_true")
   parser.add_argument("--split_lora_level", type=int, default=0, help="どれくらい細分化するかの設定 0がunetのみを層別に 1がunetを大枠で分割 2がtextencoder含めて層別")
-  parser.add_argument("--blocks_lr_setting", type=str, default=None)
-  parser.add_argument("--block_optim_args", type=str, nargs="*", default=None)
   parser.add_argument("--min_resolution", type=str, default=None)
   parser.add_argument("--area_step", type=int, default=1)
   parser.add_argument("--config", type=str, default=None)
-  parser.add_argument("--not_output_config", action="store_true")
-
-class MyNetwork_Names:
-  ex_block_weight_dic = {
-    "BASE": ["te"],
-    "IN01": ["down_0_at_0","donw_0_res_0"], "IN02": ["down_0_at_1","down_0_res_1"], "IN03": ["down_0_down"],
-    "IN04": ["down_1_at_0","donw_1_res_0"], "IN05": ["down_1_at_1","donw_1_res_1"], "IN06": ["down_1_down"],
-    "IN07": ["down_2_at_0","donw_2_res_0"], "IN08": ["down_2_at_1","donw_2_res_1"], "IN09": ["down_2_down"],
-    "IN10": ["down_3_res_0"], "IN11": ["down_3_res_1"],
-    "MID": ["mid"],
-    "OUT00": ["up_0_res_0"], "OUT01": ["up_0_res_1"], "OUT02": ["up_0_res_2", "up_0_up"],
-    "OUT03": ["up_1_at_0", "up_1_res_0"], "OUT04": ["up_1_at_1", "up_1_res_1"], "OUT05": ["up_1_at_2", "up_1_res_2", "up_1_up"],
-    "OUT06": ["up_2_at_0", "up_2_res_0"], "OUT07": ["up_2_at_1", "up_2_res_1"], "OUT08": ["up_2_at_2", "up_2_res_2", "up_2_up"],
-    "OUT09": ["up_3_at_0", "up_3_res_0"], "OUT10": ["up_3_at_1", "up_3_res_1"], "OUT11": ["up_3_at_2", "up_3_res_2"],
-  }
-
-  blocks_name_dic = { "te": "lora_te_",
-                      "unet": "lora_unet_",
-                      "mid": "lora_unet_mid_block_",
-                      "down": "lora_unet_down_blocks_",
-                      "up": "lora_unet_up_blocks_"}
-  for i in range(12):
-    blocks_name_dic[f"te_{i}"] = f"lora_te_text_model_encoder_layers_{i}_"
-  for i in range(3):
-    blocks_name_dic[f"down_{i}"] = f"lora_unet_down_blocks_{i}"
-    blocks_name_dic[f"up_{i+1}"] = f"lora_unet_up_blocks_{i+1}"
-  for i in range(4):
-    for j in range(2):
-      if i<=2: blocks_name_dic[f"down_{i}_at_{j}"] = f"lora_unet_down_blocks_{i}_attentions_{j}_"
-      blocks_name_dic[f"down_{i}_res_{j}"] = f"lora_unet_down_blocks_{i}_resnets_{j}"
-    for j in range(3):
-      if i>=1: blocks_name_dic[f"up_{i}_at_{j}"] = f"lora_unet_up_blocks_{i}_attentions_{j}_"
-      blocks_name_dic[f"up_{i}_res_{j}"] = f"lora_unet_up_blocks_{i}_resnets_{j}"
-    if i<=2:
-      blocks_name_dic[f"down_{i}_down"] = f"lora_unet_down_blocks_{i}_downsamplers_"
-      blocks_name_dic[f"up_{i}_up"] = f"lora_unet_up_blocks_{i}_upsamplers_"
-
-def create_lr_blocks(lr_setting_str=None, block_optim_args=None):
-  ex_block_weight_dic = MyNetwork_Names.ex_block_weight_dic
-  blocks_name_dic = MyNetwork_Names.blocks_name_dic
-  
-  lr_dic = {}
-  if lr_setting_str==None or lr_setting_str=="":
-    pass
-  else:
-    lr_settings = lr_setting_str.replace(" ", "").split(",")
-    for lr_setting in lr_settings:
-      key, value = lr_setting.split("=")
-      if key in ex_block_weight_dic:
-        keys = ex_block_weight_dic[key]
-      else:
-        keys = [key]
-      for key in keys:
-        if key in blocks_name_dic:
-          new_key = blocks_name_dic[key]
-          lr_dic[new_key] = float(value)
-  if len(lr_dic)==0:
-    lr_dic = None
-  
-  args_dic = {}
-  if (block_optim_args is None):
-    block_optim_args = []
-  if (len(block_optim_args)>0):
-    for my_arg in block_optim_args:
-      my_arg = my_arg.replace(" ", "")
-      splits = my_arg.split(":")
-      b_name = splits[0]
-      
-      key, _value = splits[1].split("=")
-      value_type = float
-      if len(splits)==3:
-        if _value=="str":
-          value_type = str
-        elif _value=="int":
-          value_type = int
-        _value = splits[2]
-      if _value=="true" or _value=="false":
-        value_type = bool
-      if "," in _value:
-        _value = _value.split(",")
-        for i in range(len(_value)):
-          _value[i] = value_type(_value[i])
-        value=tuple(_value)
-      else:
-        value = value_type(_value)
-
-      if b_name in ex_block_weight_dic:
-        b_names = ex_block_weight_dic[b_name]
-      else:
-        b_names = [b_name]
-      for b_name in b_names:
-        new_b_name = blocks_name_dic[b_name]
-        if not new_b_name in args_dic:
-          args_dic[new_b_name] = {}
-        args_dic[new_b_name][key] = value
-  
-  if len(args_dic)==0:
-    args_dic = None
-  return lr_dic, args_dic
 
 def create_split_names(split_flag, split_level):
   split_names = None
@@ -732,28 +446,14 @@ def create_split_names(split_flag, split_level):
     if split_level==1:
       unet_names.append(f"lora_unet_down_blocks_")
       unet_names.append(f"lora_unet_up_blocks_")
-    elif split_level==2 or split_level==0 or split_level==4:
-      if split_level>=2:
+    elif split_level==2 or split_level==0:
+      if split_level==2:
         text_encoder_names = []
         for i in range(12):
           text_encoder_names.append(f"lora_te_text_model_encoder_layers_{i}_")
-      
-      if split_level<=2:
-        for i in range(3):
-          unet_names.append(f"lora_unet_down_blocks_{i}")
-          unet_names.append(f"lora_unet_up_blocks_{i+1}")
-      
-    if split_level>=3:
-      for i in range(4):
-        for j in range(2):
-          if i<=2: unet_names.append(f"lora_unet_down_blocks_{i}_attentions_{j}_")
-          if i== 3: unet_names.append(f"lora_unet_down_blocks_{i}_resnets_{j}")
-        for j in range(3):
-          if i>=1: unet_names.append(f"lora_unet_up_blocks_{i}_attentions_{j}_")
-          if i==0: unet_names.append(f"lora_unet_up_blocks_{i}_resnets_{j}")
-        if i<=2:
-          unet_names.append(f"lora_unet_down_blocks_{i}_downsamplers_")
-
+      for i in range(3):
+        unet_names.append(f"lora_unet_down_blocks_{i}")
+        unet_names.append(f"lora_unet_up_blocks_{i+1}")
     split_names["text_encoder"] = text_encoder_names
     split_names["unet"] = unet_names
   return split_names
@@ -765,7 +465,7 @@ def get_config(parser):
     import datetime
     if os.path.splitext(args.config)[-1] == ".yaml":
       args.config = os.path.splitext(args.config)[0]
-    config_path = f"{args.config}.yaml"
+    config_path = f"./{args.config}.yaml"
     if os.path.exists(config_path):
       print(f"{config_path} から設定を読み込み中...")
       margs, rest = parser.parse_known_args()
@@ -786,41 +486,19 @@ def get_config(parser):
         args_type_dic[key] = act.type
       #データタイプの確認とargsにkeyの内容を代入していく
       for key, v in configs.items():
-        if v is not None:
-          if key in args_dic:
-            if args_dic[key] is not None:
-              new_type = type(args_dic[key])
-              if (not type(v) == new_type) and (not new_type==list):
-                  v = new_type(v)
-            else:
+        if key in args_dic:
+          if args_dic[key] is not None:
+            new_type = type(args_dic[key])
+            if (not type(v) == new_type) and (not new_type==list):
+              v = new_type(v)
+          else:
+            if v is not None:
               if not type(v) == args_type_dic[key]:
                 v = args_type_dic[key](v)
-        args_dic[key] = v
+          args_dic[key] = v
       #最後にデフォから指定が変わってるものを変更する
       for key, v in change_def_dic.items():
         args_dic[key] = v
     else:
       print(f"{config_path} が見つかりませんでした")
   return args
-
-'''
-class GradientReversalFunction(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, input_forward: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
-        ctx.save_for_backward(scale)
-        return input_forward
-    @staticmethod
-    def backward(ctx, grad_backward: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        scale, = ctx.saved_tensors
-        return scale * -grad_backward, None
-
-class GradientReversal(torch.nn.Module):
-    def __init__(self, scale: float):
-        super(GradientReversal, self).__init__()
-        self.scale = torch.tensor(scale) 
-    def forward(self, x: torch.Tensor, flag: bool = False) -> torch.Tensor:
-      if flag:
-        return x
-      else:
-        return GradientReversalFunction.apply(x, self.scale)
-'''

build/lib/library/model_util.py

@@ -0,0 +1,1180 @@
+# v1: split from train_db_fixed.py.
+# v2: support safetensors
+
+import math
+import os
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig
+from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNet2DConditionModel
+from safetensors.torch import load_file, save_file
+
+# DiffUsers版StableDiffusionのモデルパラメータ
+NUM_TRAIN_TIMESTEPS = 1000
+BETA_START = 0.00085
+BETA_END = 0.0120
+
+UNET_PARAMS_MODEL_CHANNELS = 320
+UNET_PARAMS_CHANNEL_MULT = [1, 2, 4, 4]
+UNET_PARAMS_ATTENTION_RESOLUTIONS = [4, 2, 1]
+UNET_PARAMS_IMAGE_SIZE = 64  # fixed from old invalid value `32`
+UNET_PARAMS_IN_CHANNELS = 4
+UNET_PARAMS_OUT_CHANNELS = 4
+UNET_PARAMS_NUM_RES_BLOCKS = 2
+UNET_PARAMS_CONTEXT_DIM = 768
+UNET_PARAMS_NUM_HEADS = 8
+
+VAE_PARAMS_Z_CHANNELS = 4
+VAE_PARAMS_RESOLUTION = 256
+VAE_PARAMS_IN_CHANNELS = 3
+VAE_PARAMS_OUT_CH = 3
+VAE_PARAMS_CH = 128
+VAE_PARAMS_CH_MULT = [1, 2, 4, 4]
+VAE_PARAMS_NUM_RES_BLOCKS = 2
+
+# V2
+V2_UNET_PARAMS_ATTENTION_HEAD_DIM = [5, 10, 20, 20]
+V2_UNET_PARAMS_CONTEXT_DIM = 1024
+
+# Diffusersの設定を読み込むための参照モデル
+DIFFUSERS_REF_MODEL_ID_V1 = "runwayml/stable-diffusion-v1-5"
+DIFFUSERS_REF_MODEL_ID_V2 = "stabilityai/stable-diffusion-2-1"
+
+
+# region StableDiffusion->Diffusersの変換コード
+# convert_original_stable_diffusion_to_diffusers をコピーして修正している（ASL 2.0）
+
+
+def shave_segments(path, n_shave_prefix_segments=1):
+  """
+  Removes segments. Positive values shave the first segments, negative shave the last segments.
+  """
+  if n_shave_prefix_segments >= 0:
+    return ".".join(path.split(".")[n_shave_prefix_segments:])
+  else:
+    return ".".join(path.split(".")[:n_shave_prefix_segments])
+
+
+def renew_resnet_paths(old_list, n_shave_prefix_segments=0):
+  """
+  Updates paths inside resnets to the new naming scheme (local renaming)
+  """
+  mapping = []
+  for old_item in old_list:
+    new_item = old_item.replace("in_layers.0", "norm1")
+    new_item = new_item.replace("in_layers.2", "conv1")
+
+    new_item = new_item.replace("out_layers.0", "norm2")
+    new_item = new_item.replace("out_layers.3", "conv2")
+
+    new_item = new_item.replace("emb_layers.1", "time_emb_proj")
+    new_item = new_item.replace("skip_connection", "conv_shortcut")
+
+    new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+    mapping.append({"old": old_item, "new": new_item})
+
+  return mapping
+
+
+def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0):
+  """
+  Updates paths inside resnets to the new naming scheme (local renaming)
+  """
+  mapping = []
+  for old_item in old_list:
+    new_item = old_item
+
+    new_item = new_item.replace("nin_shortcut", "conv_shortcut")
+    new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+    mapping.append({"old": old_item, "new": new_item})
+
+  return mapping
+
+
+def renew_attention_paths(old_list, n_shave_prefix_segments=0):
+  """
+  Updates paths inside attentions to the new naming scheme (local renaming)
+  """
+  mapping = []
+  for old_item in old_list:
+    new_item = old_item
+
+    #         new_item = new_item.replace('norm.weight', 'group_norm.weight')
+    #         new_item = new_item.replace('norm.bias', 'group_norm.bias')
+
+    #         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
+    #         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')
+
+    #         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+    mapping.append({"old": old_item, "new": new_item})
+
+  return mapping
+
+
+def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):
+  """
+  Updates paths inside attentions to the new naming scheme (local renaming)
+  """
+  mapping = []
+  for old_item in old_list:
+    new_item = old_item
+
+    new_item = new_item.replace("norm.weight", "group_norm.weight")
+    new_item = new_item.replace("norm.bias", "group_norm.bias")
+
+    new_item = new_item.replace("q.weight", "query.weight")
+    new_item = new_item.replace("q.bias", "query.bias")
+
+    new_item = new_item.replace("k.weight", "key.weight")
+    new_item = new_item.replace("k.bias", "key.bias")
+
+    new_item = new_item.replace("v.weight", "value.weight")
+    new_item = new_item.replace("v.bias", "value.bias")
+
+    new_item = new_item.replace("proj_out.weight", "proj_attn.weight")
+    new_item = new_item.replace("proj_out.bias", "proj_attn.bias")
+
+    new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+    mapping.append({"old": old_item, "new": new_item})
+
+  return mapping
+
+
+def assign_to_checkpoint(
+    paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None
+):
+  """
+  This does the final conversion step: take locally converted weights and apply a global renaming
+  to them. It splits attention layers, and takes into account additional replacements
+  that may arise.
+
+  Assigns the weights to the new checkpoint.
+  """
+  assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys."
+
+  # Splits the attention layers into three variables.
+  if attention_paths_to_split is not None:
+    for path, path_map in attention_paths_to_split.items():
+      old_tensor = old_checkpoint[path]
+      channels = old_tensor.shape[0] // 3
+
+      target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1)
+
+      num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3
+
+      old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:])
+      query, key, value = old_tensor.split(channels // num_heads, dim=1)
+
+      checkpoint[path_map["query"]] = query.reshape(target_shape)
+      checkpoint[path_map["key"]] = key.reshape(target_shape)
+      checkpoint[path_map["value"]] = value.reshape(target_shape)
+
+  for path in paths:
+    new_path = path["new"]
+
+    # These have already been assigned
+    if attention_paths_to_split is not None and new_path in attention_paths_to_split:
+      continue
+
+    # Global renaming happens here
+    new_path = new_path.replace("middle_block.0", "mid_block.resnets.0")
+    new_path = new_path.replace("middle_block.1", "mid_block.attentions.0")
+    new_path = new_path.replace("middle_block.2", "mid_block.resnets.1")
+
+    if additional_replacements is not None:
+      for replacement in additional_replacements:
+        new_path = new_path.replace(replacement["old"], replacement["new"])
+
+    # proj_attn.weight has to be converted from conv 1D to linear
+    if "proj_attn.weight" in new_path:
+      checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0]
+    else:
+      checkpoint[new_path] = old_checkpoint[path["old"]]
+
+
+def conv_attn_to_linear(checkpoint):
+  keys = list(checkpoint.keys())
+  attn_keys = ["query.weight", "key.weight", "value.weight"]
+  for key in keys:
+    if ".".join(key.split(".")[-2:]) in attn_keys:
+      if checkpoint[key].ndim > 2:
+        checkpoint[key] = checkpoint[key][:, :, 0, 0]
+    elif "proj_attn.weight" in key:
+      if checkpoint[key].ndim > 2:
+        checkpoint[key] = checkpoint[key][:, :, 0]
+
+
+def linear_transformer_to_conv(checkpoint):
+  keys = list(checkpoint.keys())
+  tf_keys = ["proj_in.weight", "proj_out.weight"]
+  for key in keys:
+    if ".".join(key.split(".")[-2:]) in tf_keys:
+      if checkpoint[key].ndim == 2:
+        checkpoint[key] = checkpoint[key].unsqueeze(2).unsqueeze(2)
+
+
+def convert_ldm_unet_checkpoint(v2, checkpoint, config):
+  """
+  Takes a state dict and a config, and returns a converted checkpoint.
+  """
+
+  # extract state_dict for UNet
+  unet_state_dict = {}
+  unet_key = "model.diffusion_model."
+  keys = list(checkpoint.keys())
+  for key in keys:
+    if key.startswith(unet_key):
+      unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key)
+
+  new_checkpoint = {}
+
+  new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"]
+  new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"]
+  new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"]
+  new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"]
+
+  new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"]
+  new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"]
+
+  new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"]
+  new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"]
+  new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"]
+  new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"]
+
+  # Retrieves the keys for the input blocks only
+  num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer})
+  input_blocks = {
+      layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}." in key]
+      for layer_id in range(num_input_blocks)
+  }
+
+  # Retrieves the keys for the middle blocks only
+  num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer})
+  middle_blocks = {
+      layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}." in key]
+      for layer_id in range(num_middle_blocks)
+  }
+
+  # Retrieves the keys for the output blocks only
+  num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer})
+  output_blocks = {
+      layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}." in key]
+      for layer_id in range(num_output_blocks)
+  }
+
+  for i in range(1, num_input_blocks):
+    block_id = (i - 1) // (config["layers_per_block"] + 1)
+    layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1)
+
+    resnets = [
+        key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key
+    ]
+    attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key]
+
+    if f"input_blocks.{i}.0.op.weight" in unet_state_dict:
+      new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop(
+          f"input_blocks.{i}.0.op.weight"
+      )
+      new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop(
+          f"input_blocks.{i}.0.op.bias"
+      )
+
+    paths = renew_resnet_paths(resnets)
+    meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"}
+    assign_to_checkpoint(
+        paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+    )
+
+    if len(attentions):
+      paths = renew_attention_paths(attentions)
+      meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"}
+      assign_to_checkpoint(
+          paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+      )
+
+  resnet_0 = middle_blocks[0]
+  attentions = middle_blocks[1]
+  resnet_1 = middle_blocks[2]
+
+  resnet_0_paths = renew_resnet_paths(resnet_0)
+  assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config)
+
+  resnet_1_paths = renew_resnet_paths(resnet_1)
+  assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config)
+
+  attentions_paths = renew_attention_paths(attentions)
+  meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"}
+  assign_to_checkpoint(
+      attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+  )
+
+  for i in range(num_output_blocks):
+    block_id = i // (config["layers_per_block"] + 1)
+    layer_in_block_id = i % (config["layers_per_block"] + 1)
+    output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]]
+    output_block_list = {}
+
+    for layer in output_block_layers:
+      layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1)
+      if layer_id in output_block_list:
+        output_block_list[layer_id].append(layer_name)
+      else:
+        output_block_list[layer_id] = [layer_name]
+
+    if len(output_block_list) > 1:
+      resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key]
+      attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key]
+
+      resnet_0_paths = renew_resnet_paths(resnets)
+      paths = renew_resnet_paths(resnets)
+
+      meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"}
+      assign_to_checkpoint(
+          paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+      )
+
+      # オリジナル：
+      # if ["conv.weight", "conv.bias"] in output_block_list.values():
+      #   index = list(output_block_list.values()).index(["conv.weight", "conv.bias"])
+
+      # biasとweightの順番に依存しないようにする：もっといいやり方がありそうだが
+      for l in output_block_list.values():
+        l.sort()
+
+      if ["conv.bias", "conv.weight"] in output_block_list.values():
+        index = list(output_block_list.values()).index(["conv.bias", "conv.weight"])
+        new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[
+            f"output_blocks.{i}.{index}.conv.bias"
+        ]
+        new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[
+            f"output_blocks.{i}.{index}.conv.weight"
+        ]
+
+        # Clear attentions as they have been attributed above.
+        if len(attentions) == 2:
+          attentions = []
+
+      if len(attentions):
+        paths = renew_attention_paths(attentions)
+        meta_path = {
+            "old": f"output_blocks.{i}.1",
+            "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}",
+        }
+        assign_to_checkpoint(
+            paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+        )
+    else:
+      resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1)
+      for path in resnet_0_paths:
+        old_path = ".".join(["output_blocks", str(i), path["old"]])
+        new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]])
+
+        new_checkpoint[new_path] = unet_state_dict[old_path]
+
+  # SDのv2では1*1のconv2dがlinearに変わっているので、linear->convに変換する
+  if v2:
+    linear_transformer_to_conv(new_checkpoint)
+
+  return new_checkpoint
+
+
+def convert_ldm_vae_checkpoint(checkpoint, config):
+  # extract state dict for VAE
+  vae_state_dict = {}
+  vae_key = "first_stage_model."
+  keys = list(checkpoint.keys())
+  for key in keys:
+    if key.startswith(vae_key):
+      vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key)
+  # if len(vae_state_dict) == 0:
+  #   # 渡されたcheckpointは.ckptから読み込んだcheckpointではなくvaeのstate_dict
+  #   vae_state_dict = checkpoint
+
+  new_checkpoint = {}
+
+  new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]
+  new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"]
+  new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"]
+  new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"]
+  new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"]
+  new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"]
+
+  new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"]
+  new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"]
+  new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"]
+  new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"]
+  new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"]
+  new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"]
+
+  new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]
+  new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]
+  new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"]
+  new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"]
+
+  # Retrieves the keys for the encoder down blocks only
+  num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer})
+  down_blocks = {
+      layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks)
+  }
+
+  # Retrieves the keys for the decoder up blocks only
+  num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer})
+  up_blocks = {
+      layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks)
+  }
+
+  for i in range(num_down_blocks):
+    resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key]
+
+    if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:
+      new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop(
+          f"encoder.down.{i}.downsample.conv.weight"
+      )
+      new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop(
+          f"encoder.down.{i}.downsample.conv.bias"
+      )
+
+    paths = renew_vae_resnet_paths(resnets)
+    meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+  mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]
+  num_mid_res_blocks = 2
+  for i in range(1, num_mid_res_blocks + 1):
+    resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key]
+
+    paths = renew_vae_resnet_paths(resnets)
+    meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+  mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key]
+  paths = renew_vae_attention_paths(mid_attentions)
+  meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+  assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+  conv_attn_to_linear(new_checkpoint)
+
+  for i in range(num_up_blocks):
+    block_id = num_up_blocks - 1 - i
+    resnets = [
+        key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key
+    ]
+
+    if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:
+      new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[
+          f"decoder.up.{block_id}.upsample.conv.weight"
+      ]
+      new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[
+          f"decoder.up.{block_id}.upsample.conv.bias"
+      ]
+
+    paths = renew_vae_resnet_paths(resnets)
+    meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+  mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]
+  num_mid_res_blocks = 2
+  for i in range(1, num_mid_res_blocks + 1):
+    resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key]
+
+    paths = renew_vae_resnet_paths(resnets)
+    meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+  mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key]
+  paths = renew_vae_attention_paths(mid_attentions)
+  meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+  assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+  conv_attn_to_linear(new_checkpoint)
+  return new_checkpoint
+
+
+def create_unet_diffusers_config(v2):
+  """
+  Creates a config for the diffusers based on the config of the LDM model.
+  """
+  # unet_params = original_config.model.params.unet_config.params
+
+  block_out_channels = [UNET_PARAMS_MODEL_CHANNELS * mult for mult in UNET_PARAMS_CHANNEL_MULT]
+
+  down_block_types = []
+  resolution = 1
+  for i in range(len(block_out_channels)):
+    block_type = "CrossAttnDownBlock2D" if resolution in UNET_PARAMS_ATTENTION_RESOLUTIONS else "DownBlock2D"
+    down_block_types.append(block_type)
+    if i != len(block_out_channels) - 1:
+      resolution *= 2
+
+  up_block_types = []
+  for i in range(len(block_out_channels)):
+    block_type = "CrossAttnUpBlock2D" if resolution in UNET_PARAMS_ATTENTION_RESOLUTIONS else "UpBlock2D"
+    up_block_types.append(block_type)
+    resolution //= 2
+
+  config = dict(
+      sample_size=UNET_PARAMS_IMAGE_SIZE,
+      in_channels=UNET_PARAMS_IN_CHANNELS,
+      out_channels=UNET_PARAMS_OUT_CHANNELS,
+      down_block_types=tuple(down_block_types),
+      up_block_types=tuple(up_block_types),
+      block_out_channels=tuple(block_out_channels),
+      layers_per_block=UNET_PARAMS_NUM_RES_BLOCKS,
+      cross_attention_dim=UNET_PARAMS_CONTEXT_DIM if not v2 else V2_UNET_PARAMS_CONTEXT_DIM,
+      attention_head_dim=UNET_PARAMS_NUM_HEADS if not v2 else V2_UNET_PARAMS_ATTENTION_HEAD_DIM,
+  )
+
+  return config
+
+
+def create_vae_diffusers_config():
+  """
+  Creates a config for the diffusers based on the config of the LDM model.
+  """
+  # vae_params = original_config.model.params.first_stage_config.params.ddconfig
+  # _ = original_config.model.params.first_stage_config.params.embed_dim
+  block_out_channels = [VAE_PARAMS_CH * mult for mult in VAE_PARAMS_CH_MULT]
+  down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels)
+  up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels)
+
+  config = dict(
+      sample_size=VAE_PARAMS_RESOLUTION,
+      in_channels=VAE_PARAMS_IN_CHANNELS,
+      out_channels=VAE_PARAMS_OUT_CH,
+      down_block_types=tuple(down_block_types),
+      up_block_types=tuple(up_block_types),
+      block_out_channels=tuple(block_out_channels),
+      latent_channels=VAE_PARAMS_Z_CHANNELS,
+      layers_per_block=VAE_PARAMS_NUM_RES_BLOCKS,
+  )
+  return config
+
+
+def convert_ldm_clip_checkpoint_v1(checkpoint):
+  keys = list(checkpoint.keys())
+  text_model_dict = {}
+  for key in keys:
+    if key.startswith("cond_stage_model.transformer"):
+      text_model_dict[key[len("cond_stage_model.transformer."):]] = checkpoint[key]
+  return text_model_dict
+
+
+def convert_ldm_clip_checkpoint_v2(checkpoint, max_length):
+  # 嫌になるくらい違うぞ！
+  def convert_key(key):
+    if not key.startswith("cond_stage_model"):
+      return None
+
+    # common conversion
+    key = key.replace("cond_stage_model.model.transformer.", "text_model.encoder.")
+    key = key.replace("cond_stage_model.model.", "text_model.")
+
+    if "resblocks" in key:
+      # resblocks conversion
+      key = key.replace(".resblocks.", ".layers.")
+      if ".ln_" in key:
+        key = key.replace(".ln_", ".layer_norm")
+      elif ".mlp." in key:
+        key = key.replace(".c_fc.", ".fc1.")
+        key = key.replace(".c_proj.", ".fc2.")
+      elif '.attn.out_proj' in key:
+        key = key.replace(".attn.out_proj.", ".self_attn.out_proj.")
+      elif '.attn.in_proj' in key:
+        key = None                  # 特殊なので後で処理する
+      else:
+        raise ValueError(f"unexpected key in SD: {key}")
+    elif '.positional_embedding' in key:
+      key = key.replace(".positional_embedding", ".embeddings.position_embedding.weight")
+    elif '.text_projection' in key:
+      key = None    # 使われない???
+    elif '.logit_scale' in key:
+      key = None    # 使われない???
+    elif '.token_embedding' in key:
+      key = key.replace(".token_embedding.weight", ".embeddings.token_embedding.weight")
+    elif '.ln_final' in key:
+      key = key.replace(".ln_final", ".final_layer_norm")
+    return key
+
+  keys = list(checkpoint.keys())
+  new_sd = {}
+  for key in keys:
+    # remove resblocks 23
+    if '.resblocks.23.' in key:
+      continue
+    new_key = convert_key(key)
+    if new_key is None:
+      continue
+    new_sd[new_key] = checkpoint[key]
+
+  # attnの変換
+  for key in keys:
+    if '.resblocks.23.' in key:
+      continue
+    if '.resblocks' in key and '.attn.in_proj_' in key:
+      # 三つに分割
+      values = torch.chunk(checkpoint[key], 3)
+
+      key_suffix = ".weight" if "weight" in key else ".bias"
+      key_pfx = key.replace("cond_stage_model.model.transformer.resblocks.", "text_model.encoder.layers.")
+      key_pfx = key_pfx.replace("_weight", "")
+      key_pfx = key_pfx.replace("_bias", "")
+      key_pfx = key_pfx.replace(".attn.in_proj", ".self_attn.")
+      new_sd[key_pfx + "q_proj" + key_suffix] = values[0]
+      new_sd[key_pfx + "k_proj" + key_suffix] = values[1]
+      new_sd[key_pfx + "v_proj" + key_suffix] = values[2]
+
+  # rename or add position_ids
+  ANOTHER_POSITION_IDS_KEY = "text_model.encoder.text_model.embeddings.position_ids"
+  if ANOTHER_POSITION_IDS_KEY in new_sd:
+    # waifu diffusion v1.4
+    position_ids = new_sd[ANOTHER_POSITION_IDS_KEY]
+    del new_sd[ANOTHER_POSITION_IDS_KEY]
+  else:
+    position_ids = torch.Tensor([list(range(max_length))]).to(torch.int64)
+
+  new_sd["text_model.embeddings.position_ids"] = position_ids
+  return new_sd
+
+# endregion
+
+
+# region Diffusers->StableDiffusion の変換コード
+# convert_diffusers_to_original_stable_diffusion をコピーして修正している（ASL 2.0）
+
+def conv_transformer_to_linear(checkpoint):
+  keys = list(checkpoint.keys())
+  tf_keys = ["proj_in.weight", "proj_out.weight"]
+  for key in keys:
+    if ".".join(key.split(".")[-2:]) in tf_keys:
+      if checkpoint[key].ndim > 2:
+        checkpoint[key] = checkpoint[key][:, :, 0, 0]
+
+
+def convert_unet_state_dict_to_sd(v2, unet_state_dict):
+  unet_conversion_map = [
+      # (stable-diffusion, HF Diffusers)
+      ("time_embed.0.weight", "time_embedding.linear_1.weight"),
+      ("time_embed.0.bias", "time_embedding.linear_1.bias"),
+      ("time_embed.2.weight", "time_embedding.linear_2.weight"),
+      ("time_embed.2.bias", "time_embedding.linear_2.bias"),
+      ("input_blocks.0.0.weight", "conv_in.weight"),
+      ("input_blocks.0.0.bias", "conv_in.bias"),
+      ("out.0.weight", "conv_norm_out.weight"),
+      ("out.0.bias", "conv_norm_out.bias"),
+      ("out.2.weight", "conv_out.weight"),
+      ("out.2.bias", "conv_out.bias"),
+  ]
+
+  unet_conversion_map_resnet = [
+      # (stable-diffusion, HF Diffusers)
+      ("in_layers.0", "norm1"),
+      ("in_layers.2", "conv1"),
+      ("out_layers.0", "norm2"),
+      ("out_layers.3", "conv2"),
+      ("emb_layers.1", "time_emb_proj"),
+      ("skip_connection", "conv_shortcut"),
+  ]
+
+  unet_conversion_map_layer = []
+  for i in range(4):
+      # loop over downblocks/upblocks
+
+    for j in range(2):
+        # loop over resnets/attentions for downblocks
+      hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
+      sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
+      unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
+
+      if i < 3:
+        # no attention layers in down_blocks.3
+        hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
+        sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
+        unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+
+    for j in range(3):
+      # loop over resnets/attentions for upblocks
+      hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
+      sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
+      unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
+
+      if i > 0:
+        # no attention layers in up_blocks.0
+        hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
+        sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
+        unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
+
+    if i < 3:
+      # no downsample in down_blocks.3
+      hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
+      sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
+      unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
+
+      # no upsample in up_blocks.3
+      hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+      sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."
+      unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
+
+  hf_mid_atn_prefix = "mid_block.attentions.0."
+  sd_mid_atn_prefix = "middle_block.1."
+  unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
+
+  for j in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{j}."
+    sd_mid_res_prefix = f"middle_block.{2*j}."
+    unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+  # buyer beware: this is a *brittle* function,
+  # and correct output requires that all of these pieces interact in
+  # the exact order in which I have arranged them.
+  mapping = {k: k for k in unet_state_dict.keys()}
+  for sd_name, hf_name in unet_conversion_map:
+    mapping[hf_name] = sd_name
+  for k, v in mapping.items():
+    if "resnets" in k:
+      for sd_part, hf_part in unet_conversion_map_resnet:
+        v = v.replace(hf_part, sd_part)
+      mapping[k] = v
+  for k, v in mapping.items():
+    for sd_part, hf_part in unet_conversion_map_layer:
+      v = v.replace(hf_part, sd_part)
+    mapping[k] = v
+  new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()}
+
+  if v2:
+    conv_transformer_to_linear(new_state_dict)
+
+  return new_state_dict
+
+
+# ================#
+# VAE Conversion #
+# ================#
+
+def reshape_weight_for_sd(w):
+    # convert HF linear weights to SD conv2d weights
+  return w.reshape(*w.shape, 1, 1)
+
+
+def convert_vae_state_dict(vae_state_dict):
+  vae_conversion_map = [
+      # (stable-diffusion, HF Diffusers)
+      ("nin_shortcut", "conv_shortcut"),
+      ("norm_out", "conv_norm_out"),
+      ("mid.attn_1.", "mid_block.attentions.0."),
+  ]
+
+  for i in range(4):
+    # down_blocks have two resnets
+    for j in range(2):
+      hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
+      sd_down_prefix = f"encoder.down.{i}.block.{j}."
+      vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
+
+    if i < 3:
+      hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
+      sd_downsample_prefix = f"down.{i}.downsample."
+      vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
+
+      hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+      sd_upsample_prefix = f"up.{3-i}.upsample."
+      vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
+
+    # up_blocks have three resnets
+    # also, up blocks in hf are numbered in reverse from sd
+    for j in range(3):
+      hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
+      sd_up_prefix = f"decoder.up.{3-i}.block.{j}."
+      vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
+
+  # this part accounts for mid blocks in both the encoder and the decoder
+  for i in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{i}."
+    sd_mid_res_prefix = f"mid.block_{i+1}."
+    vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+  vae_conversion_map_attn = [
+      # (stable-diffusion, HF Diffusers)
+      ("norm.", "group_norm."),
+      ("q.", "query."),
+      ("k.", "key."),
+      ("v.", "value."),
+      ("proj_out.", "proj_attn."),
+  ]
+
+  mapping = {k: k for k in vae_state_dict.keys()}
+  for k, v in mapping.items():
+    for sd_part, hf_part in vae_conversion_map:
+      v = v.replace(hf_part, sd_part)
+    mapping[k] = v
+  for k, v in mapping.items():
+    if "attentions" in k:
+      for sd_part, hf_part in vae_conversion_map_attn:
+        v = v.replace(hf_part, sd_part)
+      mapping[k] = v
+  new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
+  weights_to_convert = ["q", "k", "v", "proj_out"]
+  for k, v in new_state_dict.items():
+    for weight_name in weights_to_convert:
+      if f"mid.attn_1.{weight_name}.weight" in k:
+        # print(f"Reshaping {k} for SD format")
+        new_state_dict[k] = reshape_weight_for_sd(v)
+
+  return new_state_dict
+
+
+# endregion
+
+# region 自作のモデル読み書きなど
+
+def is_safetensors(path):
+  return os.path.splitext(path)[1].lower() == '.safetensors'
+
+
+def load_checkpoint_with_text_encoder_conversion(ckpt_path):
+  # text encoderの格納形式が違うモデルに対応する ('text_model'がない)
+  TEXT_ENCODER_KEY_REPLACEMENTS = [
+      ('cond_stage_model.transformer.embeddings.', 'cond_stage_model.transformer.text_model.embeddings.'),
+      ('cond_stage_model.transformer.encoder.', 'cond_stage_model.transformer.text_model.encoder.'),
+      ('cond_stage_model.transformer.final_layer_norm.', 'cond_stage_model.transformer.text_model.final_layer_norm.')
+  ]
+
+  if is_safetensors(ckpt_path):
+    checkpoint = None
+    state_dict = load_file(ckpt_path, "cpu")
+  else:
+    checkpoint = torch.load(ckpt_path, map_location="cpu")
+    if "state_dict" in checkpoint:
+      state_dict = checkpoint["state_dict"]
+    else:
+      state_dict = checkpoint
+      checkpoint = None
+
+  key_reps = []
+  for rep_from, rep_to in TEXT_ENCODER_KEY_REPLACEMENTS:
+    for key in state_dict.keys():
+      if key.startswith(rep_from):
+        new_key = rep_to + key[len(rep_from):]
+        key_reps.append((key, new_key))
+
+  for key, new_key in key_reps:
+    state_dict[new_key] = state_dict[key]
+    del state_dict[key]
+
+  return checkpoint, state_dict
+
+
+# TODO dtype指定の動作が怪しいので確認する text_encoderを指定形式で作れるか未確認
+def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, dtype=None):
+  _, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path)
+  if dtype is not None:
+    for k, v in state_dict.items():
+      if type(v) is torch.Tensor:
+        state_dict[k] = v.to(dtype)
+
+  # Convert the UNet2DConditionModel model.
+  unet_config = create_unet_diffusers_config(v2)
+  converted_unet_checkpoint = convert_ldm_unet_checkpoint(v2, state_dict, unet_config)
+
+  unet = UNet2DConditionModel(**unet_config)
+  info = unet.load_state_dict(converted_unet_checkpoint)
+  print("loading u-net:", info)
+
+  # Convert the VAE model.
+  vae_config = create_vae_diffusers_config()
+  converted_vae_checkpoint = convert_ldm_vae_checkpoint(state_dict, vae_config)
+
+  vae = AutoencoderKL(**vae_config)
+  info = vae.load_state_dict(converted_vae_checkpoint)
+  print("loading vae:", info)
+
+  # convert text_model
+  if v2:
+    converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v2(state_dict, 77)
+    cfg = CLIPTextConfig(
+        vocab_size=49408,
+        hidden_size=1024,
+        intermediate_size=4096,
+        num_hidden_layers=23,
+        num_attention_heads=16,
+        max_position_embeddings=77,
+        hidden_act="gelu",
+        layer_norm_eps=1e-05,
+        dropout=0.0,
+        attention_dropout=0.0,
+        initializer_range=0.02,
+        initializer_factor=1.0,
+        pad_token_id=1,
+        bos_token_id=0,
+        eos_token_id=2,
+        model_type="clip_text_model",
+        projection_dim=512,
+        torch_dtype="float32",
+        transformers_version="4.25.0.dev0",
+    )
+    text_model = CLIPTextModel._from_config(cfg)
+    info = text_model.load_state_dict(converted_text_encoder_checkpoint)
+  else:
+    converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v1(state_dict)
+    text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+    info = text_model.load_state_dict(converted_text_encoder_checkpoint)
+  print("loading text encoder:", info)
+
+  return text_model, vae, unet
+
+
+def convert_text_encoder_state_dict_to_sd_v2(checkpoint, make_dummy_weights=False):
+  def convert_key(key):
+    # position_idsの除去
+    if ".position_ids" in key:
+      return None
+
+    # common
+    key = key.replace("text_model.encoder.", "transformer.")
+    key = key.replace("text_model.", "")
+    if "layers" in key:
+      # resblocks conversion
+      key = key.replace(".layers.", ".resblocks.")
+      if ".layer_norm" in key:
+        key = key.replace(".layer_norm", ".ln_")
+      elif ".mlp." in key:
+        key = key.replace(".fc1.", ".c_fc.")
+        key = key.replace(".fc2.", ".c_proj.")
+      elif '.self_attn.out_proj' in key:
+        key = key.replace(".self_attn.out_proj.", ".attn.out_proj.")
+      elif '.self_attn.' in key:
+        key = None                  # 特殊なので後で処理する
+      else:
+        raise ValueError(f"unexpected key in DiffUsers model: {key}")
+    elif '.position_embedding' in key:
+      key = key.replace("embeddings.position_embedding.weight", "positional_embedding")
+    elif '.token_embedding' in key:
+      key = key.replace("embeddings.token_embedding.weight", "token_embedding.weight")
+    elif 'final_layer_norm' in key:
+      key = key.replace("final_layer_norm", "ln_final")
+    return key
+
+  keys = list(checkpoint.keys())
+  new_sd = {}
+  for key in keys:
+    new_key = convert_key(key)
+    if new_key is None:
+      continue
+    new_sd[new_key] = checkpoint[key]
+
+  # attnの変換
+  for key in keys:
+    if 'layers' in key and 'q_proj' in key:
+      # 三つを結合
+      key_q = key
+      key_k = key.replace("q_proj", "k_proj")
+      key_v = key.replace("q_proj", "v_proj")
+
+      value_q = checkpoint[key_q]
+      value_k = checkpoint[key_k]
+      value_v = checkpoint[key_v]
+      value = torch.cat([value_q, value_k, value_v])
+
+      new_key = key.replace("text_model.encoder.layers.", "transformer.resblocks.")
+      new_key = new_key.replace(".self_attn.q_proj.", ".attn.in_proj_")
+      new_sd[new_key] = value
+
+  # 最後の層などを捏造するか
+  if make_dummy_weights:
+    print("make dummy weights for resblock.23, text_projection and logit scale.")
+    keys = list(new_sd.keys())
+    for key in keys:
+      if key.startswith("transformer.resblocks.22."):
+        new_sd[key.replace(".22.", ".23.")] = new_sd[key].clone()          # copyしないとsafetensorsの保存で落ちる
+
+    # Diffusersに含まれない重みを作っておく
+    new_sd['text_projection'] = torch.ones((1024, 1024), dtype=new_sd[keys[0]].dtype, device=new_sd[keys[0]].device)
+    new_sd['logit_scale'] = torch.tensor(1)
+
+  return new_sd
+
+
+def save_stable_diffusion_checkpoint(v2, output_file, text_encoder, unet, ckpt_path, epochs, steps, save_dtype=None, vae=None):
+  if ckpt_path is not None:
+    # epoch/stepを参照する。またVAEがメモリ上にないときなど、もう一度VAEを含めて読み込む
+    checkpoint, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path)
+    if checkpoint is None:                # safetensors または state_dictのckpt
+      checkpoint = {}
+      strict = False
+    else:
+      strict = True
+    if "state_dict" in state_dict:
+      del state_dict["state_dict"]
+  else:
+    # 新しく作る
+    assert vae is not None, "VAE is required to save a checkpoint without a given checkpoint"
+    checkpoint = {}
+    state_dict = {}
+    strict = False
+
+  def update_sd(prefix, sd):
+    for k, v in sd.items():
+      key = prefix + k
+      assert not strict or key in state_dict, f"Illegal key in save SD: {key}"
+      if save_dtype is not None:
+        v = v.detach().clone().to("cpu").to(save_dtype)
+      state_dict[key] = v
+
+  # Convert the UNet model
+  unet_state_dict = convert_unet_state_dict_to_sd(v2, unet.state_dict())
+  update_sd("model.diffusion_model.", unet_state_dict)
+
+  # Convert the text encoder model
+  if v2:
+    make_dummy = ckpt_path is None                 # 参照元のcheckpointがない場合は最後の層を前の層から複製��て作るなどダミーの重みを入れる
+    text_enc_dict = convert_text_encoder_state_dict_to_sd_v2(text_encoder.state_dict(), make_dummy)
+    update_sd("cond_stage_model.model.", text_enc_dict)
+  else:
+    text_enc_dict = text_encoder.state_dict()
+    update_sd("cond_stage_model.transformer.", text_enc_dict)
+
+  # Convert the VAE
+  if vae is not None:
+    vae_dict = convert_vae_state_dict(vae.state_dict())
+    update_sd("first_stage_model.", vae_dict)
+
+  # Put together new checkpoint
+  key_count = len(state_dict.keys())
+  new_ckpt = {'state_dict': state_dict}
+
+  if 'epoch' in checkpoint:
+    epochs += checkpoint['epoch']
+  if 'global_step' in checkpoint:
+    steps += checkpoint['global_step']
+
+  new_ckpt['epoch'] = epochs
+  new_ckpt['global_step'] = steps
+
+  if is_safetensors(output_file):
+    # TODO Tensor以外のdictの値を削除したほうがいいか
+    save_file(state_dict, output_file)
+  else:
+    torch.save(new_ckpt, output_file)
+
+  return key_count
+
+
+def save_diffusers_checkpoint(v2, output_dir, text_encoder, unet, pretrained_model_name_or_path, vae=None, use_safetensors=False):
+  if pretrained_model_name_or_path is None:
+    # load default settings for v1/v2
+    if v2:
+      pretrained_model_name_or_path = DIFFUSERS_REF_MODEL_ID_V2
+    else:
+      pretrained_model_name_or_path = DIFFUSERS_REF_MODEL_ID_V1
+
+  scheduler = DDIMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
+  tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
+  if vae is None:
+    vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
+
+  pipeline = StableDiffusionPipeline(
+      unet=unet,
+      text_encoder=text_encoder,
+      vae=vae,
+      scheduler=scheduler,
+      tokenizer=tokenizer,
+      safety_checker=None,
+      feature_extractor=None,
+      requires_safety_checker=None,
+  )
+  pipeline.save_pretrained(output_dir, safe_serialization=use_safetensors)
+
+
+VAE_PREFIX = "first_stage_model."
+
+
+def load_vae(vae_id, dtype):
+  print(f"load VAE: {vae_id}")
+  if os.path.isdir(vae_id) or not os.path.isfile(vae_id):
+    # Diffusers local/remote
+    try:
+      vae = AutoencoderKL.from_pretrained(vae_id, subfolder=None, torch_dtype=dtype)
+    except EnvironmentError as e:
+      print(f"exception occurs in loading vae: {e}")
+      print("retry with subfolder='vae'")
+      vae = AutoencoderKL.from_pretrained(vae_id, subfolder="vae", torch_dtype=dtype)
+    return vae
+
+  # local
+  vae_config = create_vae_diffusers_config()
+
+  if vae_id.endswith(".bin"):
+    # SD 1.5 VAE on Huggingface
+    converted_vae_checkpoint = torch.load(vae_id, map_location="cpu")
+  else:
+    # StableDiffusion
+    vae_model = (load_file(vae_id, "cpu") if is_safetensors(vae_id)
+                 else torch.load(vae_id, map_location="cpu"))
+    vae_sd = vae_model['state_dict'] if 'state_dict' in vae_model else vae_model
+
+    # vae only or full model
+    full_model = False
+    for vae_key in vae_sd:
+      if vae_key.startswith(VAE_PREFIX):
+        full_model = True
+        break
+    if not full_model:
+      sd = {}
+      for key, value in vae_sd.items():
+        sd[VAE_PREFIX + key] = value
+      vae_sd = sd
+      del sd
+
+    # Convert the VAE model.
+    converted_vae_checkpoint = convert_ldm_vae_checkpoint(vae_sd, vae_config)
+
+  vae = AutoencoderKL(**vae_config)
+  vae.load_state_dict(converted_vae_checkpoint)
+  return vae
+
+# endregion
+
+
+def make_bucket_resolutions(max_reso, min_size=256, max_size=1024, divisible=64):
+  max_width, max_height = max_reso
+  max_area = (max_width // divisible) * (max_height // divisible)
+
+  resos = set()
+
+  size = int(math.sqrt(max_area)) * divisible
+  resos.add((size, size))
+
+  size = min_size
+  while size <= max_size:
+    width = size
+    height = min(max_size, (max_area // (width // divisible)) * divisible)
+    resos.add((width, height))
+    resos.add((height, width))
+
+    # # make additional resos
+    # if width >= height and width - divisible >= min_size:
+    #   resos.add((width - divisible, height))
+    #   resos.add((height, width - divisible))
+    # if height >= width and height - divisible >= min_size:
+    #   resos.add((width, height - divisible))
+    #   resos.add((height - divisible, width))
+
+    size += divisible
+
+  resos = list(resos)
+  resos.sort()
+  return resos
+
+
+if __name__ == '__main__':
+  resos = make_bucket_resolutions((512, 768))
+  print(len(resos))
+  print(resos)
+  aspect_ratios = [w / h for w, h in resos]
+  print(aspect_ratios)
+
+  ars = set()
+  for ar in aspect_ratios:
+    if ar in ars:
+      print("error! duplicate ar:", ar)
+    ars.add(ar)

build/lib/library/train_util.py

@@ -0,0 +1,1796 @@
+# common functions for training
+
+import argparse
+import json
+import shutil
+import time
+from typing import Dict, List, NamedTuple, Tuple
+from accelerate import Accelerator
+from torch.autograd.function import Function
+import glob
+import math
+import os
+import random
+import hashlib
+import subprocess
+from io import BytesIO
+
+from tqdm import tqdm
+import torch
+from torchvision import transforms
+from transformers import CLIPTokenizer
+import diffusers
+from diffusers import DDPMScheduler, StableDiffusionPipeline
+import albumentations as albu
+import numpy as np
+from PIL import Image
+import cv2
+from einops import rearrange
+from torch import einsum
+import safetensors.torch
+
+import library.model_util as model_util
+
+# Tokenizer: checkpointから読み込むのではなくあらかじめ提供されているものを使う
+TOKENIZER_PATH = "openai/clip-vit-large-patch14"
+V2_STABLE_DIFFUSION_PATH = "stabilityai/stable-diffusion-2"     # ここからtokenizerだけ使う v2とv2.1はtokenizer仕様は同じ
+
+# checkpointファイル名
+EPOCH_STATE_NAME = "{}-{:06d}-state"
+EPOCH_FILE_NAME = "{}-{:06d}"
+EPOCH_DIFFUSERS_DIR_NAME = "{}-{:06d}"
+LAST_STATE_NAME = "{}-state"
+DEFAULT_EPOCH_NAME = "epoch"
+DEFAULT_LAST_OUTPUT_NAME = "last"
+
+# region dataset
+
+IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".webp", ".bmp"]
+# , ".PNG", ".JPG", ".JPEG", ".WEBP", ".BMP"]         # Linux?
+
+
+class ImageInfo():
+  def __init__(self, image_key: str, num_repeats: int, caption: str, is_reg: bool, absolute_path: str) -> None:
+    self.image_key: str = image_key
+    self.num_repeats: int = num_repeats
+    self.caption: str = caption
+    self.is_reg: bool = is_reg
+    self.absolute_path: str = absolute_path
+    self.image_size: Tuple[int, int] = None
+    self.resized_size: Tuple[int, int] = None
+    self.bucket_reso: Tuple[int, int] = None
+    self.latents: torch.Tensor = None
+    self.latents_flipped: torch.Tensor = None
+    self.latents_npz: str = None
+    self.latents_npz_flipped: str = None
+
+
+class BucketManager():
+  def __init__(self, no_upscale, max_reso, min_size, max_size, reso_steps) -> None:
+    self.no_upscale = no_upscale
+    if max_reso is None:
+      self.max_reso = None
+      self.max_area = None
+    else:
+      self.max_reso = max_reso
+      self.max_area = max_reso[0] * max_reso[1]
+    self.min_size = min_size
+    self.max_size = max_size
+    self.reso_steps = reso_steps
+
+    self.resos = []
+    self.reso_to_id = {}
+    self.buckets = []                     # 前処理時は (image_key, image)、学習時は image_key
+
+  def add_image(self, reso, image):
+    bucket_id = self.reso_to_id[reso]
+    self.buckets[bucket_id].append(image)
+
+  def shuffle(self):
+    for bucket in self.buckets:
+      random.shuffle(bucket)
+
+  def sort(self):
+    # 解像度順にソートする（表示時、メタデータ格納時の見栄えをよくするためだけ）。bucketsも入れ替えてreso_to_idも振り直す
+    sorted_resos = self.resos.copy()
+    sorted_resos.sort()
+
+    sorted_buckets = []
+    sorted_reso_to_id = {}
+    for i, reso in enumerate(sorted_resos):
+      bucket_id = self.reso_to_id[reso]
+      sorted_buckets.append(self.buckets[bucket_id])
+      sorted_reso_to_id[reso] = i
+
+    self.resos = sorted_resos
+    self.buckets = sorted_buckets
+    self.reso_to_id = sorted_reso_to_id
+
+  def make_buckets(self):
+    resos = model_util.make_bucket_resolutions(self.max_reso, self.min_size, self.max_size, self.reso_steps)
+    self.set_predefined_resos(resos)
+
+  def set_predefined_resos(self, resos):
+    # 規定サイズから選ぶ場合の解像度、aspect ratioの情報を格納しておく
+    self.predefined_resos = resos.copy()
+    self.predefined_resos_set = set(resos)
+    self.predefined_aspect_ratios = np.array([w / h for w, h in resos])
+
+  def add_if_new_reso(self, reso):
+    if reso not in self.reso_to_id:
+      bucket_id = len(self.resos)
+      self.reso_to_id[reso] = bucket_id
+      self.resos.append(reso)
+      self.buckets.append([])
+      # print(reso, bucket_id, len(self.buckets))
+
+  def round_to_steps(self, x):
+    x = int(x + .5)
+    return x - x % self.reso_steps
+
+  def select_bucket(self, image_width, image_height):
+    aspect_ratio = image_width / image_height
+    if not self.no_upscale:
+      # 同じaspect ratioがあるかもしれないので（fine tuningで、no_upscale=Trueで前処理した場合）、解像度が同じものを優先する
+      reso = (image_width, image_height)
+      if reso in self.predefined_resos_set:
+        pass
+      else:
+        ar_errors = self.predefined_aspect_ratios - aspect_ratio
+        predefined_bucket_id = np.abs(ar_errors).argmin()          # 当該解像度以外でaspect ratio errorが最も少ないもの
+        reso = self.predefined_resos[predefined_bucket_id]
+
+      ar_reso = reso[0] / reso[1]
+      if aspect_ratio > ar_reso:                   # 横が長い→縦を合わせる
+        scale = reso[1] / image_height
+      else:
+        scale = reso[0] / image_width
+
+      resized_size = (int(image_width * scale + .5), int(image_height * scale + .5))
+      # print("use predef", image_width, image_height, reso, resized_size)
+    else:
+      if image_width * image_height > self.max_area:
+        # 画像が大きすぎるのでアスペクト比を保ったまま縮小することを前提にbucketを決める
+        resized_width = math.sqrt(self.max_area * aspect_ratio)
+        resized_height = self.max_area / resized_width
+        assert abs(resized_width / resized_height - aspect_ratio) < 1e-2, "aspect is illegal"
+
+        # リサイズ後の短辺または長辺をreso_steps単位にする：aspect ratioの差が少ないほうを選ぶ
+        # 元のbucketingと同じロジック
+        b_width_rounded = self.round_to_steps(resized_width)
+        b_height_in_wr = self.round_to_steps(b_width_rounded / aspect_ratio)
+        ar_width_rounded = b_width_rounded / b_height_in_wr
+
+        b_height_rounded = self.round_to_steps(resized_height)
+        b_width_in_hr = self.round_to_steps(b_height_rounded * aspect_ratio)
+        ar_height_rounded = b_width_in_hr / b_height_rounded
+
+        # print(b_width_rounded, b_height_in_wr, ar_width_rounded)
+        # print(b_width_in_hr, b_height_rounded, ar_height_rounded)
+
+        if abs(ar_width_rounded - aspect_ratio) < abs(ar_height_rounded - aspect_ratio):
+          resized_size = (b_width_rounded, int(b_width_rounded / aspect_ratio + .5))
+        else:
+          resized_size = (int(b_height_rounded * aspect_ratio + .5), b_height_rounded)
+        # print(resized_size)
+      else:
+        resized_size = (image_width, image_height)              # リサイズは不要
+
+      # 画像のサイズ未満をbucketのサイズとする（paddingせずにcroppingする）
+      bucket_width = resized_size[0] - resized_size[0] % self.reso_steps
+      bucket_height = resized_size[1] - resized_size[1] % self.reso_steps
+      # print("use arbitrary", image_width, image_height, resized_size, bucket_width, bucket_height)
+
+      reso = (bucket_width, bucket_height)
+
+    self.add_if_new_reso(reso)
+
+    ar_error = (reso[0] / reso[1]) - aspect_ratio
+    return reso, resized_size, ar_error
+
+
+class BucketBatchIndex(NamedTuple):
+  bucket_index: int
+  bucket_batch_size: int
+  batch_index: int
+
+
+class BaseDataset(torch.utils.data.Dataset):
+  def __init__(self, tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens, resolution, flip_aug: bool, color_aug: bool, face_crop_aug_range, random_crop, debug_dataset: bool) -> None:
+    super().__init__()
+    self.tokenizer: CLIPTokenizer = tokenizer
+    self.max_token_length = max_token_length
+    self.shuffle_caption = shuffle_caption
+    self.shuffle_keep_tokens = shuffle_keep_tokens
+    # width/height is used when enable_bucket==False
+    self.width, self.height = (None, None) if resolution is None else resolution
+    self.face_crop_aug_range = face_crop_aug_range
+    self.flip_aug = flip_aug
+    self.color_aug = color_aug
+    self.debug_dataset = debug_dataset
+    self.random_crop = random_crop
+    self.token_padding_disabled = False
+    self.dataset_dirs_info = {}
+    self.reg_dataset_dirs_info = {}
+    self.tag_frequency = {}
+
+    self.enable_bucket = False
+    self.bucket_manager: BucketManager = None                         # not initialized
+    self.min_bucket_reso = None
+    self.max_bucket_reso = None
+    self.bucket_reso_steps = None
+    self.bucket_no_upscale = None
+    self.bucket_info = None                                           # for metadata
+
+    self.tokenizer_max_length = self.tokenizer.model_max_length if max_token_length is None else max_token_length + 2
+
+    self.current_epoch: int = 0            # インスタンスがepochごとに新しく作られるようなので外側から渡さないとダメ
+    self.dropout_rate: float = 0
+    self.dropout_every_n_epochs: int = None
+    self.tag_dropout_rate: float = 0
+
+    # augmentation
+    flip_p = 0.5 if flip_aug else 0.0
+    if color_aug:
+      # わりと弱めの色合いaugmentation：brightness/contrastあたりは画像のpixel valueの最大値・最小値を変えてしまうのでよくないのではという想定でgamma/hueあたりを触る
+      self.aug = albu.Compose([
+          albu.OneOf([
+              albu.HueSaturationValue(8, 0, 0, p=.5),
+              albu.RandomGamma((95, 105), p=.5),
+          ], p=.33),
+          albu.HorizontalFlip(p=flip_p)
+      ], p=1.)
+    elif flip_aug:
+      self.aug = albu.Compose([
+          albu.HorizontalFlip(p=flip_p)
+      ], p=1.)
+    else:
+      self.aug = None
+
+    self.image_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5]), ])
+
+    self.image_data: Dict[str, ImageInfo] = {}
+
+    self.replacements = {}
+
+  def set_current_epoch(self, epoch):
+    self.current_epoch = epoch
+
+  def set_caption_dropout(self, dropout_rate, dropout_every_n_epochs, tag_dropout_rate):
+    # コンストラクタで渡さないのはTextual Inversionで意識したくないから（ということにしておく）
+    self.dropout_rate = dropout_rate
+    self.dropout_every_n_epochs = dropout_every_n_epochs
+    self.tag_dropout_rate = tag_dropout_rate
+
+  def set_tag_frequency(self, dir_name, captions):
+    frequency_for_dir = self.tag_frequency.get(dir_name, {})
+    self.tag_frequency[dir_name] = frequency_for_dir
+    for caption in captions:
+      for tag in caption.split(","):
+        if tag and not tag.isspace():
+          tag = tag.lower()
+          frequency = frequency_for_dir.get(tag, 0)
+          frequency_for_dir[tag] = frequency + 1
+
+  def disable_token_padding(self):
+    self.token_padding_disabled = True
+
+  def add_replacement(self, str_from, str_to):
+    self.replacements[str_from] = str_to
+
+  def process_caption(self, caption):
+    # dropoutの決定：tag dropがこのメソッド内にあるのでここで行うのが良い
+    is_drop_out = self.dropout_rate > 0 and random.random() < self.dropout_rate
+    is_drop_out = is_drop_out or self.dropout_every_n_epochs and self.current_epoch % self.dropout_every_n_epochs == 0
+
+    if is_drop_out:
+      caption = ""
+    else:
+      if self.shuffle_caption or self.tag_dropout_rate > 0:
+        def dropout_tags(tokens):
+          if self.tag_dropout_rate <= 0:
+            return tokens
+          l = []
+          for token in tokens:
+            if random.random() >= self.tag_dropout_rate:
+              l.append(token)
+          return l
+
+        tokens = [t.strip() for t in caption.strip().split(",")]
+        if self.shuffle_keep_tokens is None:
+          if self.shuffle_caption:
+            random.shuffle(tokens)
+
+          tokens = dropout_tags(tokens)
+        else:
+          if len(tokens) > self.shuffle_keep_tokens:
+            keep_tokens = tokens[:self.shuffle_keep_tokens]
+            tokens = tokens[self.shuffle_keep_tokens:]
+
+            if self.shuffle_caption:
+              random.shuffle(tokens)
+
+            tokens = dropout_tags(tokens)
+
+            tokens = keep_tokens + tokens
+        caption = ", ".join(tokens)
+
+      # textual inversion対応
+      for str_from, str_to in self.replacements.items():
+        if str_from == "":
+          # replace all
+          if type(str_to) == list:
+            caption = random.choice(str_to)
+          else:
+            caption = str_to
+        else:
+          caption = caption.replace(str_from, str_to)
+
+    return caption
+
+  def get_input_ids(self, caption):
+    input_ids = self.tokenizer(caption, padding="max_length", truncation=True,
+                               max_length=self.tokenizer_max_length, return_tensors="pt").input_ids
+
+    if self.tokenizer_max_length > self.tokenizer.model_max_length:
+      input_ids = input_ids.squeeze(0)
+      iids_list = []
+      if self.tokenizer.pad_token_id == self.tokenizer.eos_token_id:
+        # v1
+        # 77以上の時は "<BOS> .... <EOS> <EOS> <EOS>" でトータル227とかになっているので、"<BOS>...<EOS>"の三連に変換する
+        # 1111氏のやつは , で区切る、とかしているようだが　とりあえず単純に
+        for i in range(1, self.tokenizer_max_length - self.tokenizer.model_max_length + 2, self.tokenizer.model_max_length - 2):  # (1, 152, 75)
+          ids_chunk = (input_ids[0].unsqueeze(0),
+                       input_ids[i:i + self.tokenizer.model_max_length - 2],
+                       input_ids[-1].unsqueeze(0))
+          ids_chunk = torch.cat(ids_chunk)
+          iids_list.append(ids_chunk)
+      else:
+        # v2
+        # 77以上の時は "<BOS> .... <EOS> <PAD> <PAD>..." でトータル227とかになっているので、"<BOS>...<EOS> <PAD> <PAD> ..."の三連に変換する
+        for i in range(1, self.tokenizer_max_length - self.tokenizer.model_max_length + 2, self.tokenizer.model_max_length - 2):
+          ids_chunk = (input_ids[0].unsqueeze(0),       # BOS
+                       input_ids[i:i + self.tokenizer.model_max_length - 2],
+                       input_ids[-1].unsqueeze(0))      # PAD or EOS
+          ids_chunk = torch.cat(ids_chunk)
+
+          # 末尾が <EOS> <PAD> または <PAD> <PAD> の場合は、何もしなくてよい
+          # 末尾が x <PAD/EOS> の場合は末尾を <EOS> に変える（x <EOS> なら結果的に変化なし）
+          if ids_chunk[-2] != self.tokenizer.eos_token_id and ids_chunk[-2] != self.tokenizer.pad_token_id:
+            ids_chunk[-1] = self.tokenizer.eos_token_id
+          # 先頭が <BOS> <PAD> ... の場合は <BOS> <EOS> <PAD> ... に変える
+          if ids_chunk[1] == self.tokenizer.pad_token_id:
+            ids_chunk[1] = self.tokenizer.eos_token_id
+
+          iids_list.append(ids_chunk)
+
+      input_ids = torch.stack(iids_list)      # 3,77
+    return input_ids
+
+  def register_image(self, info: ImageInfo):
+    self.image_data[info.image_key] = info
+
+  def make_buckets(self):
+    '''
+    bucketingを行わない場合も呼び出し必須（ひとつだけbucketを作る）
+    min_size and max_size are ignored when enable_bucket is False
+    '''
+    print("loading image sizes.")
+    for info in tqdm(self.image_data.values()):
+      if info.image_size is None:
+        info.image_size = self.get_image_size(info.absolute_path)
+
+    if self.enable_bucket:
+      print("make buckets")
+    else:
+      print("prepare dataset")
+
+    # bucketを作成し、画像をbucketに振り分ける
+    if self.enable_bucket:
+      if self.bucket_manager is None:                         # fine tuningの場合でmetadataに定義がある場合は、すでに初期化済み
+        self.bucket_manager = BucketManager(self.bucket_no_upscale, (self.width, self.height),
+                                            self.min_bucket_reso, self.max_bucket_reso, self.bucket_reso_steps)
+        if not self.bucket_no_upscale:
+          self.bucket_manager.make_buckets()
+        else:
+          print("min_bucket_reso and max_bucket_reso are ignored if bucket_no_upscale is set, because bucket reso is defined by image size automatically / bucket_no_upscaleが指定された場合は、bucketの解像度は画像サイズから自動計算されるため、min_bucket_resoとmax_bucket_resoは無視されます")
+
+      img_ar_errors = []
+      for image_info in self.image_data.values():
+        image_width, image_height = image_info.image_size
+        image_info.bucket_reso, image_info.resized_size, ar_error = self.bucket_manager.select_bucket(image_width, image_height)
+
+        # print(image_info.image_key, image_info.bucket_reso)
+        img_ar_errors.append(abs(ar_error))
+
+      self.bucket_manager.sort()
+    else:
+      self.bucket_manager = BucketManager(False, (self.width, self.height), None, None, None)
+      self.bucket_manager.set_predefined_resos([(self.width, self.height)])  # ひとつの固定サイズbucketのみ
+      for image_info in self.image_data.values():
+        image_width, image_height = image_info.image_size
+        image_info.bucket_reso, image_info.resized_size, _ = self.bucket_manager.select_bucket(image_width, image_height)
+
+    for image_info in self.image_data.values():
+      for _ in range(image_info.num_repeats):
+        self.bucket_manager.add_image(image_info.bucket_reso, image_info.image_key)
+
+    # bucket情報を表示、格納する
+    if self.enable_bucket:
+      self.bucket_info = {"buckets": {}}
+      print("number of images (including repeats) / 各bucketの画像枚数（繰り返し回数を含む）")
+      for i, (reso, bucket) in enumerate(zip(self.bucket_manager.resos, self.bucket_manager.buckets)):
+        count = len(bucket)
+        if count > 0:
+          self.bucket_info["buckets"][i] = {"resolution": reso, "count": len(bucket)}
+          print(f"bucket {i}: resolution {reso}, count: {len(bucket)}")
+
+      img_ar_errors = np.array(img_ar_errors)
+      mean_img_ar_error = np.mean(np.abs(img_ar_errors))
+      self.bucket_info["mean_img_ar_error"] = mean_img_ar_error
+      print(f"mean ar error (without repeats): {mean_img_ar_error}")
+
+    # データ参照用indexを作る。このindexはdatasetのshuffleに用いられる
+    self.buckets_indices: List(BucketBatchIndex) = []
+    for bucket_index, bucket in enumerate(self.bucket_manager.buckets):
+      batch_count = int(math.ceil(len(bucket) / self.batch_size))
+      for batch_index in range(batch_count):
+        self.buckets_indices.append(BucketBatchIndex(bucket_index, self.batch_size, batch_index))
+
+      # ↓以下はbucketごとのbatch件数があまりにも増えて混乱を招くので元に戻す
+      # 　学習時はステップ数がランダムなので、同一画像が同一batch内にあってもそれほど悪影響はないであろう、と考えられる
+      #
+      # # bucketが細分化されることにより、ひとつのbucketに一種類の画像のみというケースが増え、つまりそれは
+      # # ひとつのbatchが同じ画像で占められることになるので、さすがに良くないであろう
+      # # そのためバッチサイズを画像種類までに制限する
+      # # ただそれでも同一画像が同一バッチに含まれる可能性はあるので、繰り返し回数が少ないほうがshuffleの品質は良くなることは間違いない？
+      # # TO DO 正則化画像をepochまたがりで利用する仕組み
+      # num_of_image_types = len(set(bucket))
+      # bucket_batch_size = min(self.batch_size, num_of_image_types)
+      # batch_count = int(math.ceil(len(bucket) / bucket_batch_size))
+      # # print(bucket_index, num_of_image_types, bucket_batch_size, batch_count)
+      # for batch_index in range(batch_count):
+      #   self.buckets_indices.append(BucketBatchIndex(bucket_index, bucket_batch_size, batch_index))
+      # ↑ここまで
+
+    self.shuffle_buckets()
+    self._length = len(self.buckets_indices)
+
+  def shuffle_buckets(self):
+    random.shuffle(self.buckets_indices)
+    self.bucket_manager.shuffle()
+
+  def load_image(self, image_path):
+    image = Image.open(image_path)
+    if not image.mode == "RGB":
+      image = image.convert("RGB")
+    img = np.array(image, np.uint8)
+    return img
+
+  def trim_and_resize_if_required(self, image, reso, resized_size):
+    image_height, image_width = image.shape[0:2]
+
+    if image_width != resized_size[0] or image_height != resized_size[1]:
+      # リサイズする
+      image = cv2.resize(image, resized_size, interpolation=cv2.INTER_AREA)       # INTER_AREAでやりたいのでcv2でリサイズ
+
+    image_height, image_width = image.shape[0:2]
+    if image_width > reso[0]:
+      trim_size = image_width - reso[0]
+      p = trim_size // 2 if not self.random_crop else random.randint(0, trim_size)
+      # print("w", trim_size, p)
+      image = image[:, p:p + reso[0]]
+    if image_height > reso[1]:
+      trim_size = image_height - reso[1]
+      p = trim_size // 2 if not self.random_crop else random.randint(0, trim_size)
+      # print("h", trim_size, p)
+      image = image[p:p + reso[1]]
+
+    assert image.shape[0] == reso[1] and image.shape[1] == reso[0], f"internal error, illegal trimmed size: {image.shape}, {reso}"
+    return image
+
+  def cache_latents(self, vae):
+    # TODO ここを高速化したい
+    print("caching latents.")
+    for info in tqdm(self.image_data.values()):
+      if info.latents_npz is not None:
+        info.latents = self.load_latents_from_npz(info, False)
+        info.latents = torch.FloatTensor(info.latents)
+        info.latents_flipped = self.load_latents_from_npz(info, True)             # might be None
+        if info.latents_flipped is not None:
+          info.latents_flipped = torch.FloatTensor(info.latents_flipped)
+        continue
+
+      image = self.load_image(info.absolute_path)
+      image = self.trim_and_resize_if_required(image, info.bucket_reso, info.resized_size)
+
+      img_tensor = self.image_transforms(image)
+      img_tensor = img_tensor.unsqueeze(0).to(device=vae.device, dtype=vae.dtype)
+      info.latents = vae.encode(img_tensor).latent_dist.sample().squeeze(0).to("cpu")
+
+      if self.flip_aug:
+        image = image[:, ::-1].copy()     # cannot convert to Tensor without copy
+        img_tensor = self.image_transforms(image)
+        img_tensor = img_tensor.unsqueeze(0).to(device=vae.device, dtype=vae.dtype)
+        info.latents_flipped = vae.encode(img_tensor).latent_dist.sample().squeeze(0).to("cpu")
+
+  def get_image_size(self, image_path):
+    image = Image.open(image_path)
+    return image.size
+
+  def load_image_with_face_info(self, image_path: str):
+    img = self.load_image(image_path)
+
+    face_cx = face_cy = face_w = face_h = 0
+    if self.face_crop_aug_range is not None:
+      tokens = os.path.splitext(os.path.basename(image_path))[0].split('_')
+      if len(tokens) >= 5:
+        face_cx = int(tokens[-4])
+        face_cy = int(tokens[-3])
+        face_w = int(tokens[-2])
+        face_h = int(tokens[-1])
+
+    return img, face_cx, face_cy, face_w, face_h
+
+  # いい感じに切り出す
+  def crop_target(self, image, face_cx, face_cy, face_w, face_h):
+    height, width = image.shape[0:2]
+    if height == self.height and width == self.width:
+      return image
+
+    # 画像サイズはsizeより大きいのでリサイズする
+    face_size = max(face_w, face_h)
+    min_scale = max(self.height / height, self.width / width)        # 画像がモデル入力サイズぴったりになる倍率（最小の倍率）
+    min_scale = min(1.0, max(min_scale, self.size / (face_size * self.face_crop_aug_range[1])))             # 指定した顔最小サイズ
+    max_scale = min(1.0, max(min_scale, self.size / (face_size * self.face_crop_aug_range[0])))             # 指定した顔最大サイズ
+    if min_scale >= max_scale:          # range指定がmin==max
+      scale = min_scale
+    else:
+      scale = random.uniform(min_scale, max_scale)
+
+    nh = int(height * scale + .5)
+    nw = int(width * scale + .5)
+    assert nh >= self.height and nw >= self.width, f"internal error. small scale {scale}, {width}*{height}"
+    image = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_AREA)
+    face_cx = int(face_cx * scale + .5)
+    face_cy = int(face_cy * scale + .5)
+    height, width = nh, nw
+
+    # 顔を中心として448*640とかへ切り出す
+    for axis, (target_size, length, face_p) in enumerate(zip((self.height, self.width), (height, width), (face_cy, face_cx))):
+      p1 = face_p - target_size // 2                # 顔を中心に持ってくるための切り出し位置
+
+      if self.random_crop:
+        # 背景も含めるために顔を中心に置く確率を高めつつずらす
+        range = max(length - face_p, face_p)        # 画像の端から顔中心までの距離の長いほう
+        p1 = p1 + (random.randint(0, range) + random.randint(0, range)) - range     # -range ~ +range までのいい感じの乱数
+      else:
+        # range指定があるときのみ、すこしだけランダムに（わりと適当）
+        if self.face_crop_aug_range[0] != self.face_crop_aug_range[1]:
+          if face_size > self.size // 10 and face_size >= 40:
+            p1 = p1 + random.randint(-face_size // 20, +face_size // 20)
+
+      p1 = max(0, min(p1, length - target_size))
+
+      if axis == 0:
+        image = image[p1:p1 + target_size, :]
+      else:
+        image = image[:, p1:p1 + target_size]
+
+    return image
+
+  def load_latents_from_npz(self, image_info: ImageInfo, flipped):
+    npz_file = image_info.latents_npz_flipped if flipped else image_info.latents_npz
+    if npz_file is None:
+      return None
+    return np.load(npz_file)['arr_0']
+
+  def __len__(self):
+    return self._length
+
+  def __getitem__(self, index):
+    if index == 0:
+      self.shuffle_buckets()
+
+    bucket = self.bucket_manager.buckets[self.buckets_indices[index].bucket_index]
+    bucket_batch_size = self.buckets_indices[index].bucket_batch_size
+    image_index = self.buckets_indices[index].batch_index * bucket_batch_size
+
+    loss_weights = []
+    captions = []
+    input_ids_list = []
+    latents_list = []
+    images = []
+
+    for image_key in bucket[image_index:image_index + bucket_batch_size]:
+      image_info = self.image_data[image_key]
+      loss_weights.append(self.prior_loss_weight if image_info.is_reg else 1.0)
+
+      # image/latentsを処理する
+      if image_info.latents is not None:
+        latents = image_info.latents if not self.flip_aug or random.random() < .5 else image_info.latents_flipped
+        image = None
+      elif image_info.latents_npz is not None:
+        latents = self.load_latents_from_npz(image_info, self.flip_aug and random.random() >= .5)
+        latents = torch.FloatTensor(latents)
+        image = None
+      else:
+        # 画像を読み込み、必要ならcropする
+        img, face_cx, face_cy, face_w, face_h = self.load_image_with_face_info(image_info.absolute_path)
+        im_h, im_w = img.shape[0:2]
+
+        if self.enable_bucket:
+          img = self.trim_and_resize_if_required(img, image_info.bucket_reso, image_info.resized_size)
+        else:
+          if face_cx > 0:                   # 顔位置情報あり
+            img = self.crop_target(img, face_cx, face_cy, face_w, face_h)
+          elif im_h > self.height or im_w > self.width:
+            assert self.random_crop, f"image too large, but cropping and bucketing are disabled / 画像サイズが大きいのでface_crop_aug_rangeかrandom_crop、またはbucketを有効にしてください: {image_info.absolute_path}"
+            if im_h > self.height:
+              p = random.randint(0, im_h - self.height)
+              img = img[p:p + self.height]
+            if im_w > self.width:
+              p = random.randint(0, im_w - self.width)
+              img = img[:, p:p + self.width]
+
+          im_h, im_w = img.shape[0:2]
+          assert im_h == self.height and im_w == self.width, f"image size is small / 画像サイズが小さいようです: {image_info.absolute_path}"
+
+        # augmentation
+        if self.aug is not None:
+          img = self.aug(image=img)['image']
+
+        latents = None
+        image = self.image_transforms(img)      # -1.0~1.0のtorch.Tensorになる
+
+      images.append(image)
+      latents_list.append(latents)
+
+      caption = self.process_caption(image_info.caption)
+      captions.append(caption)
+      if not self.token_padding_disabled:                     # this option might be omitted in future
+        input_ids_list.append(self.get_input_ids(caption))
+
+    example = {}
+    example['loss_weights'] = torch.FloatTensor(loss_weights)
+
+    if self.token_padding_disabled:
+      # padding=True means pad in the batch
+      example['input_ids'] = self.tokenizer(captions, padding=True, truncation=True, return_tensors="pt").input_ids
+    else:
+      # batch processing seems to be good
+      example['input_ids'] = torch.stack(input_ids_list)
+
+    if images[0] is not None:
+      images = torch.stack(images)
+      images = images.to(memory_format=torch.contiguous_format).float()
+    else:
+      images = None
+    example['images'] = images
+
+    example['latents'] = torch.stack(latents_list) if latents_list[0] is not None else None
+
+    if self.debug_dataset:
+      example['image_keys'] = bucket[image_index:image_index + self.batch_size]
+      example['captions'] = captions
+    return example
+
+
+class DreamBoothDataset(BaseDataset):
+  def __init__(self, batch_size, train_data_dir, reg_data_dir, tokenizer, max_token_length, caption_extension, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, prior_loss_weight, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset) -> None:
+    super().__init__(tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens,
+                     resolution, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset)
+
+    assert resolution is not None, f"resolution is required / resolution（解像度）指定は必須です"
+
+    self.batch_size = batch_size
+    self.size = min(self.width, self.height)                  # 短いほう
+    self.prior_loss_weight = prior_loss_weight
+    self.latents_cache = None
+
+    self.enable_bucket = enable_bucket
+    if self.enable_bucket:
+      assert min(resolution) >= min_bucket_reso, f"min_bucket_reso must be equal or less than resolution / min_bucket_resoは最小解像度より大きくできません。解像度を大きくするかmin_bucket_resoを小さくしてください"
+      assert max(resolution) <= max_bucket_reso, f"max_bucket_reso must be equal or greater than resolution / max_bucket_resoは最大解像度より小さくできません。解像度を小さくするかmin_bucket_resoを大きくしてください"
+      self.min_bucket_reso = min_bucket_reso
+      self.max_bucket_reso = max_bucket_reso
+      self.bucket_reso_steps = bucket_reso_steps
+      self.bucket_no_upscale = bucket_no_upscale
+    else:
+      self.min_bucket_reso = None
+      self.max_bucket_reso = None
+      self.bucket_reso_steps = None                              # この情報は使われない
+      self.bucket_no_upscale = False
+
+    def read_caption(img_path):
+      # captionの候補ファイル名を作る
+      base_name = os.path.splitext(img_path)[0]
+      base_name_face_det = base_name
+      tokens = base_name.split("_")
+      if len(tokens) >= 5:
+        base_name_face_det = "_".join(tokens[:-4])
+      cap_paths = [base_name + caption_extension, base_name_face_det + caption_extension]
+
+      caption = None
+      for cap_path in cap_paths:
+        if os.path.isfile(cap_path):
+          with open(cap_path, "rt", encoding='utf-8') as f:
+            try:
+              lines = f.readlines()
+            except UnicodeDecodeError as e:
+              print(f"illegal char in file (not UTF-8) / ファイルにUTF-8以外の文字があります: {cap_path}")
+              raise e
+            assert len(lines) > 0, f"caption file is empty / キャプションファイルが空です: {cap_path}"
+            caption = lines[0].strip()
+          break
+      return caption
+
+    def load_dreambooth_dir(dir):
+      if not os.path.isdir(dir):
+        # print(f"ignore file: {dir}")
+        return 0, [], []
+
+      tokens = os.path.basename(dir).split('_')
+      try:
+        n_repeats = int(tokens[0])
+      except ValueError as e:
+        print(f"ignore directory without repeats / 繰り返し回数のないディレクトリを無視します: {dir}")
+        return 0, [], []
+
+      caption_by_folder = '_'.join(tokens[1:])
+      img_paths = glob_images(dir, "*")
+      print(f"found directory {n_repeats}_{caption_by_folder} contains {len(img_paths)} image files")
+
+      # 画像ファイルごとにプロンプトを読み込み、もしあればそちらを使う
+      captions = []
+      for img_path in img_paths:
+        cap_for_img = read_caption(img_path)
+        captions.append(caption_by_folder if cap_for_img is None else cap_for_img)
+
+      self.set_tag_frequency(os.path.basename(dir), captions)         # タグ頻度を記録
+
+      return n_repeats, img_paths, captions
+
+    print("prepare train images.")
+    train_dirs = os.listdir(train_data_dir)
+    num_train_images = 0
+    for dir in train_dirs:
+      n_repeats, img_paths, captions = load_dreambooth_dir(os.path.join(train_data_dir, dir))
+      num_train_images += n_repeats * len(img_paths)
+
+      for img_path, caption in zip(img_paths, captions):
+        info = ImageInfo(img_path, n_repeats, caption, False, img_path)
+        self.register_image(info)
+
+      self.dataset_dirs_info[os.path.basename(dir)] = {"n_repeats": n_repeats, "img_count": len(img_paths)}
+
+    print(f"{num_train_images} train images with repeating.")
+    self.num_train_images = num_train_images
+
+    # reg imageは数を数えて学習画像と同じ枚数にする
+    num_reg_images = 0
+    if reg_data_dir:
+      print("prepare reg images.")
+      reg_infos: List[ImageInfo] = []
+
+      reg_dirs = os.listdir(reg_data_dir)
+      for dir in reg_dirs:
+        n_repeats, img_paths, captions = load_dreambooth_dir(os.path.join(reg_data_dir, dir))
+        num_reg_images += n_repeats * len(img_paths)
+
+        for img_path, caption in zip(img_paths, captions):
+          info = ImageInfo(img_path, n_repeats, caption, True, img_path)
+          reg_infos.append(info)
+
+        self.reg_dataset_dirs_info[os.path.basename(dir)] = {"n_repeats": n_repeats, "img_count": len(img_paths)}
+
+      print(f"{num_reg_images} reg images.")
+      if num_train_images < num_reg_images:
+        print("some of reg images are not used / 正則化画像の数が多いので、一部使用されない正則化画像があります")
+
+      if num_reg_images == 0:
+        print("no regularization images / 正則化画像が見つかりませんでした")
+      else:
+        # num_repeatsを計算する：どうせ大した数ではないのでループで処理する
+        n = 0
+        first_loop = True
+        while n < num_train_images:
+          for info in reg_infos:
+            if first_loop:
+              self.register_image(info)
+              n += info.num_repeats
+            else:
+              info.num_repeats += 1
+              n += 1
+            if n >= num_train_images:
+              break
+          first_loop = False
+
+    self.num_reg_images = num_reg_images
+
+
+class FineTuningDataset(BaseDataset):
+  def __init__(self, json_file_name, batch_size, train_data_dir, tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, flip_aug, color_aug, face_crop_aug_range, random_crop, dataset_repeats, debug_dataset) -> None:
+    super().__init__(tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens,
+                     resolution, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset)
+
+    # メタデータを読み込む
+    if os.path.exists(json_file_name):
+      print(f"loading existing metadata: {json_file_name}")
+      with open(json_file_name, "rt", encoding='utf-8') as f:
+        metadata = json.load(f)
+    else:
+      raise ValueError(f"no metadata / メタデータファイルがありません: {json_file_name}")
+
+    self.metadata = metadata
+    self.train_data_dir = train_data_dir
+    self.batch_size = batch_size
+
+    tags_list = []
+    for image_key, img_md in metadata.items():
+      # path情報を作る
+      if os.path.exists(image_key):
+        abs_path = image_key
+      else:
+        # わりといい加減だがいい方法が思いつかん
+        abs_path = glob_images(train_data_dir, image_key)
+        assert len(abs_path) >= 1, f"no image / 画像がありません: {image_key}"
+        abs_path = abs_path[0]
+
+      caption = img_md.get('caption')
+      tags = img_md.get('tags')
+      if caption is None:
+        caption = tags
+      elif tags is not None and len(tags) > 0:
+        caption = caption + ', ' + tags
+        tags_list.append(tags)
+      assert caption is not None and len(caption) > 0, f"caption or tag is required / キャプションまたはタグは必須です:{abs_path}"
+
+      image_info = ImageInfo(image_key, dataset_repeats, caption, False, abs_path)
+      image_info.image_size = img_md.get('train_resolution')
+
+      if not self.color_aug and not self.random_crop:
+        # if npz exists, use them
+        image_info.latents_npz, image_info.latents_npz_flipped = self.image_key_to_npz_file(image_key)
+
+      self.register_image(image_info)
+    self.num_train_images = len(metadata) * dataset_repeats
+    self.num_reg_images = 0
+
+    # TODO do not record tag freq when no tag
+    self.set_tag_frequency(os.path.basename(json_file_name), tags_list)
+    self.dataset_dirs_info[os.path.basename(json_file_name)] = {"n_repeats": dataset_repeats, "img_count": len(metadata)}
+
+    # check existence of all npz files
+    use_npz_latents = not (self.color_aug or self.random_crop)
+    if use_npz_latents:
+      npz_any = False
+      npz_all = True
+      for image_info in self.image_data.values():
+        has_npz = image_info.latents_npz is not None
+        npz_any = npz_any or has_npz
+
+        if self.flip_aug:
+          has_npz = has_npz and image_info.latents_npz_flipped is not None
+        npz_all = npz_all and has_npz
+
+        if npz_any and not npz_all:
+          break
+
+      if not npz_any:
+        use_npz_latents = False
+        print(f"npz file does not exist. ignore npz files / npzファイルが見つからないためnpzファイルを無視します")
+      elif not npz_all:
+        use_npz_latents = False
+        print(f"some of npz file does not exist. ignore npz files / いくつかのnpzファイルが見つからないためnpzファイルを無視します")
+        if self.flip_aug:
+          print("maybe no flipped files / 反転されたnpzファイルがないのかもしれません")
+    # else:
+    #   print("npz files are not used with color_aug and/or random_crop / color_augまたはrandom_cropが指定されているためnpzファイルは使用されません")
+
+    # check min/max bucket size
+    sizes = set()
+    resos = set()
+    for image_info in self.image_data.values():
+      if image_info.image_size is None:
+        sizes = None                  # not calculated
+        break
+      sizes.add(image_info.image_size[0])
+      sizes.add(image_info.image_size[1])
+      resos.add(tuple(image_info.image_size))
+
+    if sizes is None:
+      if use_npz_latents:
+        use_npz_latents = False
+        print(f"npz files exist, but no bucket info in metadata. ignore npz files / メタデータにbucket情報がないためnpzファイルを無視します")
+
+      assert resolution is not None, "if metadata doesn't have bucket info, resolution is required / メタデータにbucket情報がない場合はresolutionを指定してください"
+
+      self.enable_bucket = enable_bucket
+      if self.enable_bucket:
+        self.min_bucket_reso = min_bucket_reso
+        self.max_bucket_reso = max_bucket_reso
+        self.bucket_reso_steps = bucket_reso_steps
+        self.bucket_no_upscale = bucket_no_upscale
+    else:
+      if not enable_bucket:
+        print("metadata has bucket info, enable bucketing / メタデータにbucket情報があるためbucketを有効にします")
+      print("using bucket info in metadata / メタデータ内のbucket情報を使います")
+      self.enable_bucket = True
+
+      assert not bucket_no_upscale, "if metadata has bucket info, bucket reso is precalculated, so bucket_no_upscale cannot be used / メタデータ内にbucket情報がある場合はbucketの解像度は計算済みのため、bucket_no_upscaleは使えません"
+
+      # bucket情報を初期化しておく、make_bucketsで再作成しない
+      self.bucket_manager = BucketManager(False, None, None, None, None)
+      self.bucket_manager.set_predefined_resos(resos)
+
+    # npz情報をきれいにしておく
+    if not use_npz_latents:
+      for image_info in self.image_data.values():
+        image_info.latents_npz = image_info.latents_npz_flipped = None
+
+  def image_key_to_npz_file(self, image_key):
+    base_name = os.path.splitext(image_key)[0]
+    npz_file_norm = base_name + '.npz'
+
+    if os.path.exists(npz_file_norm):
+      # image_key is full path
+      npz_file_flip = base_name + '_flip.npz'
+      if not os.path.exists(npz_file_flip):
+        npz_file_flip = None
+      return npz_file_norm, npz_file_flip
+
+    # image_key is relative path
+    npz_file_norm = os.path.join(self.train_data_dir, image_key + '.npz')
+    npz_file_flip = os.path.join(self.train_data_dir, image_key + '_flip.npz')
+
+    if not os.path.exists(npz_file_norm):
+      npz_file_norm = None
+      npz_file_flip = None
+    elif not os.path.exists(npz_file_flip):
+      npz_file_flip = None
+
+    return npz_file_norm, npz_file_flip
+
+
+def debug_dataset(train_dataset, show_input_ids=False):
+  print(f"Total dataset length (steps) / データセットの長さ（ステップ数）: {len(train_dataset)}")
+  print("Escape for exit. / Escキーで中断、終了します")
+
+  train_dataset.set_current_epoch(1)
+  k = 0
+  for i, example in enumerate(train_dataset):
+    if example['latents'] is not None:
+      print(f"sample has latents from npz file: {example['latents'].size()}")
+    for j, (ik, cap, lw, iid) in enumerate(zip(example['image_keys'], example['captions'], example['loss_weights'], example['input_ids'])):
+      print(f'{ik}, size: {train_dataset.image_data[ik].image_size}, loss weight: {lw}, caption: "{cap}"')
+      if show_input_ids:
+        print(f"input ids: {iid}")
+      if example['images'] is not None:
+        im = example['images'][j]
+        print(f"image size: {im.size()}")
+        im = ((im.numpy() + 1.0) * 127.5).astype(np.uint8)
+        im = np.transpose(im, (1, 2, 0))                # c,H,W -> H,W,c
+        im = im[:, :, ::-1]                             # RGB -> BGR (OpenCV)
+        if os.name == 'nt':                             # only windows
+          cv2.imshow("img", im)
+        k = cv2.waitKey()
+        cv2.destroyAllWindows()
+        if k == 27:
+          break
+    if k == 27 or (example['images'] is None and i >= 8):
+      break
+
+
+def glob_images(directory, base="*"):
+  img_paths = []
+  for ext in IMAGE_EXTENSIONS:
+    if base == '*':
+      img_paths.extend(glob.glob(os.path.join(glob.escape(directory), base + ext)))
+    else:
+      img_paths.extend(glob.glob(glob.escape(os.path.join(directory, base + ext))))
+  # img_paths = list(set(img_paths))                    # 重複を排除
+  # img_paths.sort()
+  return img_paths
+
+
+def glob_images_pathlib(dir_path, recursive):
+  image_paths = []
+  if recursive:
+    for ext in IMAGE_EXTENSIONS:
+      image_paths += list(dir_path.rglob('*' + ext))
+  else:
+    for ext in IMAGE_EXTENSIONS:
+      image_paths += list(dir_path.glob('*' + ext))
+  # image_paths = list(set(image_paths))        # 重複を排除
+  # image_paths.sort()
+  return image_paths
+
+# endregion
+
+
+# region モジュール入れ替え部
+"""
+高速化のためのモジュール入れ替え
+"""
+
+# FlashAttentionを使うCrossAttention
+# based on https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/memory_efficient_attention_pytorch/flash_attention.py
+# LICENSE MIT https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/LICENSE
+
+# constants
+
+EPSILON = 1e-6
+
+# helper functions
+
+
+def exists(val):
+  return val is not None
+
+
+def default(val, d):
+  return val if exists(val) else d
+
+
+def model_hash(filename):
+  """Old model hash used by stable-diffusion-webui"""
+  try:
+    with open(filename, "rb") as file:
+      m = hashlib.sha256()
+
+      file.seek(0x100000)
+      m.update(file.read(0x10000))
+      return m.hexdigest()[0:8]
+  except FileNotFoundError:
+    return 'NOFILE'
+
+
+def calculate_sha256(filename):
+  """New model hash used by stable-diffusion-webui"""
+  hash_sha256 = hashlib.sha256()
+  blksize = 1024 * 1024
+
+  with open(filename, "rb") as f:
+    for chunk in iter(lambda: f.read(blksize), b""):
+      hash_sha256.update(chunk)
+
+  return hash_sha256.hexdigest()
+
+
+def precalculate_safetensors_hashes(tensors, metadata):
+  """Precalculate the model hashes needed by sd-webui-additional-networks to
+  save time on indexing the model later."""
+
+  # Because writing user metadata to the file can change the result of
+  # sd_models.model_hash(), only retain the training metadata for purposes of
+  # calculating the hash, as they are meant to be immutable
+  metadata = {k: v for k, v in metadata.items() if k.startswith("ss_")}
+
+  bytes = safetensors.torch.save(tensors, metadata)
+  b = BytesIO(bytes)
+
+  model_hash = addnet_hash_safetensors(b)
+  legacy_hash = addnet_hash_legacy(b)
+  return model_hash, legacy_hash
+
+
+def addnet_hash_legacy(b):
+  """Old model hash used by sd-webui-additional-networks for .safetensors format files"""
+  m = hashlib.sha256()
+
+  b.seek(0x100000)
+  m.update(b.read(0x10000))
+  return m.hexdigest()[0:8]
+
+
+def addnet_hash_safetensors(b):
+  """New model hash used by sd-webui-additional-networks for .safetensors format files"""
+  hash_sha256 = hashlib.sha256()
+  blksize = 1024 * 1024
+
+  b.seek(0)
+  header = b.read(8)
+  n = int.from_bytes(header, "little")
+
+  offset = n + 8
+  b.seek(offset)
+  for chunk in iter(lambda: b.read(blksize), b""):
+    hash_sha256.update(chunk)
+
+  return hash_sha256.hexdigest()
+
+
+def get_git_revision_hash() -> str:
+  try:
+    return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=os.path.dirname(__file__)).decode('ascii').strip()
+  except:
+    return "(unknown)"
+
+
+# flash attention forwards and backwards
+
+# https://arxiv.org/abs/2205.14135
+
+
+class FlashAttentionFunction(torch.autograd.function.Function):
+  @ staticmethod
+  @ torch.no_grad()
+  def forward(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size):
+    """ Algorithm 2 in the paper """
+
+    device = q.device
+    dtype = q.dtype
+    max_neg_value = -torch.finfo(q.dtype).max
+    qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)
+
+    o = torch.zeros_like(q)
+    all_row_sums = torch.zeros((*q.shape[:-1], 1), dtype=dtype, device=device)
+    all_row_maxes = torch.full((*q.shape[:-1], 1), max_neg_value, dtype=dtype, device=device)
+
+    scale = (q.shape[-1] ** -0.5)
+
+    if not exists(mask):
+      mask = (None,) * math.ceil(q.shape[-2] / q_bucket_size)
+    else:
+      mask = rearrange(mask, 'b n -> b 1 1 n')
+      mask = mask.split(q_bucket_size, dim=-1)
+
+    row_splits = zip(
+        q.split(q_bucket_size, dim=-2),
+        o.split(q_bucket_size, dim=-2),
+        mask,
+        all_row_sums.split(q_bucket_size, dim=-2),
+        all_row_maxes.split(q_bucket_size, dim=-2),
+    )
+
+    for ind, (qc, oc, row_mask, row_sums, row_maxes) in enumerate(row_splits):
+      q_start_index = ind * q_bucket_size - qk_len_diff
+
+      col_splits = zip(
+          k.split(k_bucket_size, dim=-2),
+          v.split(k_bucket_size, dim=-2),
+      )
+
+      for k_ind, (kc, vc) in enumerate(col_splits):
+        k_start_index = k_ind * k_bucket_size
+
+        attn_weights = einsum('... i d, ... j d -> ... i j', qc, kc) * scale
+
+        if exists(row_mask):
+          attn_weights.masked_fill_(~row_mask, max_neg_value)
+
+        if causal and q_start_index < (k_start_index + k_bucket_size - 1):
+          causal_mask = torch.ones((qc.shape[-2], kc.shape[-2]), dtype=torch.bool,
+                                   device=device).triu(q_start_index - k_start_index + 1)
+          attn_weights.masked_fill_(causal_mask, max_neg_value)
+
+        block_row_maxes = attn_weights.amax(dim=-1, keepdims=True)
+        attn_weights -= block_row_maxes
+        exp_weights = torch.exp(attn_weights)
+
+        if exists(row_mask):
+          exp_weights.masked_fill_(~row_mask, 0.)
+
+        block_row_sums = exp_weights.sum(dim=-1, keepdims=True).clamp(min=EPSILON)
+
+        new_row_maxes = torch.maximum(block_row_maxes, row_maxes)
+
+        exp_values = einsum('... i j, ... j d -> ... i d', exp_weights, vc)
+
+        exp_row_max_diff = torch.exp(row_maxes - new_row_maxes)
+        exp_block_row_max_diff = torch.exp(block_row_maxes - new_row_maxes)
+
+        new_row_sums = exp_row_max_diff * row_sums + exp_block_row_max_diff * block_row_sums
+
+        oc.mul_((row_sums / new_row_sums) * exp_row_max_diff).add_((exp_block_row_max_diff / new_row_sums) * exp_values)
+
+        row_maxes.copy_(new_row_maxes)
+        row_sums.copy_(new_row_sums)
+
+    ctx.args = (causal, scale, mask, q_bucket_size, k_bucket_size)
+    ctx.save_for_backward(q, k, v, o, all_row_sums, all_row_maxes)
+
+    return o
+
+  @ staticmethod
+  @ torch.no_grad()
+  def backward(ctx, do):
+    """ Algorithm 4 in the paper """
+
+    causal, scale, mask, q_bucket_size, k_bucket_size = ctx.args
+    q, k, v, o, l, m = ctx.saved_tensors
+
+    device = q.device
+
+    max_neg_value = -torch.finfo(q.dtype).max
+    qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)
+
+    dq = torch.zeros_like(q)
+    dk = torch.zeros_like(k)
+    dv = torch.zeros_like(v)
+
+    row_splits = zip(
+        q.split(q_bucket_size, dim=-2),
+        o.split(q_bucket_size, dim=-2),
+        do.split(q_bucket_size, dim=-2),
+        mask,
+        l.split(q_bucket_size, dim=-2),
+        m.split(q_bucket_size, dim=-2),
+        dq.split(q_bucket_size, dim=-2)
+    )
+
+    for ind, (qc, oc, doc, row_mask, lc, mc, dqc) in enumerate(row_splits):
+      q_start_index = ind * q_bucket_size - qk_len_diff
+
+      col_splits = zip(
+          k.split(k_bucket_size, dim=-2),
+          v.split(k_bucket_size, dim=-2),
+          dk.split(k_bucket_size, dim=-2),
+          dv.split(k_bucket_size, dim=-2),
+      )
+
+      for k_ind, (kc, vc, dkc, dvc) in enumerate(col_splits):
+        k_start_index = k_ind * k_bucket_size
+
+        attn_weights = einsum('... i d, ... j d -> ... i j', qc, kc) * scale
+
+        if causal and q_start_index < (k_start_index + k_bucket_size - 1):
+          causal_mask = torch.ones((qc.shape[-2], kc.shape[-2]), dtype=torch.bool,
+                                   device=device).triu(q_start_index - k_start_index + 1)
+          attn_weights.masked_fill_(causal_mask, max_neg_value)
+
+        exp_attn_weights = torch.exp(attn_weights - mc)
+
+        if exists(row_mask):
+          exp_attn_weights.masked_fill_(~row_mask, 0.)
+
+        p = exp_attn_weights / lc
+
+        dv_chunk = einsum('... i j, ... i d -> ... j d', p, doc)
+        dp = einsum('... i d, ... j d -> ... i j', doc, vc)
+
+        D = (doc * oc).sum(dim=-1, keepdims=True)
+        ds = p * scale * (dp - D)
+
+        dq_chunk = einsum('... i j, ... j d -> ... i d', ds, kc)
+        dk_chunk = einsum('... i j, ... i d -> ... j d', ds, qc)
+
+        dqc.add_(dq_chunk)
+        dkc.add_(dk_chunk)
+        dvc.add_(dv_chunk)
+
+    return dq, dk, dv, None, None, None, None
+
+
+def replace_unet_modules(unet: diffusers.models.unet_2d_condition.UNet2DConditionModel, mem_eff_attn, xformers):
+  if mem_eff_attn:
+    replace_unet_cross_attn_to_memory_efficient()
+  elif xformers:
+    replace_unet_cross_attn_to_xformers()
+
+
+def replace_unet_cross_attn_to_memory_efficient():
+  print("Replace CrossAttention.forward to use FlashAttention (not xformers)")
+  flash_func = FlashAttentionFunction
+
+  def forward_flash_attn(self, x, context=None, mask=None):
+    q_bucket_size = 512
+    k_bucket_size = 1024
+
+    h = self.heads
+    q = self.to_q(x)
+
+    context = context if context is not None else x
+    context = context.to(x.dtype)
+
+    if hasattr(self, 'hypernetwork') and self.hypernetwork is not None:
+      context_k, context_v = self.hypernetwork.forward(x, context)
+      context_k = context_k.to(x.dtype)
+      context_v = context_v.to(x.dtype)
+    else:
+      context_k = context
+      context_v = context
+
+    k = self.to_k(context_k)
+    v = self.to_v(context_v)
+    del context, x
+
+    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v))
+
+    out = flash_func.apply(q, k, v, mask, False, q_bucket_size, k_bucket_size)
+
+    out = rearrange(out, 'b h n d -> b n (h d)')
+
+    # diffusers 0.7.0~  わざわざ変えるなよ (;´Д｀)
+    out = self.to_out[0](out)
+    out = self.to_out[1](out)
+    return out
+
+  diffusers.models.attention.CrossAttention.forward = forward_flash_attn
+
+
+def replace_unet_cross_attn_to_xformers():
+  print("Replace CrossAttention.forward to use xformers")
+  try:
+    import xformers.ops
+  except ImportError:
+    raise ImportError("No xformers / xformersがインストールされていないようです")
+
+  def forward_xformers(self, x, context=None, mask=None):
+    h = self.heads
+    q_in = self.to_q(x)
+
+    context = default(context, x)
+    context = context.to(x.dtype)
+
+    if hasattr(self, 'hypernetwork') and self.hypernetwork is not None:
+      context_k, context_v = self.hypernetwork.forward(x, context)
+      context_k = context_k.to(x.dtype)
+      context_v = context_v.to(x.dtype)
+    else:
+      context_k = context
+      context_v = context
+
+    k_in = self.to_k(context_k)
+    v_in = self.to_v(context_v)
+
+    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b n h d', h=h), (q_in, k_in, v_in))
+    del q_in, k_in, v_in
+
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)        # 最適なのを選んでくれる
+
+    out = rearrange(out, 'b n h d -> b n (h d)', h=h)
+
+    # diffusers 0.7.0~
+    out = self.to_out[0](out)
+    out = self.to_out[1](out)
+    return out
+
+  diffusers.models.attention.CrossAttention.forward = forward_xformers
+# endregion
+
+
+# region arguments
+
+def add_sd_models_arguments(parser: argparse.ArgumentParser):
+  # for pretrained models
+  parser.add_argument("--v2", action='store_true',
+                      help='load Stable Diffusion v2.0 model / Stable Diffusion 2.0のモデルを読み込む')
+  parser.add_argument("--v_parameterization", action='store_true',
+                      help='enable v-parameterization training / v-parameterization学習を有効にする')
+  parser.add_argument("--pretrained_model_name_or_path", type=str, default=None,
+                      help="pretrained model to train, directory to Diffusers model or StableDiffusion checkpoint / 学習元モデル、Diffusers形式モデルのディレクトリまたはStableDiffusionのckptファイル")
+
+
+def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth: bool):
+  parser.add_argument("--output_dir", type=str, default=None,
+                      help="directory to output trained model / 学習後のモデル出力先ディレクトリ")
+  parser.add_argument("--output_name", type=str, default=None,
+                      help="base name of trained model file / 学習後のモデルの拡張子を除くファイル名")
+  parser.add_argument("--save_precision", type=str, default=None,
+                      choices=[None, "float", "fp16", "bf16"], help="precision in saving / 保存時に精度を変更して保存する")
+  parser.add_argument("--save_every_n_epochs", type=int, default=None,
+                      help="save checkpoint every N epochs / 学習中のモデルを指定エポックごとに保存する")
+  parser.add_argument("--save_n_epoch_ratio", type=int, default=None,
+                      help="save checkpoint N epoch ratio (for example 5 means save at least 5 files total) / 学習中のモデルを指定のエポック割合で保存する（たとえば5を指定すると最低5個のファイルが保存される）")
+  parser.add_argument("--save_last_n_epochs", type=int, default=None, help="save last N checkpoints / 最大Nエポック保存する")
+  parser.add_argument("--save_last_n_epochs_state", type=int, default=None,
+                      help="save last N checkpoints of state (overrides the value of --save_last_n_epochs)/ 最大Nエポックstateを保存する(--save_last_n_epochsの指定を上書きします)")
+  parser.add_argument("--save_state", action="store_true",
+                      help="save training state additionally (including optimizer states etc.) / optimizerなど学習状態も含めたstateを追加で保存する")
+  parser.add_argument("--resume", type=str, default=None, help="saved state to resume training / 学習再開するモデルのstate")
+
+  parser.add_argument("--train_batch_size", type=int, default=1, help="batch size for training / 学習時のバッチサイズ")
+  parser.add_argument("--max_token_length", type=int, default=None, choices=[None, 150, 225],
+                      help="max token length of text encoder (default for 75, 150 or 225) / text encoderのトークンの最大長（未指定で75、150または225が指定可）")
+  parser.add_argument("--use_8bit_adam", action="store_true",
+                      help="use 8bit Adam optimizer (requires bitsandbytes) / 8bit Adamオプティマイザを使う（bitsandbytesのインストールが必要）")
+  parser.add_argument("--use_lion_optimizer", action="store_true",
+                      help="use Lion optimizer (requires lion-pytorch) / Lionオプティマイザを使う（ lion-pytorch のインストールが必要）")
+  parser.add_argument("--mem_eff_attn", action="store_true",
+                      help="use memory efficient attention for CrossAttention / CrossAttentionに省メモリ版attentionを使う")
+  parser.add_argument("--xformers", action="store_true",
+                      help="use xformers for CrossAttention / CrossAttentionにxformersを使う")
+  parser.add_argument("--vae", type=str, default=None,
+                      help="path to checkpoint of vae to replace / VAEを入れ替える場合、VAEのcheckpointファイルまたはディレクトリ")
+
+  parser.add_argument("--learning_rate", type=float, default=2.0e-6, help="learning rate / 学習率")
+  parser.add_argument("--max_train_steps", type=int, default=1600, help="training steps / 学習ステップ数")
+  parser.add_argument("--max_train_epochs", type=int, default=None,
+                      help="training epochs (overrides max_train_steps) / 学習エポック数（max_train_stepsを上書きします）")
+  parser.add_argument("--max_data_loader_n_workers", type=int, default=8,
+                      help="max num workers for DataLoader (lower is less main RAM usage, faster epoch start and slower data loading) / DataLoaderの最大プロセス数（小さい値ではメインメモリの使用量が減りエポック間の待ち時間が減りますが、データ読み込みは遅くなります）")
+  parser.add_argument("--persistent_data_loader_workers", action="store_true",
+                      help="persistent DataLoader workers (useful for reduce time gap between epoch, but may use more memory) / DataLoader のワーカーを持続させる (エポック間の時間差を少なくするのに有効だが、より多くのメモリを消費する可能性がある)")
+  parser.add_argument("--seed", type=int, default=None, help="random seed for training / 学習時の乱数のseed")
+  parser.add_argument("--gradient_checkpointing", action="store_true",
+                      help="enable gradient checkpointing / grandient checkpointingを有効にする")
+  parser.add_argument("--gradient_accumulation_steps", type=int, default=1,
+                      help="Number of updates steps to accumulate before performing a backward/update pass / 学習時に逆伝播をする前に勾配を合計するステップ数")
+  parser.add_argument("--mixed_precision", type=str, default="no",
+                      choices=["no", "fp16", "bf16"], help="use mixed precision / 混合精度を使う場合、その精度")
+  parser.add_argument("--full_fp16", action="store_true", help="fp16 training including gradients / 勾配も含めてfp16で学習する")
+  parser.add_argument("--clip_skip", type=int, default=None,
+                      help="use output of nth layer from back of text encoder (n>=1) / text encoderの後ろからn番目の層の出力を用いる（nは1以上）")
+  parser.add_argument("--logging_dir", type=str, default=None,
+                      help="enable logging and output TensorBoard log to this directory / ログ出力を有効にしてこのディレクトリにTensorBoard用のログを出力する")
+  parser.add_argument("--log_prefix", type=str, default=None, help="add prefix for each log directory / ログディレクトリ名の先頭に追加する文字列")
+  parser.add_argument("--lr_scheduler", type=str, default="constant",
+                      help="scheduler to use for learning rate / 学習率のスケジューラ: linear, cosine, cosine_with_restarts, polynomial, constant (default), constant_with_warmup")
+  parser.add_argument("--lr_warmup_steps", type=int, default=0,
+                      help="Number of steps for the warmup in the lr scheduler (default is 0) / 学習率のスケジューラをウォームアップするステップ数（デフォルト0）")
+  parser.add_argument("--noise_offset", type=float, default=None,
+                      help="enable noise offset with this value (if enabled, around 0.1 is recommended) / Noise offsetを有効にしてこの値を設定する（有効にする場合は0.1程度を推奨）")
+  parser.add_argument("--lowram", action="store_true",
+                      help="enable low RAM optimization. e.g. load models to VRAM instead of RAM (for machines which have bigger VRAM than RAM such as Colab and Kaggle) / メインメモリが少ない環境向け最適化を有効にする。たとえばVRAMにモデルを読み込むなど（ColabやKaggleなどRAMに比べてVRAMが多い環境向け）")
+
+  if support_dreambooth:
+    # DreamBooth training
+    parser.add_argument("--prior_loss_weight", type=float, default=1.0,
+                        help="loss weight for regularization images / 正則化画像のlossの重み")
+
+
+def verify_training_args(args: argparse.Namespace):
+  if args.v_parameterization and not args.v2:
+    print("v_parameterization should be with v2 / v1でv_parameterizationを使用することは想定されていません")
+  if args.v2 and args.clip_skip is not None:
+    print("v2 with clip_skip will be unexpected / v2でclip_skipを使用することは想定されていません")
+
+
+def add_dataset_arguments(parser: argparse.ArgumentParser, support_dreambooth: bool, support_caption: bool, support_caption_dropout: bool):
+  # dataset common
+  parser.add_argument("--train_data_dir", type=str, default=None, help="directory for train images / 学習画像データのディレクトリ")
+  parser.add_argument("--shuffle_caption", action="store_true",
+                      help="shuffle comma-separated caption / コンマで区切られたcaptionの各要素をshuffleする")
+  parser.add_argument("--caption_extension", type=str, default=".caption", help="extension of caption files / 読み込むcaptionファイルの拡張子")
+  parser.add_argument("--caption_extention", type=str, default=None,
+                      help="extension of caption files (backward compatibility) / 読み込むcaptionファイルの拡張子（スペルミスを残してあります）")
+  parser.add_argument("--keep_tokens", type=int, default=None,
+                      help="keep heading N tokens when shuffling caption tokens / captionのシャッフル時に、先頭からこの個数のトークンをシャッフルしないで残す")
+  parser.add_argument("--color_aug", action="store_true", help="enable weak color augmentation / 学習時に色合いのaugmentationを有効にする")
+  parser.add_argument("--flip_aug", action="store_true", help="enable horizontal flip augmentation / 学習時に左右反転のaugmentationを有効にする")
+  parser.add_argument("--face_crop_aug_range", type=str, default=None,
+                      help="enable face-centered crop augmentation and its range (e.g. 2.0,4.0) / 学習時に顔を中心とした切り出しaugmentationを有効にするときは倍率を指定する（例：2.0,4.0）")
+  parser.add_argument("--random_crop", action="store_true",
+                      help="enable random crop (for style training in face-centered crop augmentation) / ランダムな切り出しを有効にする（顔を中心としたaugmentationを行うときに画風の学習用に指定する）")
+  parser.add_argument("--debug_dataset", action="store_true",
+                      help="show images for debugging (do not train) / デバッグ用に学習データを画面表示する（学習は行わない）")
+  parser.add_argument("--resolution", type=str, default=None,
+                      help="resolution in training ('size' or 'width,height') / 学習時の画像解像度（'サイズ'指定、または'幅,高さ'指定）")
+  parser.add_argument("--cache_latents", action="store_true",
+                      help="cache latents to reduce memory (augmentations must be disabled) / メモリ削減のためにlatentをcacheする（augmentationは使用不可）")
+  parser.add_argument("--enable_bucket", action="store_true",
+                      help="enable buckets for multi aspect ratio training / 複数解像度学習のためのbucketを有効にする")
+  parser.add_argument("--min_bucket_reso", type=int, default=256, help="minimum resolution for buckets / bucketの最小解像度")
+  parser.add_argument("--max_bucket_reso", type=int, default=1024, help="maximum resolution for buckets / bucketの最大解像度")
+  parser.add_argument("--bucket_reso_steps", type=int, default=64,
+                      help="steps of resolution for buckets, divisible by 8 is recommended / bucketの解像度の単位、8で割り切れる値を推奨します")
+  parser.add_argument("--bucket_no_upscale", action="store_true",
+                      help="make bucket for each image without upscaling / 画像を拡大せずbucketを作成します")
+
+  if support_caption_dropout:
+    # Textual Inversion はcaptionのdropoutをsupportしない
+    # いわゆるtensorのDropoutと紛らわしいのでprefixにcaptionを付けておく　every_n_epochsは他と平仄を合わせてdefault Noneに
+    parser.add_argument("--caption_dropout_rate", type=float, default=0,
+                        help="Rate out dropout caption(0.0~1.0) / captionをdropoutする割合")
+    parser.add_argument("--caption_dropout_every_n_epochs", type=int, default=None,
+                        help="Dropout all captions every N epochs / captionを指定エポックごとにdropoutする")
+    parser.add_argument("--caption_tag_dropout_rate", type=float, default=0,
+                        help="Rate out dropout comma separated tokens(0.0~1.0) / カンマ区切りのタグをdropoutする割合")
+
+  if support_dreambooth:
+    # DreamBooth dataset
+    parser.add_argument("--reg_data_dir", type=str, default=None, help="directory for regularization images / 正則化画像データのディレクトリ")
+
+  if support_caption:
+    # caption dataset
+    parser.add_argument("--in_json", type=str, default=None, help="json metadata for dataset / データセットのmetadataのjsonファイル")
+    parser.add_argument("--dataset_repeats", type=int, default=1,
+                        help="repeat dataset when training with captions / キャプションでの学習時にデータセットを繰り返す回数")
+
+
+def add_sd_saving_arguments(parser: argparse.ArgumentParser):
+  parser.add_argument("--save_model_as", type=str, default=None, choices=[None, "ckpt", "safetensors", "diffusers", "diffusers_safetensors"],
+                      help="format to save the model (default is same to original) / モデル保存時の形式（未指定時は元モデルと同じ）")
+  parser.add_argument("--use_safetensors", action='store_true',
+                      help="use safetensors format to save (if save_model_as is not specified) / checkpoint、モデルをsafetensors形式で保存する（save_model_as未指定時）")
+
+# endregion
+
+# region utils
+
+
+def prepare_dataset_args(args: argparse.Namespace, support_metadata: bool):
+  # backward compatibility
+  if args.caption_extention is not None:
+    args.caption_extension = args.caption_extention
+    args.caption_extention = None
+
+  if args.cache_latents:
+    assert not args.color_aug, "when caching latents, color_aug cannot be used / latentをキャッシュするときはcolor_augは使えません"
+    assert not args.random_crop, "when caching latents, random_crop cannot be used / latentをキャッシュするときはrandom_cropは使えません"
+
+  # assert args.resolution is not None, f"resolution is required / resolution（解像度）を指定してください"
+  if args.resolution is not None:
+    args.resolution = tuple([int(r) for r in args.resolution.split(',')])
+    if len(args.resolution) == 1:
+      args.resolution = (args.resolution[0], args.resolution[0])
+    assert len(args.resolution) == 2, \
+        f"resolution must be 'size' or 'width,height' / resolution（解像度）は'サイズ'または'幅','高さ'で指定してください: {args.resolution}"
+
+  if args.face_crop_aug_range is not None:
+    args.face_crop_aug_range = tuple([float(r) for r in args.face_crop_aug_range.split(',')])
+    assert len(args.face_crop_aug_range) == 2 and args.face_crop_aug_range[0] <= args.face_crop_aug_range[1], \
+        f"face_crop_aug_range must be two floats / face_crop_aug_rangeは'下限,上限'で指定してください: {args.face_crop_aug_range}"
+  else:
+    args.face_crop_aug_range = None
+
+  if support_metadata:
+    if args.in_json is not None and (args.color_aug or args.random_crop):
+      print(f"latents in npz is ignored when color_aug or random_crop is True / color_augまたはrandom_cropを有効にした場合、npzファイルのlatentsは無視されます")
+
+
+def load_tokenizer(args: argparse.Namespace):
+  print("prepare tokenizer")
+  if args.v2:
+    tokenizer = CLIPTokenizer.from_pretrained(V2_STABLE_DIFFUSION_PATH, subfolder="tokenizer")
+  else:
+    tokenizer = CLIPTokenizer.from_pretrained(TOKENIZER_PATH)
+  if args.max_token_length is not None:
+    print(f"update token length: {args.max_token_length}")
+  return tokenizer
+
+
+def prepare_accelerator(args: argparse.Namespace):
+  if args.logging_dir is None:
+    log_with = None
+    logging_dir = None
+  else:
+    log_with = "tensorboard"
+    log_prefix = "" if args.log_prefix is None else args.log_prefix
+    logging_dir = args.logging_dir + "/" + log_prefix + time.strftime('%Y%m%d%H%M%S', time.localtime())
+
+  accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.mixed_precision,
+                            log_with=log_with, logging_dir=logging_dir)
+
+  # accelerateの互換性問題を解決する
+  accelerator_0_15 = True
+  try:
+    accelerator.unwrap_model("dummy", True)
+    print("Using accelerator 0.15.0 or above.")
+  except TypeError:
+    accelerator_0_15 = False
+
+  def unwrap_model(model):
+    if accelerator_0_15:
+      return accelerator.unwrap_model(model, True)
+    return accelerator.unwrap_model(model)
+
+  return accelerator, unwrap_model
+
+
+def prepare_dtype(args: argparse.Namespace):
+  weight_dtype = torch.float32
+  if args.mixed_precision == "fp16":
+    weight_dtype = torch.float16
+  elif args.mixed_precision == "bf16":
+    weight_dtype = torch.bfloat16
+
+  save_dtype = None
+  if args.save_precision == "fp16":
+    save_dtype = torch.float16
+  elif args.save_precision == "bf16":
+    save_dtype = torch.bfloat16
+  elif args.save_precision == "float":
+    save_dtype = torch.float32
+
+  return weight_dtype, save_dtype
+
+
+def load_target_model(args: argparse.Namespace, weight_dtype):
+  load_stable_diffusion_format = os.path.isfile(args.pretrained_model_name_or_path)           # determine SD or Diffusers
+  if load_stable_diffusion_format:
+    print("load StableDiffusion checkpoint")
+    text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.pretrained_model_name_or_path)
+  else:
+    print("load Diffusers pretrained models")
+    pipe = StableDiffusionPipeline.from_pretrained(args.pretrained_model_name_or_path, tokenizer=None, safety_checker=None)
+    text_encoder = pipe.text_encoder
+    vae = pipe.vae
+    unet = pipe.unet
+    del pipe
+
+  # VAEを読み込む
+  if args.vae is not None:
+    vae = model_util.load_vae(args.vae, weight_dtype)
+    print("additional VAE loaded")
+
+  return text_encoder, vae, unet, load_stable_diffusion_format
+
+
+def patch_accelerator_for_fp16_training(accelerator):
+  org_unscale_grads = accelerator.scaler._unscale_grads_
+
+  def _unscale_grads_replacer(optimizer, inv_scale, found_inf, allow_fp16):
+    return org_unscale_grads(optimizer, inv_scale, found_inf, True)
+
+  accelerator.scaler._unscale_grads_ = _unscale_grads_replacer
+
+
+def get_hidden_states(args: argparse.Namespace, input_ids, tokenizer, text_encoder, weight_dtype=None):
+  # with no_token_padding, the length is not max length, return result immediately
+  if input_ids.size()[-1] != tokenizer.model_max_length:
+    return text_encoder(input_ids)[0]
+
+  b_size = input_ids.size()[0]
+  input_ids = input_ids.reshape((-1, tokenizer.model_max_length))     # batch_size*3, 77
+
+  if args.clip_skip is None:
+    encoder_hidden_states = text_encoder(input_ids)[0]
+  else:
+    enc_out = text_encoder(input_ids, output_hidden_states=True, return_dict=True)
+    encoder_hidden_states = enc_out['hidden_states'][-args.clip_skip]
+    encoder_hidden_states = text_encoder.text_model.final_layer_norm(encoder_hidden_states)
+
+  # bs*3, 77, 768 or 1024
+  encoder_hidden_states = encoder_hidden_states.reshape((b_size, -1, encoder_hidden_states.shape[-1]))
+
+  if args.max_token_length is not None:
+    if args.v2:
+        # v2: <BOS>...<EOS> <PAD> ... の三連を <BOS>...<EOS> <PAD> ... へ戻す　正直この実装でいいのかわからん
+      states_list = [encoder_hidden_states[:, 0].unsqueeze(1)]                              # <BOS>
+      for i in range(1, args.max_token_length, tokenizer.model_max_length):
+        chunk = encoder_hidden_states[:, i:i + tokenizer.model_max_length - 2]              # <BOS> の後から 最後の前まで
+        if i > 0:
+          for j in range(len(chunk)):
+            if input_ids[j, 1] == tokenizer.eos_token:                                      # 空、つまり <BOS> <EOS> <PAD> ...のパターン
+              chunk[j, 0] = chunk[j, 1]                                                     # 次の <PAD> の値をコピーする
+        states_list.append(chunk)  # <BOS> の後から <EOS> の前まで
+      states_list.append(encoder_hidden_states[:, -1].unsqueeze(1))                         # <EOS> か <PAD> のどちらか
+      encoder_hidden_states = torch.cat(states_list, dim=1)
+    else:
+      # v1: <BOS>...<EOS> の三連を <BOS>...<EOS> へ戻す
+      states_list = [encoder_hidden_states[:, 0].unsqueeze(1)]                              # <BOS>
+      for i in range(1, args.max_token_length, tokenizer.model_max_length):
+        states_list.append(encoder_hidden_states[:, i:i + tokenizer.model_max_length - 2])  # <BOS> の後から <EOS> の前まで
+      states_list.append(encoder_hidden_states[:, -1].unsqueeze(1))                         # <EOS>
+      encoder_hidden_states = torch.cat(states_list, dim=1)
+
+  if weight_dtype is not None:
+    # this is required for additional network training
+    encoder_hidden_states = encoder_hidden_states.to(weight_dtype)
+
+  return encoder_hidden_states
+
+
+def get_epoch_ckpt_name(args: argparse.Namespace, use_safetensors, epoch):
+  model_name = DEFAULT_EPOCH_NAME if args.output_name is None else args.output_name
+  ckpt_name = EPOCH_FILE_NAME.format(model_name, epoch) + (".safetensors" if use_safetensors else ".ckpt")
+  return model_name, ckpt_name
+
+
+def save_on_epoch_end(args: argparse.Namespace, save_func, remove_old_func, epoch_no: int, num_train_epochs: int):
+  saving = epoch_no % args.save_every_n_epochs == 0 and epoch_no < num_train_epochs
+  if saving:
+    os.makedirs(args.output_dir, exist_ok=True)
+    save_func()
+
+    if args.save_last_n_epochs is not None:
+      remove_epoch_no = epoch_no - args.save_every_n_epochs * args.save_last_n_epochs
+      remove_old_func(remove_epoch_no)
+  return saving
+
+
+def save_sd_model_on_epoch_end(args: argparse.Namespace, accelerator, src_path: str, save_stable_diffusion_format: bool, use_safetensors: bool, save_dtype: torch.dtype, epoch: int, num_train_epochs: int, global_step: int, text_encoder, unet, vae):
+  epoch_no = epoch + 1
+  model_name, ckpt_name = get_epoch_ckpt_name(args, use_safetensors, epoch_no)
+
+  if save_stable_diffusion_format:
+    def save_sd():
+      ckpt_file = os.path.join(args.output_dir, ckpt_name)
+      print(f"saving checkpoint: {ckpt_file}")
+      model_util.save_stable_diffusion_checkpoint(args.v2, ckpt_file, text_encoder, unet,
+                                                  src_path, epoch_no, global_step, save_dtype, vae)
+
+    def remove_sd(old_epoch_no):
+      _, old_ckpt_name = get_epoch_ckpt_name(args,  use_safetensors, old_epoch_no)
+      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}")
+        os.remove(old_ckpt_file)
+
+    save_func = save_sd
+    remove_old_func = remove_sd
+  else:
+    def save_du():
+      out_dir = os.path.join(args.output_dir, EPOCH_DIFFUSERS_DIR_NAME.format(model_name, epoch_no))
+      print(f"saving model: {out_dir}")
+      os.makedirs(out_dir, exist_ok=True)
+      model_util.save_diffusers_checkpoint(args.v2, out_dir, text_encoder, unet,
+                                           src_path, vae=vae, use_safetensors=use_safetensors)
+
+    def remove_du(old_epoch_no):
+      out_dir_old = os.path.join(args.output_dir, EPOCH_DIFFUSERS_DIR_NAME.format(model_name, old_epoch_no))
+      if os.path.exists(out_dir_old):
+        print(f"removing old model: {out_dir_old}")
+        shutil.rmtree(out_dir_old)
+
+    save_func = save_du
+    remove_old_func = remove_du
+
+  saving = save_on_epoch_end(args, save_func, remove_old_func, epoch_no, num_train_epochs)
+  if saving and args.save_state:
+    save_state_on_epoch_end(args, accelerator, model_name, epoch_no)
+
+
+def save_state_on_epoch_end(args: argparse.Namespace, accelerator, model_name, epoch_no):
+  print("saving state.")
+  accelerator.save_state(os.path.join(args.output_dir, EPOCH_STATE_NAME.format(model_name, epoch_no)))
+
+  last_n_epochs = args.save_last_n_epochs_state if args.save_last_n_epochs_state else args.save_last_n_epochs
+  if last_n_epochs is not None:
+    remove_epoch_no = epoch_no - args.save_every_n_epochs * last_n_epochs
+    state_dir_old = os.path.join(args.output_dir, EPOCH_STATE_NAME.format(model_name, remove_epoch_no))
+    if os.path.exists(state_dir_old):
+      print(f"removing old state: {state_dir_old}")
+      shutil.rmtree(state_dir_old)
+
+
+def save_sd_model_on_train_end(args: argparse.Namespace, src_path: str, save_stable_diffusion_format: bool, use_safetensors: bool, save_dtype: torch.dtype, epoch: int, global_step: int, text_encoder, unet, vae):
+  model_name = DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name
+
+  if save_stable_diffusion_format:
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    ckpt_name = model_name + (".safetensors" if use_safetensors else ".ckpt")
+    ckpt_file = os.path.join(args.output_dir, ckpt_name)
+
+    print(f"save trained model as StableDiffusion checkpoint to {ckpt_file}")
+    model_util.save_stable_diffusion_checkpoint(args.v2, ckpt_file, text_encoder, unet,
+                                                src_path, epoch, global_step, save_dtype, vae)
+  else:
+    out_dir = os.path.join(args.output_dir, model_name)
+    os.makedirs(out_dir, exist_ok=True)
+
+    print(f"save trained model as Diffusers to {out_dir}")
+    model_util.save_diffusers_checkpoint(args.v2, out_dir, text_encoder, unet,
+                                         src_path, vae=vae, use_safetensors=use_safetensors)
+
+
+def save_state_on_train_end(args: argparse.Namespace, accelerator):
+  print("saving last state.")
+  os.makedirs(args.output_dir, exist_ok=True)
+  model_name = DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name
+  accelerator.save_state(os.path.join(args.output_dir, LAST_STATE_NAME.format(model_name)))
+
+# endregion
+
+# region 前処理用
+
+
+class ImageLoadingDataset(torch.utils.data.Dataset):
+  def __init__(self, image_paths):
+    self.images = image_paths
+
+  def __len__(self):
+    return len(self.images)
+
+  def __getitem__(self, idx):
+    img_path = self.images[idx]
+
+    try:
+      image = Image.open(img_path).convert("RGB")
+      # convert to tensor temporarily so dataloader will accept it
+      tensor_pil = transforms.functional.pil_to_tensor(image)
+    except Exception as e:
+      print(f"Could not load image path / 画像を読み込めません: {img_path}, error: {e}")
+      return None
+
+    return (tensor_pil, img_path)
+
+
+# endregion

fine_tune.py

@@ -13,11 +13,7 @@ import diffusers
 from diffusers import DDPMScheduler
 
 import library.train_util as train_util
-import library.config_util as config_util
-from library.config_util import (
-  ConfigSanitizer,
-  BlueprintGenerator,
-)
+
 
 def collate_fn(examples):
   return examples[0]
@@ -34,36 +30,25 @@ def train(args):
 
   tokenizer = train_util.load_tokenizer(args)
 
-  blueprint_generator = BlueprintGenerator(ConfigSanitizer(False, True, True))
-  if args.dataset_config is not None:
-    print(f"Load dataset config from {args.dataset_config}")
-    user_config = config_util.load_user_config(args.dataset_config)
-    ignored = ["train_data_dir", "in_json"]
-    if any(getattr(args, attr) is not None for attr in ignored):
-      print("ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(', '.join(ignored)))
-  else:
-    user_config = {
-      "datasets": [{
-        "subsets": [{
-          "image_dir": args.train_data_dir,
-          "metadata_file": args.in_json,
-        }]
-      }]
-    }
-
-  blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
-  train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
+  train_dataset = train_util.FineTuningDataset(args.in_json, args.train_batch_size, args.train_data_dir,
+                                               tokenizer, args.max_token_length, args.shuffle_caption, args.keep_tokens,
+                                               args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                               args.bucket_reso_steps, args.bucket_no_upscale,
+                                               args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop,
+                                               args.dataset_repeats, args.debug_dataset)
+
+  # 学習データのdropout率を設定する
+  train_dataset.set_caption_dropout(args.caption_dropout_rate, args.caption_dropout_every_n_epochs, args.caption_tag_dropout_rate)
+
+  train_dataset.make_buckets()
 
   if args.debug_dataset:
-    train_util.debug_dataset(train_dataset_group)
+    train_util.debug_dataset(train_dataset)
     return
-  if len(train_dataset_group) == 0:
+  if len(train_dataset) == 0:
     print("No data found. Please verify the metadata file and train_data_dir option. / 画像がありません。メタデータおよびtrain_data_dirオプションを確認してください。")
     return
 
-  if cache_latents:
-    assert train_dataset_group.is_latent_cacheable(), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
-
   # acceleratorを準備する
   print("prepare accelerator")
   accelerator, unwrap_model = train_util.prepare_accelerator(args)
@@ -124,7 +109,7 @@ def train(args):
     vae.requires_grad_(False)
     vae.eval()
     with torch.no_grad():
-      train_dataset_group.cache_latents(vae)
+      train_dataset.cache_latents(vae)
     vae.to("cpu")
     if torch.cuda.is_available():
       torch.cuda.empty_cache()
@@ -164,13 +149,33 @@ def train(args):
 
   # 学習に必要なクラスを準備する
   print("prepare optimizer, data loader etc.")
-  _, _, optimizer = train_util.get_optimizer(args, trainable_params=params_to_optimize)
+
+  # 8-bit Adamを使う
+  if args.use_8bit_adam:
+    try:
+      import bitsandbytes as bnb
+    except ImportError:
+      raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
+    print("use 8-bit Adam optimizer")
+    optimizer_class = bnb.optim.AdamW8bit
+  elif args.use_lion_optimizer:
+    try:
+      import lion_pytorch
+    except ImportError:
+      raise ImportError("No lion_pytorch / lion_pytorch がインストールされていないようです")
+    print("use Lion optimizer")
+    optimizer_class = lion_pytorch.Lion
+  else:
+    optimizer_class = torch.optim.AdamW
+
+  # betaやweight decayはdiffusers DreamBoothもDreamBooth SDもデフォルト値のようなのでオプションはとりあえず省略
+  optimizer = optimizer_class(params_to_optimize, lr=args.learning_rate)
 
   # dataloaderを準備する
   # 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, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
 
   # 学習ステップ数を計算する
   if args.max_train_epochs is not None:
@@ -178,9 +183,8 @@ 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,
-                                              num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
-                                              num_cycles=args.lr_scheduler_num_cycles, power=args.lr_scheduler_power)
+  lr_scheduler = diffusers.optimization.get_scheduler(
+      args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps, num_training_steps=args.max_train_steps * args.gradient_accumulation_steps)
 
   # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
   if args.full_fp16:
@@ -214,7 +218,7 @@ def train(args):
   # 学習する
   total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
   print("running training / 学習開始")
-  print(f"  num examples / サンプル数: {train_dataset_group.num_train_images}")
+  print(f"  num examples / サンプル数: {train_dataset.num_train_images}")
   print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
   print(f"  num epochs / epoch数: {num_train_epochs}")
   print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
@@ -233,7 +237,7 @@ def train(args):
 
   for epoch in range(num_train_epochs):
     print(f"epoch {epoch+1}/{num_train_epochs}")
-    train_dataset_group.set_current_epoch(epoch + 1)
+    train_dataset.set_current_epoch(epoch + 1)
 
     for m in training_models:
       m.train()
@@ -282,11 +286,11 @@ def train(args):
         loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="mean")
 
         accelerator.backward(loss)
-        if accelerator.sync_gradients and args.max_grad_norm != 0.0:
+        if accelerator.sync_gradients:
           params_to_clip = []
           for m in training_models:
             params_to_clip.extend(m.parameters())
-          accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+          accelerator.clip_grad_norm_(params_to_clip, 1.0)  # args.max_grad_norm)
 
         optimizer.step()
         lr_scheduler.step()
@@ -297,16 +301,11 @@ def train(args):
         progress_bar.update(1)
         global_step += 1
 
-        train_util.sample_images(accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
-
       current_loss = loss.detach().item()        # 平均なのでbatch sizeは関係ないはず
       if args.logging_dir is not None:
-        logs = {"loss": current_loss, "lr": float(lr_scheduler.get_last_lr()[0])}
-        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 = {"loss": current_loss, "lr": lr_scheduler.get_last_lr()[0]}
         accelerator.log(logs, step=global_step)
 
-      # TODO moving averageにする
       loss_total += current_loss
       avr_loss = loss_total / (step+1)
       logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
@@ -316,7 +315,7 @@ def train(args):
         break
 
     if args.logging_dir is not None:
-      logs = {"loss/epoch": loss_total / len(train_dataloader)}
+      logs = {"epoch_loss": loss_total / len(train_dataloader)}
       accelerator.log(logs, step=epoch+1)
 
     accelerator.wait_for_everyone()
@@ -326,8 +325,6 @@ def train(args):
       train_util.save_sd_model_on_epoch_end(args, accelerator, src_path, save_stable_diffusion_format, use_safetensors,
                                             save_dtype, epoch, num_train_epochs, global_step,  unwrap_model(text_encoder), unwrap_model(unet), vae)
 
-    train_util.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
-    
   is_main_process = accelerator.is_main_process
   if is_main_process:
     unet = unwrap_model(unet)
@@ -354,8 +351,6 @@ if __name__ == '__main__':
   train_util.add_dataset_arguments(parser, False, True, True)
   train_util.add_training_arguments(parser, False)
   train_util.add_sd_saving_arguments(parser)
-  train_util.add_optimizer_arguments(parser)
-  config_util.add_config_arguments(parser)
 
   parser.add_argument("--diffusers_xformers", action='store_true',
                       help='use xformers by diffusers / Diffusersでxformersを使用する')

gen_img_diffusers.py

@@ -47,7 +47,7 @@ VGG(
 """
 
 import json
-from typing import Any, List, NamedTuple, Optional, Tuple, Union, Callable
+from typing import List, Optional, Union
 import glob
 import importlib
 import inspect
@@ -60,6 +60,7 @@ import math
 import os
 import random
 import re
+from typing import Any, Callable, List, Optional, Union
 
 import diffusers
 import numpy as np
@@ -80,9 +81,6 @@ from PIL import Image
 from PIL.PngImagePlugin import PngInfo
 
 import library.model_util as model_util
-import library.train_util as train_util
-import tools.original_control_net as original_control_net
-from tools.original_control_net import ControlNetInfo
 
 # Tokenizer: checkpointから読み込むのではなくあらかじめ提供されているものを使う
 TOKENIZER_PATH = "openai/clip-vit-large-patch14"
@@ -489,9 +487,6 @@ class PipelineLike():
       self.vgg16_feat_model = torchvision.models._utils.IntermediateLayerGetter(vgg16_model.features, return_layers=return_layers)
       self.vgg16_normalize = transforms.Normalize(mean=VGG16_IMAGE_MEAN, std=VGG16_IMAGE_STD)
 
-    # ControlNet
-    self.control_nets: List[ControlNetInfo] = []
-
   # Textual Inversion
   def add_token_replacement(self, target_token_id, rep_token_ids):
     self.token_replacements[target_token_id] = rep_token_ids
@@ -505,11 +500,7 @@ class PipelineLike():
         new_tokens.append(token)
     return new_tokens
 
-  def set_control_nets(self, ctrl_nets):
-    self.control_nets = ctrl_nets
-
   # region xformersとか使う部分：独自に書き換えるので関係なし
-
   def enable_xformers_memory_efficient_attention(self):
     r"""
     Enable memory efficient attention as implemented in xformers.
@@ -590,8 +581,6 @@ class PipelineLike():
       latents: Optional[torch.FloatTensor] = None,
       max_embeddings_multiples: Optional[int] = 3,
       output_type: Optional[str] = "pil",
-      vae_batch_size: float = None,
-      return_latents: bool = False,
       # return_dict: bool = True,
       callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
       is_cancelled_callback: Optional[Callable[[], bool]] = None,
@@ -683,9 +672,6 @@ class PipelineLike():
     else:
       raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
 
-    vae_batch_size = batch_size if vae_batch_size is None else (
-        int(vae_batch_size) if vae_batch_size >= 1 else max(1, int(batch_size * vae_batch_size)))
-
     if strength < 0 or strength > 1:
       raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
 
@@ -766,7 +752,7 @@ class PipelineLike():
       text_embeddings_clip = self.clip_model.get_text_features(clip_text_input)
       text_embeddings_clip = text_embeddings_clip / text_embeddings_clip.norm(p=2, dim=-1, keepdim=True)      # prompt複数件でもOK
 
-    if self.clip_image_guidance_scale > 0 or self.vgg16_guidance_scale > 0 and clip_guide_images is not None or self.control_nets:
+    if self.clip_image_guidance_scale > 0 or self.vgg16_guidance_scale > 0 and clip_guide_images is not None:
       if isinstance(clip_guide_images, PIL.Image.Image):
         clip_guide_images = [clip_guide_images]
 
@@ -779,7 +765,7 @@ class PipelineLike():
         image_embeddings_clip = image_embeddings_clip / image_embeddings_clip.norm(p=2, dim=-1, keepdim=True)
         if len(image_embeddings_clip) == 1:
           image_embeddings_clip = image_embeddings_clip.repeat((batch_size, 1, 1, 1))
-      elif self.vgg16_guidance_scale > 0:
+      else:
         size = (width // VGG16_INPUT_RESIZE_DIV, height // VGG16_INPUT_RESIZE_DIV)            # とりあえず1/4に（小さいか?）
         clip_guide_images = [preprocess_vgg16_guide_image(im, size) for im in clip_guide_images]
         clip_guide_images = torch.cat(clip_guide_images, dim=0)
@@ -788,10 +774,6 @@ class PipelineLike():
         image_embeddings_vgg16 = self.vgg16_feat_model(clip_guide_images)['feat']
         if len(image_embeddings_vgg16) == 1:
           image_embeddings_vgg16 = image_embeddings_vgg16.repeat((batch_size, 1, 1, 1))
-      else:
-        # ControlNetのhintにguide imageを流用する
-        # 前処理はControlNet側で行う
-        pass
 
     # set timesteps
     self.scheduler.set_timesteps(num_inference_steps, self.device)
@@ -799,6 +781,7 @@ class PipelineLike():
     latents_dtype = text_embeddings.dtype
     init_latents_orig = None
     mask = None
+    noise = None
 
     if init_image is None:
       # get the initial random noise unless the user supplied it
@@ -830,8 +813,6 @@ class PipelineLike():
       if isinstance(init_image[0], PIL.Image.Image):
         init_image = [preprocess_image(im) for im in init_image]
         init_image = torch.cat(init_image)
-      if isinstance(init_image, list):
-        init_image = torch.stack(init_image)
 
       # mask image to tensor
       if mask_image is not None:
@@ -842,24 +823,9 @@ class PipelineLike():
 
       # encode the init image into latents and scale the latents
       init_image = init_image.to(device=self.device, dtype=latents_dtype)
-      if init_image.size()[2:] == (height // 8, width // 8):
-        init_latents = init_image
-      else:
-        if vae_batch_size >= batch_size:
-          init_latent_dist = self.vae.encode(init_image).latent_dist
-          init_latents = init_latent_dist.sample(generator=generator)
-        else:
-          if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-          init_latents = []
-          for i in tqdm(range(0, batch_size, vae_batch_size)):
-            init_latent_dist = self.vae.encode(init_image[i:i + vae_batch_size]
-                                               if vae_batch_size > 1 else init_image[i].unsqueeze(0)).latent_dist
-            init_latents.append(init_latent_dist.sample(generator=generator))
-          init_latents = torch.cat(init_latents)
-
-        init_latents = 0.18215 * init_latents
-
+      init_latent_dist = self.vae.encode(init_image).latent_dist
+      init_latents = init_latent_dist.sample(generator=generator)
+      init_latents = 0.18215 * init_latents
       if len(init_latents) == 1:
         init_latents = init_latents.repeat((batch_size, 1, 1, 1))
       init_latents_orig = init_latents
@@ -898,21 +864,12 @@ class PipelineLike():
       extra_step_kwargs["eta"] = eta
 
     num_latent_input = (3 if negative_scale is not None else 2) if do_classifier_free_guidance else 1
-
-    if self.control_nets:
-      guided_hints = original_control_net.get_guided_hints(self.control_nets, num_latent_input, batch_size, clip_guide_images)
-
     for i, t in enumerate(tqdm(timesteps)):
       # expand the latents if we are doing classifier free guidance
       latent_model_input = latents.repeat((num_latent_input, 1, 1, 1))
       latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-
       # predict the noise residual
-      if self.control_nets:
-        noise_pred = original_control_net.call_unet_and_control_net(
-            i, num_latent_input, self.unet, self.control_nets, guided_hints, i / len(timesteps), latent_model_input, t, text_embeddings).sample
-      else:
-        noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
+      noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
 
       # perform guidance
       if do_classifier_free_guidance:
@@ -954,19 +911,8 @@ class PipelineLike():
         if is_cancelled_callback is not None and is_cancelled_callback():
           return None
 
-    if return_latents:
-      return (latents, False)
-
     latents = 1 / 0.18215 * latents
-    if vae_batch_size >= batch_size:
-      image = self.vae.decode(latents).sample
-    else:
-      if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-      images = []
-      for i in tqdm(range(0, batch_size, vae_batch_size)):
-        images.append(self.vae.decode(latents[i:i + vae_batch_size] if vae_batch_size > 1 else latents[i].unsqueeze(0)).sample)
-      image = torch.cat(images)
+    image = self.vae.decode(latents).sample
 
     image = (image / 2 + 0.5).clamp(0, 1)
 
@@ -1649,11 +1595,10 @@ def get_unweighted_text_embeddings(
       if pad == eos:                        # v1
         text_input_chunk[:, -1] = text_input[0, -1]
       else:                                 # v2
-        for j in range(len(text_input_chunk)):
-          if text_input_chunk[j, -1] != eos and text_input_chunk[j, -1] != pad:     # 最後に普通の文字がある
-            text_input_chunk[j, -1] = eos
-          if text_input_chunk[j, 1] == pad:                                         # BOSだけであとはPAD
-            text_input_chunk[j, 1] = eos
+        if text_input_chunk[:, -1] != eos and text_input_chunk[:, -1] != pad:     # 最後に普通の文字がある
+          text_input_chunk[:, -1] = eos
+        if text_input_chunk[:, 1] == pad:                                         # BOSだけであとはPAD
+          text_input_chunk[:, 1] = eos
 
       if clip_skip is None or clip_skip == 1:
         text_embedding = pipe.text_encoder(text_input_chunk)[0]
@@ -1854,7 +1799,7 @@ def preprocess_mask(mask):
   mask = mask.convert("L")
   w, h = mask.size
   w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
-  mask = mask.resize((w // 8, h // 8), resample=PIL.Image.BILINEAR) # LANCZOS)
+  mask = mask.resize((w // 8, h // 8), resample=PIL.Image.LANCZOS)
   mask = np.array(mask).astype(np.float32) / 255.0
   mask = np.tile(mask, (4, 1, 1))
   mask = mask[None].transpose(0, 1, 2, 3)  # what does this step do?
@@ -1872,35 +1817,6 @@ def preprocess_mask(mask):
 #   return text_encoder
 
 
-class BatchDataBase(NamedTuple):
-  # バッチ分割が必要ないデータ
-  step: int
-  prompt: str
-  negative_prompt: str
-  seed: int
-  init_image: Any
-  mask_image: Any
-  clip_prompt: str
-  guide_image: Any
-
-
-class BatchDataExt(NamedTuple):
-  # バッチ分割が必要なデータ
-  width: int
-  height: int
-  steps: int
-  scale:  float
-  negative_scale: float
-  strength: float
-  network_muls: Tuple[float]
-
-
-class BatchData(NamedTuple):
-  return_latents: bool
-  base: BatchDataBase
-  ext: BatchDataExt
-
-
 def main(args):
   if args.fp16:
     dtype = torch.float16
@@ -1965,7 +1881,10 @@ def main(args):
   # tokenizerを読み込む
   print("loading tokenizer")
   if use_stable_diffusion_format:
-    tokenizer = train_util.load_tokenizer(args)
+    if args.v2:
+      tokenizer = CLIPTokenizer.from_pretrained(V2_STABLE_DIFFUSION_PATH, subfolder="tokenizer")
+    else:
+      tokenizer = CLIPTokenizer.from_pretrained(TOKENIZER_PATH)  # , model_max_length=max_token_length + 2)
 
   # schedulerを用意する
   sched_init_args = {}
@@ -2076,13 +1995,11 @@ def main(args):
   # networkを組み込む
   if args.network_module:
     networks = []
-    network_default_muls = []
     for i, network_module in enumerate(args.network_module):
       print("import network module:", network_module)
       imported_module = importlib.import_module(network_module)
 
       network_mul = 1.0 if args.network_mul is None or len(args.network_mul) <= i else args.network_mul[i]
-      network_default_muls.append(network_mul)
 
       net_kwargs = {}
       if args.network_args and i < len(args.network_args):
@@ -2097,7 +2014,7 @@ def main(args):
         network_weight = args.network_weights[i]
         print("load network weights from:", network_weight)
 
-        if model_util.is_safetensors(network_weight) and args.network_show_meta:
+        if model_util.is_safetensors(network_weight):
           from safetensors.torch import safe_open
           with safe_open(network_weight, framework="pt") as f:
             metadata = f.metadata()
@@ -2120,18 +2037,6 @@ def main(args):
   else:
     networks = []
 
-  # ControlNetの処理
-  control_nets: List[ControlNetInfo] = []
-  if args.control_net_models:
-    for i, model in enumerate(args.control_net_models):
-      prep_type = None if not args.control_net_preps or len(args.control_net_preps) <= i else args.control_net_preps[i]
-      weight = 1.0 if not args.control_net_weights or len(args.control_net_weights) <= i else args.control_net_weights[i]
-      ratio = 1.0 if not args.control_net_ratios or len(args.control_net_ratios) <= i else args.control_net_ratios[i]
-
-      ctrl_unet, ctrl_net = original_control_net.load_control_net(args.v2, unet, model)
-      prep = original_control_net.load_preprocess(prep_type)
-      control_nets.append(ControlNetInfo(ctrl_unet, ctrl_net, prep, weight, ratio))
-
   if args.opt_channels_last:
     print(f"set optimizing: channels last")
     text_encoder.to(memory_format=torch.channels_last)
@@ -2145,14 +2050,9 @@ def main(args):
     if vgg16_model is not None:
       vgg16_model.to(memory_format=torch.channels_last)
 
-    for cn in control_nets:
-      cn.unet.to(memory_format=torch.channels_last)
-      cn.net.to(memory_format=torch.channels_last)
-
   pipe = PipelineLike(device, vae, text_encoder, tokenizer, unet, scheduler, args.clip_skip,
                       clip_model, args.clip_guidance_scale, args.clip_image_guidance_scale,
                       vgg16_model, args.vgg16_guidance_scale, args.vgg16_guidance_layer)
-  pipe.set_control_nets(control_nets)
   print("pipeline is ready.")
 
   if args.diffusers_xformers:
@@ -2277,34 +2177,18 @@ def main(args):
       mask_images = l
 
   # 画像サイズにオプション指定があるときはリサイズする
-  if args.W is not None and args.H is not None:
-    if init_images is not None:
-      print(f"resize img2img source images to {args.W}*{args.H}")
-      init_images = resize_images(init_images, (args.W, args.H))
+  if init_images is not None and args.W is not None and args.H is not None:
+    print(f"resize img2img source images to {args.W}*{args.H}")
+    init_images = resize_images(init_images, (args.W, args.H))
     if mask_images is not None:
       print(f"resize img2img mask images to {args.W}*{args.H}")
       mask_images = resize_images(mask_images, (args.W, args.H))
 
-  if networks and mask_images:
-    # mask を領域情報として流用する、現在は1枚だけ対応
-    # TODO 複数のnetwork classの混在時の考慮
-    print("use mask as region")
-    # import cv2
-    # for i in range(3):
-    #   cv2.imshow("msk", np.array(mask_images[0])[:,:,i])
-    #   cv2.waitKey()
-    #   cv2.destroyAllWindows()
-    networks[0].__class__.set_regions(networks, np.array(mask_images[0]))
-    mask_images = None
-
   prev_image = None               # for VGG16 guided
   if args.guide_image_path is not None:
-    print(f"load image for CLIP/VGG16/ControlNet guidance: {args.guide_image_path}")
-    guide_images = []
-    for p in args.guide_image_path:
-      guide_images.extend(load_images(p))
-
-    print(f"loaded {len(guide_images)} guide images for guidance")
+    print(f"load image for CLIP/VGG16 guidance: {args.guide_image_path}")
+    guide_images = load_images(args.guide_image_path)
+    print(f"loaded {len(guide_images)} guide images for CLIP/VGG16 guidance")
     if len(guide_images) == 0:
       print(f"No guide image, use previous generated image. / ガイド画像がありません。直前に生成した画像を使います: {args.image_path}")
       guide_images = None
@@ -2335,46 +2219,33 @@ def main(args):
     iter_seed = random.randint(0, 0x7fffffff)
 
     # バッチ処理の関数
-    def process_batch(batch: List[BatchData], highres_fix, highres_1st=False):
+    def process_batch(batch, highres_fix, highres_1st=False):
       batch_size = len(batch)
 
       # highres_fixの処理
       if highres_fix and not highres_1st:
-        # 1st stageのバッチを作成して呼び出す：サイズを小さくして呼び出す
-        print("process 1st stage")
+        # 1st stageのバッチを作成して呼び出す
+        print("process 1st stage1")
         batch_1st = []
-        for _, base, ext in batch:
-          width_1st = int(ext.width * args.highres_fix_scale + .5)
-          height_1st = int(ext.height * args.highres_fix_scale + .5)
+        for params1, (width, height, steps, scale, negative_scale, strength) in batch:
+          width_1st = int(width * args.highres_fix_scale + .5)
+          height_1st = int(height * args.highres_fix_scale + .5)
           width_1st = width_1st - width_1st % 32
           height_1st = height_1st - height_1st % 32
-
-          ext_1st = BatchDataExt(width_1st, height_1st, args.highres_fix_steps, ext.scale,
-                                 ext.negative_scale, ext.strength, ext.network_muls)
-          batch_1st.append(BatchData(args.highres_fix_latents_upscaling, base, ext_1st))
+          batch_1st.append((params1, (width_1st, height_1st, args.highres_fix_steps, scale, negative_scale, strength)))
         images_1st = process_batch(batch_1st, True, True)
 
         # 2nd stageのバッチを作成して以下処理する
-        print("process 2nd stage")
-        if args.highres_fix_latents_upscaling:
-          org_dtype = images_1st.dtype
-          if images_1st.dtype == torch.bfloat16:
-            images_1st = images_1st.to(torch.float)                 # interpolateがbf16をサポートしていない
-          images_1st = torch.nn.functional.interpolate(
-              images_1st, (batch[0].ext.height // 8, batch[0].ext.width // 8), mode='bilinear')  # , antialias=True)
-          images_1st = images_1st.to(org_dtype)
-
+        print("process 2nd stage1")
         batch_2nd = []
-        for i, (bd, image) in enumerate(zip(batch, images_1st)):
-          if not args.highres_fix_latents_upscaling:
-            image = image.resize((bd.ext.width, bd.ext.height), resample=PIL.Image.LANCZOS)      # img2imgとして設定
-          bd_2nd = BatchData(False, BatchDataBase(*bd.base[0:3], bd.base.seed+1, image, None, *bd.base[6:]), bd.ext)
-          batch_2nd.append(bd_2nd)
+        for i, (b1, image) in enumerate(zip(batch, images_1st)):
+          image = image.resize((width, height), resample=PIL.Image.LANCZOS)
+          (step, prompt, negative_prompt, seed, _, _, clip_prompt, guide_image), params2 = b1
+          batch_2nd.append(((step, prompt, negative_prompt, seed+1, image, None, clip_prompt, guide_image), params2))
         batch = batch_2nd
 
-      # このバッチの情報を取り出す
-      return_latents, (step_first, _, _, _, init_image, mask_image, _, guide_image), \
-          (width, height, steps, scale, negative_scale, strength, network_muls) = batch[0]
+      (step_first, _, _, _, init_image, mask_image, _, guide_image), (width,
+                                                                      height, steps, scale, negative_scale, strength) = batch[0]
       noise_shape = (LATENT_CHANNELS, height // DOWNSAMPLING_FACTOR, width // DOWNSAMPLING_FACTOR)
 
       prompts = []
@@ -2407,7 +2278,7 @@ def main(args):
       all_images_are_same = True
       all_masks_are_same = True
       all_guide_images_are_same = True
-      for i, (_, (_, prompt, negative_prompt, seed, init_image, mask_image, clip_prompt, guide_image), _) in enumerate(batch):
+      for i, ((_, prompt, negative_prompt, seed, init_image, mask_image, clip_prompt, guide_image), _) in enumerate(batch):
         prompts.append(prompt)
         negative_prompts.append(negative_prompt)
         seeds.append(seed)
@@ -2424,13 +2295,9 @@ def main(args):
             all_masks_are_same = mask_images[-2] is mask_image
 
         if guide_image is not None:
-          if type(guide_image) is list:
-            guide_images.extend(guide_image)
-            all_guide_images_are_same = False
-          else:
-            guide_images.append(guide_image)
-            if i > 0 and all_guide_images_are_same:
-              all_guide_images_are_same = guide_images[-2] is guide_image
+          guide_images.append(guide_image)
+          if i > 0 and all_guide_images_are_same:
+            all_guide_images_are_same = guide_images[-2] is guide_image
 
         # make start code
         torch.manual_seed(seed)
@@ -2453,24 +2320,10 @@ def main(args):
       if guide_images is not None and all_guide_images_are_same:
         guide_images = guide_images[0]
 
-      # ControlNet使用時はguide imageをリサイズする
-      if control_nets:
-        # TODO resample��メソッド
-        guide_images = guide_images if type(guide_images) == list else [guide_images]
-        guide_images = [i.resize((width, height), resample=PIL.Image.LANCZOS) for i in guide_images]
-        if len(guide_images) == 1:
-          guide_images = guide_images[0]
-
       # generate
-      if networks:
-        for n, m in zip(networks, network_muls if network_muls else network_default_muls):
-          n.set_multiplier(m)
-
       images = pipe(prompts, negative_prompts, init_images, mask_images, height, width, steps, scale, negative_scale, strength, latents=start_code,
-                    output_type='pil', max_embeddings_multiples=max_embeddings_multiples, img2img_noise=i2i_noises,
-                    vae_batch_size=args.vae_batch_size, return_latents=return_latents,
-                    clip_prompts=clip_prompts, clip_guide_images=guide_images)[0]
-      if highres_1st and not args.highres_fix_save_1st:             # return images or latents
+                    output_type='pil', max_embeddings_multiples=max_embeddings_multiples, img2img_noise=i2i_noises, clip_prompts=clip_prompts, clip_guide_images=guide_images)[0]
+      if highres_1st and not args.highres_fix_save_1st:
         return images
 
       # save image
@@ -2545,7 +2398,6 @@ def main(args):
       strength = 0.8 if args.strength is None else args.strength
       negative_prompt = ""
       clip_prompt = None
-      network_muls = None
 
       prompt_args = prompt.strip().split(' --')
       prompt = prompt_args[0]
@@ -2609,15 +2461,6 @@ def main(args):
             clip_prompt = m.group(1)
             print(f"clip prompt: {clip_prompt}")
             continue
-
-          m = re.match(r'am ([\d\.\-,]+)', parg, re.IGNORECASE)
-          if m:               # network multiplies
-            network_muls = [float(v) for v in m.group(1).split(",")]
-            while len(network_muls) < len(networks):
-              network_muls.append(network_muls[-1])
-            print(f"network mul: {network_muls}")
-            continue
-
         except ValueError as ex:
           print(f"Exception in parsing / 解析エラー: {parg}")
           print(ex)
@@ -2655,12 +2498,7 @@ def main(args):
           mask_image = mask_images[global_step % len(mask_images)]
 
         if guide_images is not None:
-          if control_nets:                                                        # 複数件の場合あり
-            c = len(control_nets)
-            p = global_step % (len(guide_images) // c)
-            guide_image = guide_images[p * c:p * c + c]
-          else:
-            guide_image = guide_images[global_step % len(guide_images)]
+          guide_image = guide_images[global_step % len(guide_images)]
         elif args.clip_image_guidance_scale > 0 or args.vgg16_guidance_scale > 0:
           if prev_image is None:
             print("Generate 1st image without guide image.")
@@ -2668,9 +2506,10 @@ def main(args):
             print("Use previous image as guide image.")
             guide_image = prev_image
 
-        b1 = BatchData(False, BatchDataBase(global_step, prompt, negative_prompt, seed, init_image, mask_image, clip_prompt, guide_image),
-                       BatchDataExt(width, height, steps, scale, negative_scale, strength, tuple(network_muls) if network_muls else None))
-        if len(batch_data) > 0 and batch_data[-1].ext != b1.ext:  # バッチ分割必要？
+        # TODO named tupleか何かにする
+        b1 = ((global_step, prompt, negative_prompt, seed, init_image, mask_image, clip_prompt, guide_image),
+              (width, height, steps, scale, negative_scale, strength))
+        if len(batch_data) > 0 and batch_data[-1][1] != b1[1]:  # バッチ分割必要？
           process_batch(batch_data, highres_fix)
           batch_data.clear()
 
@@ -2714,8 +2553,6 @@ if __name__ == '__main__':
   parser.add_argument("--H", type=int, default=None, help="image height, in pixel space / 生成画像高さ")
   parser.add_argument("--W", type=int, default=None, help="image width, in pixel space / 生成画像幅")
   parser.add_argument("--batch_size", type=int, default=1, help="batch size / バッチサイズ")
-  parser.add_argument("--vae_batch_size", type=float, default=None,
-                      help="batch size for VAE, < 1.0 for ratio / VAE処理時のバッチサイズ、1未満の値の場合は通常バッチサイズの比率")
   parser.add_argument("--steps", type=int, default=50, help="number of ddim sampling steps / サンプリングステップ数")
   parser.add_argument('--sampler', type=str, default='ddim',
                       choices=['ddim', 'pndm', 'lms', 'euler', 'euler_a', 'heun', 'dpm_2', 'dpm_2_a', 'dpmsolver',
@@ -2727,8 +2564,6 @@ if __name__ == '__main__':
   parser.add_argument("--ckpt", type=str, default=None, help="path to checkpoint of model / モデルのcheckpointファイルまたはディレクトリ")
   parser.add_argument("--vae", type=str, default=None,
                       help="path to checkpoint of vae to replace / VAEを入れ替える場合、VAEのcheckpointファイルまたはディレクトリ")
-  parser.add_argument("--tokenizer_cache_dir", type=str, default=None,
-                      help="directory for caching Tokenizer (for offline training) / Tokenizerをキャッシュするディレクトリ（ネット接続なしでの学習のため）")
   # parser.add_argument("--replace_clip_l14_336", action='store_true',
   #                     help="Replace CLIP (Text Encoder) to l/14@336 / CLIP(Text Encoder)をl/14@336に入れ替える")
   parser.add_argument("--seed", type=int, default=None,
@@ -2743,15 +2578,12 @@ if __name__ == '__main__':
   parser.add_argument("--opt_channels_last", action='store_true',
                       help='set channels last option to model / モデルにchannels lastを指定し最適化する')
   parser.add_argument("--network_module", type=str, default=None, nargs='*',
-                      help='additional network module to use / 追加ネットワークを使う時そのモジュール名')
+                      help='Hypernetwork module to use / Hypernetworkを使う時そのモジュール名')
   parser.add_argument("--network_weights", type=str, default=None, nargs='*',
-                      help='additional network weights to load / 追加ネットワークの重み')
-  parser.add_argument("--network_mul", type=float, default=None, nargs='*',
-                      help='additional network multiplier / 追加ネットワークの効果の倍率')
+                      help='Hypernetwork weights to load / Hypernetworkの重み')
+  parser.add_argument("--network_mul", type=float, default=None, nargs='*', help='Hypernetwork multiplier / Hypernetworkの効果の倍率')
   parser.add_argument("--network_args", type=str, default=None, nargs='*',
                       help='additional argmuments for network (key=value) / ネットワークへの追加の引数')
-  parser.add_argument("--network_show_meta", action='store_true',
-                      help='show metadata of network model / ネットワークモデルのメタデータを表示する')
   parser.add_argument("--textual_inversion_embeddings", type=str, default=None, nargs='*',
                       help='Embeddings files of Textual Inversion / Textual Inversionのembeddings')
   parser.add_argument("--clip_skip", type=int, default=None, help='layer number from bottom to use in CLIP / CLIPの後ろからn層目の出力を使う')
@@ -2765,26 +2597,15 @@ if __name__ == '__main__':
                       help='enable VGG16 guided SD by image, scale for guidance / 画像によるVGG16 guided SDを有効にしてこのscaleを適用する')
   parser.add_argument("--vgg16_guidance_layer", type=int, default=20,
                       help='layer of VGG16 to calculate contents guide (1~30, 20 for conv4_2) / VGG16のcontents guideに使うレイヤー番号 (1~30、20はconv4_2)')
-  parser.add_argument("--guide_image_path", type=str, default=None, nargs="*",
-                      help="image to CLIP guidance / CLIP guided SDでガイドに使う画像")
+  parser.add_argument("--guide_image_path", type=str, default=None, help="image to CLIP guidance / CLIP guided SDでガイドに使う画像")
   parser.add_argument("--highres_fix_scale", type=float, default=None,
                       help="enable highres fix, reso scale for 1st stage / highres fixを有効にして最初の解像度をこのscaleにする")
   parser.add_argument("--highres_fix_steps", type=int, default=28,
                       help="1st stage steps for highres fix / highres fixの最初のステージのステップ数")
   parser.add_argument("--highres_fix_save_1st", action='store_true',
                       help="save 1st stage images for highres fix / highres fixの最初のステージの画像を保存する")
-  parser.add_argument("--highres_fix_latents_upscaling", action='store_true',
-                      help="use latents upscaling for highres fix / highres fixでlatentで拡大する")
   parser.add_argument("--negative_scale", type=float, default=None,
                       help="set another guidance scale for negative prompt / ネガティブプロンプトのscaleを指定する")
 
-  parser.add_argument("--control_net_models", type=str, default=None, nargs='*',
-                      help='ControlNet models to use / 使用するControlNetのモデル名')
-  parser.add_argument("--control_net_preps", type=str, default=None, nargs='*',
-                      help='ControlNet preprocess to use / 使用するControlNetのプリプロセス名')
-  parser.add_argument("--control_net_weights", type=float, default=None, nargs='*', help='ControlNet weights / ControlNetの重み')
-  parser.add_argument("--control_net_ratios", type=float, default=None, nargs='*',
-                      help='ControlNet guidance ratio for steps / ControlNetでガイドするステップ比率')
-
   args = parser.parse_args()
   main(args)

library.egg-info/PKG-INFO

@@ -0,0 +1,4 @@
+Metadata-Version: 2.1
+Name: library
+Version: 0.0.0
+License-File: LICENSE.md

library.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+LICENSE.md
+README.md
+setup.py
+library/__init__.py
+library/model_util.py
+library/train_util.py
+library.egg-info/PKG-INFO
+library.egg-info/SOURCES.txt
+library.egg-info/dependency_links.txt
+library.egg-info/top_level.txt

library.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

library.egg-info/top_level.txt

@@ -0,0 +1 @@
+library

library/model_util.py

@@ -4,7 +4,7 @@
 import math
 import os
 import torch
-from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig, logging
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig
 from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNet2DConditionModel
 from safetensors.torch import load_file, save_file
 
@@ -916,11 +916,7 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, dtype=None):
     info = text_model.load_state_dict(converted_text_encoder_checkpoint)
   else:
     converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v1(state_dict)
-
-    logging.set_verbosity_error()                                                       # don't show annoying warning
     text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
-    logging.set_verbosity_warning()
-
     info = text_model.load_state_dict(converted_text_encoder_checkpoint)
   print("loading text encoder:", info)
 

library/train_util.py

@@ -1,21 +1,12 @@
 # common functions for training
 
 import argparse
-import importlib
 import json
-import re
 import shutil
 import time
-from typing import (
-    Dict,
-    List,
-    NamedTuple,
-    Optional,
-    Sequence,
-    Tuple,
-    Union,
-)
+from typing import Dict, List, NamedTuple, Tuple
 from accelerate import Accelerator
+from torch.autograd.function import Function
 import glob
 import math
 import os
@@ -26,16 +17,10 @@ from io import BytesIO
 
 from tqdm import tqdm
 import torch
-from torch.optim import Optimizer
 from torchvision import transforms
 from transformers import CLIPTokenizer
-import transformers
 import diffusers
-from diffusers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION
-from diffusers import (StableDiffusionPipeline, DDPMScheduler,
-                       EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler,
-                       LMSDiscreteScheduler, PNDMScheduler, DDIMScheduler, EulerDiscreteScheduler, HeunDiscreteScheduler,
-                       KDPM2DiscreteScheduler, KDPM2AncestralDiscreteScheduler)
+from diffusers import DDPMScheduler, StableDiffusionPipeline
 import albumentations as albu
 import numpy as np
 from PIL import Image
@@ -210,95 +195,23 @@ class BucketBatchIndex(NamedTuple):
   batch_index: int
 
 
-class AugHelper:
-  def __init__(self):
-    # prepare all possible augmentators
-    color_aug_method = albu.OneOf([
-        albu.HueSaturationValue(8, 0, 0, p=.5),
-        albu.RandomGamma((95, 105), p=.5),
-    ], p=.33)
-    flip_aug_method = albu.HorizontalFlip(p=0.5)
-
-    # key: (use_color_aug, use_flip_aug)
-    self.augmentors = {
-        (True, True): albu.Compose([
-            color_aug_method,
-            flip_aug_method,
-        ], p=1.),
-        (True, False): albu.Compose([
-            color_aug_method,
-        ], p=1.),
-        (False, True): albu.Compose([
-            flip_aug_method,
-        ], p=1.),
-        (False, False): None
-    }
-
-  def get_augmentor(self, use_color_aug: bool, use_flip_aug: bool) -> Optional[albu.Compose]:
-    return self.augmentors[(use_color_aug, use_flip_aug)]
-
-
-class BaseSubset:
-  def __init__(self, image_dir: Optional[str], num_repeats: int, shuffle_caption: bool, keep_tokens: int, color_aug: bool, flip_aug: bool, face_crop_aug_range: Optional[Tuple[float, float]], random_crop: bool, caption_dropout_rate: float, caption_dropout_every_n_epochs: int, caption_tag_dropout_rate: float) -> None:
-    self.image_dir = image_dir
-    self.num_repeats = num_repeats
-    self.shuffle_caption = shuffle_caption
-    self.keep_tokens = keep_tokens
-    self.color_aug = color_aug
-    self.flip_aug = flip_aug
-    self.face_crop_aug_range = face_crop_aug_range
-    self.random_crop = random_crop
-    self.caption_dropout_rate = caption_dropout_rate
-    self.caption_dropout_every_n_epochs = caption_dropout_every_n_epochs
-    self.caption_tag_dropout_rate = caption_tag_dropout_rate
-
-    self.img_count = 0
-
-
-class DreamBoothSubset(BaseSubset):
-  def __init__(self, image_dir: str, is_reg: bool, class_tokens: Optional[str], caption_extension: str, num_repeats, shuffle_caption, keep_tokens, color_aug, flip_aug, face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate) -> None:
-    assert image_dir is not None, "image_dir must be specified / image_dirは指定が必須です"
-
-    super().__init__(image_dir, num_repeats, shuffle_caption, keep_tokens, color_aug, flip_aug,
-                     face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate)
-
-    self.is_reg = is_reg
-    self.class_tokens = class_tokens
-    self.caption_extension = caption_extension
-
-  def __eq__(self, other) -> bool:
-    if not isinstance(other, DreamBoothSubset):
-      return NotImplemented
-    return self.image_dir == other.image_dir
-
-
-class FineTuningSubset(BaseSubset):
-  def __init__(self, image_dir, metadata_file: str, num_repeats, shuffle_caption, keep_tokens, color_aug, flip_aug, face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate) -> None:
-    assert metadata_file is not None, "metadata_file must be specified / metadata_fileは指定が必須です"
-
-    super().__init__(image_dir, num_repeats, shuffle_caption, keep_tokens, color_aug, flip_aug,
-                     face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate)
-
-    self.metadata_file = metadata_file
-
-  def __eq__(self, other) -> bool:
-    if not isinstance(other, FineTuningSubset):
-      return NotImplemented
-    return self.metadata_file == other.metadata_file
-
-
 class BaseDataset(torch.utils.data.Dataset):
-  def __init__(self, tokenizer: CLIPTokenizer, max_token_length: int, resolution: Optional[Tuple[int, int]], debug_dataset: bool) -> None:
+  def __init__(self, tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens, resolution, flip_aug: bool, color_aug: bool, face_crop_aug_range, random_crop, debug_dataset: bool) -> None:
     super().__init__()
-    self.tokenizer = tokenizer
+    self.tokenizer: CLIPTokenizer = tokenizer
     self.max_token_length = max_token_length
+    self.shuffle_caption = shuffle_caption
+    self.shuffle_keep_tokens = shuffle_keep_tokens
     # width/height is used when enable_bucket==False
     self.width, self.height = (None, None) if resolution is None else resolution
+    self.face_crop_aug_range = face_crop_aug_range
+    self.flip_aug = flip_aug
+    self.color_aug = color_aug
     self.debug_dataset = debug_dataset
-
-    self.subsets: List[Union[DreamBoothSubset, FineTuningSubset]] = []
-
+    self.random_crop = random_crop
     self.token_padding_disabled = False
+    self.dataset_dirs_info = {}
+    self.reg_dataset_dirs_info = {}
     self.tag_frequency = {}
 
     self.enable_bucket = False
@@ -312,28 +225,49 @@ class BaseDataset(torch.utils.data.Dataset):
     self.tokenizer_max_length = self.tokenizer.model_max_length if max_token_length is None else max_token_length + 2
 
     self.current_epoch: int = 0            # インスタンスがepochごとに新しく作られるようなので外側から渡さないとダメ
+    self.dropout_rate: float = 0
+    self.dropout_every_n_epochs: int = None
+    self.tag_dropout_rate: float = 0
 
     # augmentation
-    self.aug_helper = AugHelper()
+    flip_p = 0.5 if flip_aug else 0.0
+    if color_aug:
+      # わりと弱めの色合いaugmentation：brightness/contrastあたりは画像のpixel valueの最大値・最小値を変えてしまうのでよくないのではという想定でgamma/hueあたりを触る
+      self.aug = albu.Compose([
+          albu.OneOf([
+              albu.HueSaturationValue(8, 0, 0, p=.5),
+              albu.RandomGamma((95, 105), p=.5),
+          ], p=.33),
+          albu.HorizontalFlip(p=flip_p)
+      ], p=1.)
+    elif flip_aug:
+      self.aug = albu.Compose([
+          albu.HorizontalFlip(p=flip_p)
+      ], p=1.)
+    else:
+      self.aug = None
 
     self.image_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5]), ])
 
     self.image_data: Dict[str, ImageInfo] = {}
-    self.image_to_subset: Dict[str, Union[DreamBoothSubset, FineTuningSubset]] = {}
 
     self.replacements = {}
 
   def set_current_epoch(self, epoch):
     self.current_epoch = epoch
-    self.shuffle_buckets()
+
+  def set_caption_dropout(self, dropout_rate, dropout_every_n_epochs, tag_dropout_rate):
+    # コンストラクタで渡さないのはTextual Inversionで意識したくないから（ということにしておく）
+    self.dropout_rate = dropout_rate
+    self.dropout_every_n_epochs = dropout_every_n_epochs
+    self.tag_dropout_rate = tag_dropout_rate
 
   def set_tag_frequency(self, dir_name, captions):
     frequency_for_dir = self.tag_frequency.get(dir_name, {})
     self.tag_frequency[dir_name] = frequency_for_dir
     for caption in captions:
       for tag in caption.split(","):
-        tag = tag.strip()
-        if tag:
+        if tag and not tag.isspace():
           tag = tag.lower()
           frequency = frequency_for_dir.get(tag, 0)
           frequency_for_dir[tag] = frequency + 1
@@ -344,36 +278,42 @@ class BaseDataset(torch.utils.data.Dataset):
   def add_replacement(self, str_from, str_to):
     self.replacements[str_from] = str_to
 
-  def process_caption(self, subset: BaseSubset, caption):
+  def process_caption(self, caption):
     # dropoutの決定：tag dropがこのメソッド内にあるのでここで行うのが良い
-    is_drop_out = subset.caption_dropout_rate > 0 and random.random() < subset.caption_dropout_rate
-    is_drop_out = is_drop_out or subset.caption_dropout_every_n_epochs > 0 and self.current_epoch % subset.caption_dropout_every_n_epochs == 0
+    is_drop_out = self.dropout_rate > 0 and random.random() < self.dropout_rate
+    is_drop_out = is_drop_out or self.dropout_every_n_epochs and self.current_epoch % self.dropout_every_n_epochs == 0
 
     if is_drop_out:
       caption = ""
     else:
-      if subset.shuffle_caption or subset.caption_tag_dropout_rate > 0:
+      if self.shuffle_caption or self.tag_dropout_rate > 0:
         def dropout_tags(tokens):
-          if subset.caption_tag_dropout_rate <= 0:
+          if self.tag_dropout_rate <= 0:
             return tokens
           l = []
           for token in tokens:
-            if random.random() >= subset.caption_tag_dropout_rate:
+            if random.random() >= self.tag_dropout_rate:
               l.append(token)
           return l
 
-        fixed_tokens = []
-        flex_tokens = [t.strip() for t in caption.strip().split(",")]
-        if subset.keep_tokens > 0:
-          fixed_tokens = flex_tokens[:subset.keep_tokens]
-          flex_tokens = flex_tokens[subset.keep_tokens:]
+        tokens = [t.strip() for t in caption.strip().split(",")]
+        if self.shuffle_keep_tokens is None:
+          if self.shuffle_caption:
+            random.shuffle(tokens)
+
+          tokens = dropout_tags(tokens)
+        else:
+          if len(tokens) > self.shuffle_keep_tokens:
+            keep_tokens = tokens[:self.shuffle_keep_tokens]
+            tokens = tokens[self.shuffle_keep_tokens:]
 
-        if subset.shuffle_caption:
-          random.shuffle(flex_tokens)
+            if self.shuffle_caption:
+              random.shuffle(tokens)
 
-        flex_tokens = dropout_tags(flex_tokens)
+            tokens = dropout_tags(tokens)
 
-        caption = ", ".join(fixed_tokens + flex_tokens)
+            tokens = keep_tokens + tokens
+        caption = ", ".join(tokens)
 
       # textual inversion対応
       for str_from, str_to in self.replacements.items():
@@ -427,9 +367,8 @@ class BaseDataset(torch.utils.data.Dataset):
       input_ids = torch.stack(iids_list)      # 3,77
     return input_ids
 
-  def register_image(self, info: ImageInfo, subset: BaseSubset):
+  def register_image(self, info: ImageInfo):
     self.image_data[info.image_key] = info
-    self.image_to_subset[info.image_key] = subset
 
   def make_buckets(self):
     '''
@@ -528,7 +467,7 @@ class BaseDataset(torch.utils.data.Dataset):
     img = np.array(image, np.uint8)
     return img
 
-  def trim_and_resize_if_required(self, subset: BaseSubset, image, reso, resized_size):
+  def trim_and_resize_if_required(self, image, reso, resized_size):
     image_height, image_width = image.shape[0:2]
 
     if image_width != resized_size[0] or image_height != resized_size[1]:
@@ -538,27 +477,22 @@ class BaseDataset(torch.utils.data.Dataset):
     image_height, image_width = image.shape[0:2]
     if image_width > reso[0]:
       trim_size = image_width - reso[0]
-      p = trim_size // 2 if not subset.random_crop else random.randint(0, trim_size)
+      p = trim_size // 2 if not self.random_crop else random.randint(0, trim_size)
       # print("w", trim_size, p)
       image = image[:, p:p + reso[0]]
     if image_height > reso[1]:
       trim_size = image_height - reso[1]
-      p = trim_size // 2 if not subset.random_crop else random.randint(0, trim_size)
+      p = trim_size // 2 if not self.random_crop else random.randint(0, trim_size)
       # print("h", trim_size, p)
       image = image[p:p + reso[1]]
 
     assert image.shape[0] == reso[1] and image.shape[1] == reso[0], f"internal error, illegal trimmed size: {image.shape}, {reso}"
     return image
 
-  def is_latent_cacheable(self):
-    return all([not subset.color_aug and not subset.random_crop for subset in self.subsets])
-
   def cache_latents(self, vae):
     # TODO ここを高速化したい
     print("caching latents.")
     for info in tqdm(self.image_data.values()):
-      subset = self.image_to_subset[info.image_key]
-
       if info.latents_npz is not None:
         info.latents = self.load_latents_from_npz(info, False)
         info.latents = torch.FloatTensor(info.latents)
@@ -568,13 +502,13 @@ class BaseDataset(torch.utils.data.Dataset):
         continue
 
       image = self.load_image(info.absolute_path)
-      image = self.trim_and_resize_if_required(subset, image, info.bucket_reso, info.resized_size)
+      image = self.trim_and_resize_if_required(image, info.bucket_reso, info.resized_size)
 
       img_tensor = self.image_transforms(image)
       img_tensor = img_tensor.unsqueeze(0).to(device=vae.device, dtype=vae.dtype)
       info.latents = vae.encode(img_tensor).latent_dist.sample().squeeze(0).to("cpu")
 
-      if subset.flip_aug:
+      if self.flip_aug:
         image = image[:, ::-1].copy()     # cannot convert to Tensor without copy
         img_tensor = self.image_transforms(image)
         img_tensor = img_tensor.unsqueeze(0).to(device=vae.device, dtype=vae.dtype)
@@ -584,11 +518,11 @@ class BaseDataset(torch.utils.data.Dataset):
     image = Image.open(image_path)
     return image.size
 
-  def load_image_with_face_info(self, subset: BaseSubset, image_path: str):
+  def load_image_with_face_info(self, image_path: str):
     img = self.load_image(image_path)
 
     face_cx = face_cy = face_w = face_h = 0
-    if subset.face_crop_aug_range is not None:
+    if self.face_crop_aug_range is not None:
       tokens = os.path.splitext(os.path.basename(image_path))[0].split('_')
       if len(tokens) >= 5:
         face_cx = int(tokens[-4])
@@ -599,7 +533,7 @@ class BaseDataset(torch.utils.data.Dataset):
     return img, face_cx, face_cy, face_w, face_h
 
   # いい感じに切り出す
-  def crop_target(self, subset: BaseSubset, image, face_cx, face_cy, face_w, face_h):
+  def crop_target(self, image, face_cx, face_cy, face_w, face_h):
     height, width = image.shape[0:2]
     if height == self.height and width == self.width:
       return image
@@ -607,8 +541,8 @@ class BaseDataset(torch.utils.data.Dataset):
     # 画像サイズはsizeより大きいのでリサイズする
     face_size = max(face_w, face_h)
     min_scale = max(self.height / height, self.width / width)        # 画像がモデル入力サイズぴったりになる倍率（最小の倍率）
-    min_scale = min(1.0, max(min_scale, self.size / (face_size * subset.face_crop_aug_range[1])))             # 指定した顔最小サイズ
-    max_scale = min(1.0, max(min_scale, self.size / (face_size * subset.face_crop_aug_range[0])))             # 指定した顔最大サイズ
+    min_scale = min(1.0, max(min_scale, self.size / (face_size * self.face_crop_aug_range[1])))             # 指定した顔最小サイズ
+    max_scale = min(1.0, max(min_scale, self.size / (face_size * self.face_crop_aug_range[0])))             # 指定した顔最大サイズ
     if min_scale >= max_scale:          # range指定がmin==max
       scale = min_scale
     else:
@@ -626,13 +560,13 @@ class BaseDataset(torch.utils.data.Dataset):
     for axis, (target_size, length, face_p) in enumerate(zip((self.height, self.width), (height, width), (face_cy, face_cx))):
       p1 = face_p - target_size // 2                # 顔を中心に持ってくるための切り出し位置
 
-      if subset.random_crop:
+      if self.random_crop:
         # 背景も含めるために顔を中心に置く確率を高めつつずらす
         range = max(length - face_p, face_p)        # 画像の端から顔中心までの距離の長いほう
         p1 = p1 + (random.randint(0, range) + random.randint(0, range)) - range     # -range ~ +range までのいい感じの乱数
       else:
         # range指定があるときのみ、すこしだけランダムに（わりと適当）
-        if subset.face_crop_aug_range[0] != subset.face_crop_aug_range[1]:
+        if self.face_crop_aug_range[0] != self.face_crop_aug_range[1]:
           if face_size > self.size // 10 and face_size >= 40:
             p1 = p1 + random.randint(-face_size // 20, +face_size // 20)
 
@@ -655,6 +589,9 @@ class BaseDataset(torch.utils.data.Dataset):
     return self._length
 
   def __getitem__(self, index):
+    if index == 0:
+      self.shuffle_buckets()
+
     bucket = self.bucket_manager.buckets[self.buckets_indices[index].bucket_index]
     bucket_batch_size = self.buckets_indices[index].bucket_batch_size
     image_index = self.buckets_indices[index].batch_index * bucket_batch_size
@@ -667,29 +604,28 @@ class BaseDataset(torch.utils.data.Dataset):
 
     for image_key in bucket[image_index:image_index + bucket_batch_size]:
       image_info = self.image_data[image_key]
-      subset = self.image_to_subset[image_key]
       loss_weights.append(self.prior_loss_weight if image_info.is_reg else 1.0)
 
       # image/latentsを処理する
       if image_info.latents is not None:
-        latents = image_info.latents if not subset.flip_aug or random.random() < .5 else image_info.latents_flipped
+        latents = image_info.latents if not self.flip_aug or random.random() < .5 else image_info.latents_flipped
         image = None
       elif image_info.latents_npz is not None:
-        latents = self.load_latents_from_npz(image_info, subset.flip_aug and random.random() >= .5)
+        latents = self.load_latents_from_npz(image_info, self.flip_aug and random.random() >= .5)
         latents = torch.FloatTensor(latents)
         image = None
       else:
         # 画像を読み込み、必要ならcropする
-        img, face_cx, face_cy, face_w, face_h = self.load_image_with_face_info(subset, image_info.absolute_path)
+        img, face_cx, face_cy, face_w, face_h = self.load_image_with_face_info(image_info.absolute_path)
         im_h, im_w = img.shape[0:2]
 
         if self.enable_bucket:
-          img = self.trim_and_resize_if_required(subset, img, image_info.bucket_reso, image_info.resized_size)
+          img = self.trim_and_resize_if_required(img, image_info.bucket_reso, image_info.resized_size)
         else:
           if face_cx > 0:                   # 顔位置情報あり
-            img = self.crop_target(subset, img, face_cx, face_cy, face_w, face_h)
+            img = self.crop_target(img, face_cx, face_cy, face_w, face_h)
           elif im_h > self.height or im_w > self.width:
-            assert subset.random_crop, f"image too large, but cropping and bucketing are disabled / 画像サイズが大きいのでface_crop_aug_rangeかrandom_crop、またはbucketを有効にしてください: {image_info.absolute_path}"
+            assert self.random_crop, f"image too large, but cropping and bucketing are disabled / 画像サイズが大きいのでface_crop_aug_rangeかrandom_crop、またはbucketを有効にしてください: {image_info.absolute_path}"
             if im_h > self.height:
               p = random.randint(0, im_h - self.height)
               img = img[p:p + self.height]
@@ -701,9 +637,8 @@ class BaseDataset(torch.utils.data.Dataset):
           assert im_h == self.height and im_w == self.width, f"image size is small / 画像サイズが小さいようです: {image_info.absolute_path}"
 
         # augmentation
-        aug = self.aug_helper.get_augmentor(subset.color_aug, subset.flip_aug)
-        if aug is not None:
-          img = aug(image=img)['image']
+        if self.aug is not None:
+          img = self.aug(image=img)['image']
 
         latents = None
         image = self.image_transforms(img)      # -1.0~1.0のtorch.Tensorになる
@@ -711,7 +646,7 @@ class BaseDataset(torch.utils.data.Dataset):
       images.append(image)
       latents_list.append(latents)
 
-      caption = self.process_caption(subset, image_info.caption)
+      caption = self.process_caption(image_info.caption)
       captions.append(caption)
       if not self.token_padding_disabled:                     # this option might be omitted in future
         input_ids_list.append(self.get_input_ids(caption))
@@ -742,8 +677,9 @@ class BaseDataset(torch.utils.data.Dataset):
 
 
 class DreamBoothDataset(BaseDataset):
-  def __init__(self, subsets: Sequence[DreamBoothSubset], batch_size: int, tokenizer, max_token_length, resolution, enable_bucket: bool, min_bucket_reso: int, max_bucket_reso: int, bucket_reso_steps: int, bucket_no_upscale: bool, prior_loss_weight: float, debug_dataset) -> None:
-    super().__init__(tokenizer, max_token_length, resolution, debug_dataset)
+  def __init__(self, batch_size, train_data_dir, reg_data_dir, tokenizer, max_token_length, caption_extension, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, prior_loss_weight, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset) -> None:
+    super().__init__(tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens,
+                     resolution, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset)
 
     assert resolution is not None, f"resolution is required / resolution（解像度）指定は必須です"
 
@@ -766,7 +702,7 @@ class DreamBoothDataset(BaseDataset):
       self.bucket_reso_steps = None                              # この情報は使われない
       self.bucket_no_upscale = False
 
-    def read_caption(img_path, caption_extension):
+    def read_caption(img_path):
       # captionの候補ファイル名を作る
       base_name = os.path.splitext(img_path)[0]
       base_name_face_det = base_name
@@ -789,181 +725,153 @@ class DreamBoothDataset(BaseDataset):
           break
       return caption
 
-    def load_dreambooth_dir(subset: DreamBoothSubset):
-      if not os.path.isdir(subset.image_dir):
-        print(f"not directory: {subset.image_dir}")
-        return [], []
+    def load_dreambooth_dir(dir):
+      if not os.path.isdir(dir):
+        # print(f"ignore file: {dir}")
+        return 0, [], []
 
-      img_paths = glob_images(subset.image_dir, "*")
-      print(f"found directory {subset.image_dir} contains {len(img_paths)} image files")
+      tokens = os.path.basename(dir).split('_')
+      try:
+        n_repeats = int(tokens[0])
+      except ValueError as e:
+        print(f"ignore directory without repeats / 繰り返し回数のないディレクトリを無視します: {dir}")
+        return 0, [], []
+
+      caption_by_folder = '_'.join(tokens[1:])
+      img_paths = glob_images(dir, "*")
+      print(f"found directory {n_repeats}_{caption_by_folder} contains {len(img_paths)} image files")
 
       # 画像ファイルごとにプロンプトを読み込み、もしあればそちらを使う
       captions = []
       for img_path in img_paths:
-        cap_for_img = read_caption(img_path, subset.caption_extension)
-        if cap_for_img is None and subset.class_tokens is None:
-          print(f"neither caption file nor class tokens are found. use empty caption for {img_path}")
-          captions.append("")
-        else:
-          captions.append(subset.class_tokens if cap_for_img is None else cap_for_img)
+        cap_for_img = read_caption(img_path)
+        captions.append(caption_by_folder if cap_for_img is None else cap_for_img)
 
-      self.set_tag_frequency(os.path.basename(subset.image_dir), captions)         # タグ頻度を記録
+      self.set_tag_frequency(os.path.basename(dir), captions)         # タグ頻度を記録
 
-      return img_paths, captions
+      return n_repeats, img_paths, captions
 
-    print("prepare images.")
+    print("prepare train images.")
+    train_dirs = os.listdir(train_data_dir)
     num_train_images = 0
-    num_reg_images = 0
-    reg_infos: List[ImageInfo] = []
-    for subset in subsets:
-      if subset.num_repeats < 1:
-        print(
-            f"ignore subset with image_dir='{subset.image_dir}': num_repeats is less than 1 / num_repeatsが1を下回っているためサブセットを無視します: {subset.num_repeats}")
-        continue
-
-      if subset in self.subsets:
-        print(
-            f"ignore duplicated subset with image_dir='{subset.image_dir}': use the first one / 既にサブセットが登録されているため、重複した後発のサブセットを無視します")
-        continue
-
-      img_paths, captions = load_dreambooth_dir(subset)
-      if len(img_paths) < 1:
-        print(f"ignore subset with image_dir='{subset.image_dir}': no images found / 画像が見つからないためサブセットを無視します")
-        continue
-
-      if subset.is_reg:
-        num_reg_images += subset.num_repeats * len(img_paths)
-      else:
-        num_train_images += subset.num_repeats * len(img_paths)
+    for dir in train_dirs:
+      n_repeats, img_paths, captions = load_dreambooth_dir(os.path.join(train_data_dir, dir))
+      num_train_images += n_repeats * len(img_paths)
 
       for img_path, caption in zip(img_paths, captions):
-        info = ImageInfo(img_path, subset.num_repeats, caption, subset.is_reg, img_path)
-        if subset.is_reg:
-          reg_infos.append(info)
-        else:
-          self.register_image(info, subset)
+        info = ImageInfo(img_path, n_repeats, caption, False, img_path)
+        self.register_image(info)
 
-      subset.img_count = len(img_paths)
-      self.subsets.append(subset)
+      self.dataset_dirs_info[os.path.basename(dir)] = {"n_repeats": n_repeats, "img_count": len(img_paths)}
 
     print(f"{num_train_images} train images with repeating.")
     self.num_train_images = num_train_images
 
-    print(f"{num_reg_images} reg images.")
-    if num_train_images < num_reg_images:
-      print("some of reg images are not used / 正則化画像の数が多いので、一部使用されない正則化画像があります")
-
-    if num_reg_images == 0:
-      print("no regularization images / 正則化画像が見つかりませんでした")
-    else:
-      # num_repeatsを計算する：どうせ大した数ではないのでループで処理する
-      n = 0
-      first_loop = True
-      while n < num_train_images:
-        for info in reg_infos:
-          if first_loop:
-            self.register_image(info, subset)
-            n += info.num_repeats
-          else:
-            info.num_repeats += 1
-            n += 1
-          if n >= num_train_images:
-            break
-        first_loop = False
-
-    self.num_reg_images = num_reg_images
-
-
-class FineTuningDataset(BaseDataset):
-  def __init__(self, subsets: Sequence[FineTuningSubset], batch_size: int, tokenizer, max_token_length, resolution, enable_bucket: bool, min_bucket_reso: int, max_bucket_reso: int, bucket_reso_steps: int, bucket_no_upscale: bool, debug_dataset) -> None:
-    super().__init__(tokenizer, max_token_length, resolution, debug_dataset)
+    # reg imageは数を数えて学習画像と同じ枚数にする
+    num_reg_images = 0
+    if reg_data_dir:
+      print("prepare reg images.")
+      reg_infos: List[ImageInfo] = []
 
-    self.batch_size = batch_size
+      reg_dirs = os.listdir(reg_data_dir)
+      for dir in reg_dirs:
+        n_repeats, img_paths, captions = load_dreambooth_dir(os.path.join(reg_data_dir, dir))
+        num_reg_images += n_repeats * len(img_paths)
 
-    self.num_train_images = 0
-    self.num_reg_images = 0
+        for img_path, caption in zip(img_paths, captions):
+          info = ImageInfo(img_path, n_repeats, caption, True, img_path)
+          reg_infos.append(info)
 
-    for subset in subsets:
-      if subset.num_repeats < 1:
-        print(
-            f"ignore subset with metadata_file='{subset.metadata_file}': num_repeats is less than 1 / num_repeatsが1を下回っているためサブセットを無視します: {subset.num_repeats}")
-        continue
+        self.reg_dataset_dirs_info[os.path.basename(dir)] = {"n_repeats": n_repeats, "img_count": len(img_paths)}
 
-      if subset in self.subsets:
-        print(
-            f"ignore duplicated subset with metadata_file='{subset.metadata_file}': use the first one / 既にサブセットが登録されているため、重複した後発のサブセットを無視します")
-        continue
+      print(f"{num_reg_images} reg images.")
+      if num_train_images < num_reg_images:
+        print("some of reg images are not used / 正則化画像の数が多いので、一部使用されない正則化画像があります")
 
-      # メタデータを読み込む
-      if os.path.exists(subset.metadata_file):
-        print(f"loading existing metadata: {subset.metadata_file}")
-        with open(subset.metadata_file, "rt", encoding='utf-8') as f:
-          metadata = json.load(f)
+      if num_reg_images == 0:
+        print("no regularization images / 正則化画像が見つかりませんでした")
       else:
-        raise ValueError(f"no metadata / メタデータファイルがありません: {subset.metadata_file}")
+        # num_repeatsを計算する：どうせ大した数ではないのでループで処理する
+        n = 0
+        first_loop = True
+        while n < num_train_images:
+          for info in reg_infos:
+            if first_loop:
+              self.register_image(info)
+              n += info.num_repeats
+            else:
+              info.num_repeats += 1
+              n += 1
+            if n >= num_train_images:
+              break
+          first_loop = False
 
-      if len(metadata) < 1:
-        print(f"ignore subset with '{subset.metadata_file}': no image entries found / 画像に関するデータが見つからないためサブセットを無視します")
-        continue
-
-      tags_list = []
-      for image_key, img_md in metadata.items():
-        # path情報を作る
-        if os.path.exists(image_key):
-          abs_path = image_key
-        else:
-          npz_path = os.path.join(subset.image_dir, image_key + ".npz")
-          if os.path.exists(npz_path):
-            abs_path = npz_path
-          else:
-            # わりといい加減だがいい方法が思いつかん
-            abs_path = glob_images(subset.image_dir, image_key)
-            assert len(abs_path) >= 1, f"no image / 画像がありません: {image_key}"
-            abs_path = abs_path[0]
-
-        caption = img_md.get('caption')
-        tags = img_md.get('tags')
-        if caption is None:
-          caption = tags
-        elif tags is not None and len(tags) > 0:
-          caption = caption + ', ' + tags
-          tags_list.append(tags)
-
-        if caption is None:
-          caption = ""
+    self.num_reg_images = num_reg_images
 
-        image_info = ImageInfo(image_key, subset.num_repeats, caption, False, abs_path)
-        image_info.image_size = img_md.get('train_resolution')
 
-        if not subset.color_aug and not subset.random_crop:
-          # if npz exists, use them
-          image_info.latents_npz, image_info.latents_npz_flipped = self.image_key_to_npz_file(subset, image_key)
+class FineTuningDataset(BaseDataset):
+  def __init__(self, json_file_name, batch_size, train_data_dir, tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, flip_aug, color_aug, face_crop_aug_range, random_crop, dataset_repeats, debug_dataset) -> None:
+    super().__init__(tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens,
+                     resolution, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset)
+
+    # メタデータを読み込む
+    if os.path.exists(json_file_name):
+      print(f"loading existing metadata: {json_file_name}")
+      with open(json_file_name, "rt", encoding='utf-8') as f:
+        metadata = json.load(f)
+    else:
+      raise ValueError(f"no metadata / メタデータファイルがありません: {json_file_name}")
 
-        self.register_image(image_info, subset)
+    self.metadata = metadata
+    self.train_data_dir = train_data_dir
+    self.batch_size = batch_size
 
-      self.num_train_images += len(metadata) * subset.num_repeats
+    tags_list = []
+    for image_key, img_md in metadata.items():
+      # path情報を作る
+      if os.path.exists(image_key):
+        abs_path = image_key
+      else:
+        # わりといい加減だがいい方法が思いつかん
+        abs_path = glob_images(train_data_dir, image_key)
+        assert len(abs_path) >= 1, f"no image / 画像がありません: {image_key}"
+        abs_path = abs_path[0]
+
+      caption = img_md.get('caption')
+      tags = img_md.get('tags')
+      if caption is None:
+        caption = tags
+      elif tags is not None and len(tags) > 0:
+        caption = caption + ', ' + tags
+        tags_list.append(tags)
+      assert caption is not None and len(caption) > 0, f"caption or tag is required / キャプションまたはタグは必須です:{abs_path}"
+
+      image_info = ImageInfo(image_key, dataset_repeats, caption, False, abs_path)
+      image_info.image_size = img_md.get('train_resolution')
+
+      if not self.color_aug and not self.random_crop:
+        # if npz exists, use them
+        image_info.latents_npz, image_info.latents_npz_flipped = self.image_key_to_npz_file(image_key)
+
+      self.register_image(image_info)
+    self.num_train_images = len(metadata) * dataset_repeats
+    self.num_reg_images = 0
 
-      # TODO do not record tag freq when no tag
-      self.set_tag_frequency(os.path.basename(subset.metadata_file), tags_list)
-      subset.img_count = len(metadata)
-      self.subsets.append(subset)
+    # TODO do not record tag freq when no tag
+    self.set_tag_frequency(os.path.basename(json_file_name), tags_list)
+    self.dataset_dirs_info[os.path.basename(json_file_name)] = {"n_repeats": dataset_repeats, "img_count": len(metadata)}
 
     # check existence of all npz files
-    use_npz_latents = all([not (subset.color_aug or subset.random_crop) for subset in self.subsets])
+    use_npz_latents = not (self.color_aug or self.random_crop)
     if use_npz_latents:
-      flip_aug_in_subset = False
       npz_any = False
       npz_all = True
-
       for image_info in self.image_data.values():
-        subset = self.image_to_subset[image_info.image_key]
-
         has_npz = image_info.latents_npz is not None
         npz_any = npz_any or has_npz
 
-        if subset.flip_aug:
+        if self.flip_aug:
           has_npz = has_npz and image_info.latents_npz_flipped is not None
-          flip_aug_in_subset = True
         npz_all = npz_all and has_npz
 
         if npz_any and not npz_all:
@@ -975,7 +883,7 @@ class FineTuningDataset(BaseDataset):
       elif not npz_all:
         use_npz_latents = False
         print(f"some of npz file does not exist. ignore npz files / いくつ���のnpzファイルが見つからないためnpzファイルを無視します")
-        if flip_aug_in_subset:
+        if self.flip_aug:
           print("maybe no flipped files / 反転されたnpzファイルがないのかもしれません")
     # else:
     #   print("npz files are not used with color_aug and/or random_crop / color_augまたはrandom_cropが指定されているためnpzファイルは使用されません")
@@ -1021,7 +929,7 @@ class FineTuningDataset(BaseDataset):
       for image_info in self.image_data.values():
         image_info.latents_npz = image_info.latents_npz_flipped = None
 
-  def image_key_to_npz_file(self, subset: FineTuningSubset, image_key):
+  def image_key_to_npz_file(self, image_key):
     base_name = os.path.splitext(image_key)[0]
     npz_file_norm = base_name + '.npz'
 
@@ -1033,8 +941,8 @@ class FineTuningDataset(BaseDataset):
       return npz_file_norm, npz_file_flip
 
     # image_key is relative path
-    npz_file_norm = os.path.join(subset.image_dir, image_key + '.npz')
-    npz_file_flip = os.path.join(subset.image_dir, image_key + '_flip.npz')
+    npz_file_norm = os.path.join(self.train_data_dir, image_key + '.npz')
+    npz_file_flip = os.path.join(self.train_data_dir, image_key + '_flip.npz')
 
     if not os.path.exists(npz_file_norm):
       npz_file_norm = None
@@ -1045,60 +953,13 @@ class FineTuningDataset(BaseDataset):
     return npz_file_norm, npz_file_flip
 
 
-# behave as Dataset mock
-class DatasetGroup(torch.utils.data.ConcatDataset):
-  def __init__(self, datasets: Sequence[Union[DreamBoothDataset, FineTuningDataset]]):
-    self.datasets: List[Union[DreamBoothDataset, FineTuningDataset]]
-
-    super().__init__(datasets)
-
-    self.image_data = {}
-    self.num_train_images = 0
-    self.num_reg_images = 0
-
-    # simply concat together
-    # TODO: handling image_data key duplication among dataset
-    #   In practical, this is not the big issue because image_data is accessed from outside of dataset only for debug_dataset.
-    for dataset in datasets:
-      self.image_data.update(dataset.image_data)
-      self.num_train_images += dataset.num_train_images
-      self.num_reg_images += dataset.num_reg_images
-
-  def add_replacement(self, str_from, str_to):
-    for dataset in self.datasets:
-      dataset.add_replacement(str_from, str_to)
-
-  # def make_buckets(self):
-  #   for dataset in self.datasets:
-  #     dataset.make_buckets()
-
-  def cache_latents(self, vae):
-    for i, dataset in enumerate(self.datasets):
-      print(f"[Dataset {i}]")
-      dataset.cache_latents(vae)
-
-  def is_latent_cacheable(self) -> bool:
-    return all([dataset.is_latent_cacheable() for dataset in self.datasets])
-
-  def set_current_epoch(self, epoch):
-    for dataset in self.datasets:
-      dataset.set_current_epoch(epoch)
-
-  def disable_token_padding(self):
-    for dataset in self.datasets:
-      dataset.disable_token_padding()
-
-
 def debug_dataset(train_dataset, show_input_ids=False):
   print(f"Total dataset length (steps) / データセットの長さ（ステップ数）: {len(train_dataset)}")
   print("Escape for exit. / Escキーで中断、終了します")
 
   train_dataset.set_current_epoch(1)
   k = 0
-  indices = list(range(len(train_dataset)))
-  random.shuffle(indices)
-  for i, idx in enumerate(indices):
-    example = train_dataset[idx]
+  for i, example in enumerate(train_dataset):
     if example['latents'] is not None:
       print(f"sample has latents from npz file: {example['latents'].size()}")
     for j, (ik, cap, lw, iid) in enumerate(zip(example['image_keys'], example['captions'], example['loss_weights'], example['input_ids'])):
@@ -1503,35 +1364,6 @@ def add_sd_models_arguments(parser: argparse.ArgumentParser):
                       help='enable v-parameterization training / v-parameterization学習を有効にする')
   parser.add_argument("--pretrained_model_name_or_path", type=str, default=None,
                       help="pretrained model to train, directory to Diffusers model or StableDiffusion checkpoint / 学習元モデル、Diffusers形式モデルのディレクトリまたはStableDiffusionのckptファイル")
-  parser.add_argument("--tokenizer_cache_dir", type=str, default=None,
-                      help="directory for caching Tokenizer (for offline training) / Tokenizerをキャッシュするディレクトリ（ネット接続なしでの学習のため）")
-
-
-def add_optimizer_arguments(parser: argparse.ArgumentParser):
-  parser.add_argument("--optimizer_type", type=str, default="",
-                      help="Optimizer to use / オプティマイザの種類: AdamW (default), AdamW8bit, Lion, SGDNesterov, SGDNesterov8bit, DAdaptation, AdaFactor")
-
-  # backward compatibility
-  parser.add_argument("--use_8bit_adam", action="store_true",
-                      help="use 8bit AdamW optimizer (requires bitsandbytes) / 8bit Adamオプティマイザを使う（bitsandbytesのインストールが必要）")
-  parser.add_argument("--use_lion_optimizer", action="store_true",
-                      help="use Lion optimizer (requires lion-pytorch) / Lionオプティマイザを使う（ lion-pytorch のインストールが必要）")
-
-  parser.add_argument("--learning_rate", type=float, default=2.0e-6, help="learning rate / 学習率")
-  parser.add_argument("--max_grad_norm", default=1.0, type=float,
-                      help="Max gradient norm, 0 for no clipping / 勾配正規化の最大norm、0でclippingを行わない")
-
-  parser.add_argument("--optimizer_args", type=str, default=None, nargs='*',
-                      help="additional arguments for optimizer (like \"weight_decay=0.01 betas=0.9,0.999 ...\") / オプティマイザの追加引数（例： \"weight_decay=0.01 betas=0.9,0.999 ...\"）")
-
-  parser.add_argument("--lr_scheduler", type=str, default="constant",
-                      help="scheduler to use for learning rate / 学習率のスケジューラ: linear, cosine, cosine_with_restarts, polynomial, constant (default), constant_with_warmup, adafactor")
-  parser.add_argument("--lr_warmup_steps", type=int, default=0,
-                      help="Number of steps for the warmup in the lr scheduler (default is 0) / 学習率のスケジューラをウォームアップするステップ数（デフォルト0）")
-  parser.add_argument("--lr_scheduler_num_cycles", type=int, default=1,
-                      help="Number of restarts for cosine scheduler with restarts / cosine with restartsスケジューラでのリスタート回数")
-  parser.add_argument("--lr_scheduler_power", type=float, default=1,
-                      help="Polynomial power for polynomial scheduler / polynomialスケジューラでのpolynomial power")
 
 
 def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth: bool):
@@ -1555,6 +1387,10 @@ def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth:
   parser.add_argument("--train_batch_size", type=int, default=1, help="batch size for training / 学習時のバッチサイズ")
   parser.add_argument("--max_token_length", type=int, default=None, choices=[None, 150, 225],
                       help="max token length of text encoder (default for 75, 150 or 225) / text encoderのトークンの最大長（未指定で75、150または225が指定可）")
+  parser.add_argument("--use_8bit_adam", action="store_true",
+                      help="use 8bit Adam optimizer (requires bitsandbytes) / 8bit Adamオプティマイザを使う（bitsandbytesのインストールが必要）")
+  parser.add_argument("--use_lion_optimizer", action="store_true",
+                      help="use Lion optimizer (requires lion-pytorch) / Lionオプティマイザを使う（ lion-pytorch のインストールが必要）")
   parser.add_argument("--mem_eff_attn", action="store_true",
                       help="use memory efficient attention for CrossAttention / CrossAttentionに省メモリ版attentionを使う")
   parser.add_argument("--xformers", action="store_true",
@@ -1562,6 +1398,7 @@ def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth:
   parser.add_argument("--vae", type=str, default=None,
                       help="path to checkpoint of vae to replace / VAEを入れ替える場合、VAEのcheckpointファイルまたはディレクトリ")
 
+  parser.add_argument("--learning_rate", type=float, default=2.0e-6, help="learning rate / 学習率")
   parser.add_argument("--max_train_steps", type=int, default=1600, help="training steps / 学習ステップ数")
   parser.add_argument("--max_train_epochs", type=int, default=None,
                       help="training epochs (overrides max_train_steps) / 学習エポック数（max_train_stepsを上書きします）")
@@ -1582,23 +1419,15 @@ def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth:
   parser.add_argument("--logging_dir", type=str, default=None,
                       help="enable logging and output TensorBoard log to this directory / ログ出力を有効にしてこのディレクトリにTensorBoard用のログを出力する")
   parser.add_argument("--log_prefix", type=str, default=None, help="add prefix for each log directory / ログディレクトリ名の先頭に追加する文字列")
+  parser.add_argument("--lr_scheduler", type=str, default="constant",
+                      help="scheduler to use for learning rate / 学習率のスケジューラ: linear, cosine, cosine_with_restarts, polynomial, constant (default), constant_with_warmup")
+  parser.add_argument("--lr_warmup_steps", type=int, default=0,
+                      help="Number of steps for the warmup in the lr scheduler (default is 0) / 学習率のスケジューラをウォームアップするステップ数（デフォルト0）")
   parser.add_argument("--noise_offset", type=float, default=None,
                       help="enable noise offset with this value (if enabled, around 0.1 is recommended) / Noise offsetを有効にしてこの値を設定する（有効にする場合は0.1程度を推奨）")
   parser.add_argument("--lowram", action="store_true",
                       help="enable low RAM optimization. e.g. load models to VRAM instead of RAM (for machines which have bigger VRAM than RAM such as Colab and Kaggle) / メインメモリが少ない環境向け最適化を有効にする。たとえばVRAMにモデルを読み込むなど（ColabやKaggleなどRAMに比べてVRAMが多い環境向け）")
 
-  parser.add_argument("--sample_every_n_steps", type=int, default=None,
-                      help="generate sample images every N steps / 学習中のモデルで指定ステップごとにサンプル出力する")
-  parser.add_argument("--sample_every_n_epochs", type=int, default=None,
-                      help="generate sample images every N epochs (overwrites n_steps) / 学習中のモデルで指定エポックごとにサンプル出力する（ステップ数指定を上書きします）")
-  parser.add_argument("--sample_prompts", type=str, default=None,
-                      help="file for prompts to generate sample images / 学習中モデルのサンプル出力用プロンプトのファイル")
-  parser.add_argument('--sample_sampler', type=str, default='ddim',
-                      choices=['ddim', 'pndm', 'lms', 'euler', 'euler_a', 'heun', 'dpm_2', 'dpm_2_a', 'dpmsolver',
-                               'dpmsolver++', 'dpmsingle',
-                               'k_lms', 'k_euler', 'k_euler_a', 'k_dpm_2', 'k_dpm_2_a'],
-                      help=f'sampler (scheduler) type for sample images / サンプル出力時のサンプラー（スケジューラ）の種類')
-
   if support_dreambooth:
     # DreamBooth training
     parser.add_argument("--prior_loss_weight", type=float, default=1.0,
@@ -1620,8 +1449,8 @@ def add_dataset_arguments(parser: argparse.ArgumentParser, support_dreambooth: b
   parser.add_argument("--caption_extension", type=str, default=".caption", help="extension of caption files / 読み込むcaptionファイルの拡張子")
   parser.add_argument("--caption_extention", type=str, default=None,
                       help="extension of caption files (backward compatibility) / 読み込むcaptionファイルの拡張子（スペルミスを残してあります）")
-  parser.add_argument("--keep_tokens", type=int, default=0,
-                      help="keep heading N tokens when shuffling caption tokens (token means comma separated strings) / captionのシャッフル時に、先頭からこの個数のトークンをシャッフルしないで残す（トークンはカンマ区切りの各部分を意味する）")
+  parser.add_argument("--keep_tokens", type=int, default=None,
+                      help="keep heading N tokens when shuffling caption tokens / captionのシャッフル時に、先頭からこの個数のトークンをシャッフルしないで残す")
   parser.add_argument("--color_aug", action="store_true", help="enable weak color augmentation / 学習時に色合いのaugmentationを有効にする")
   parser.add_argument("--flip_aug", action="store_true", help="enable horizontal flip augmentation / 学習時に左右反転のaugmentationを有効にする")
   parser.add_argument("--face_crop_aug_range", type=str, default=None,
@@ -1646,11 +1475,11 @@ def add_dataset_arguments(parser: argparse.ArgumentParser, support_dreambooth: b
   if support_caption_dropout:
     # Textual Inversion はcaptionのdropoutをsupportしない
     # いわゆるtensorのDropoutと紛らわしいのでprefixにcaptionを付けておく　every_n_epochsは他と平仄を合わせてdefault Noneに
-    parser.add_argument("--caption_dropout_rate", type=float, default=0.0,
+    parser.add_argument("--caption_dropout_rate", type=float, default=0,
                         help="Rate out dropout caption(0.0~1.0) / captionをdropoutする割合")
-    parser.add_argument("--caption_dropout_every_n_epochs", type=int, default=0,
+    parser.add_argument("--caption_dropout_every_n_epochs", type=int, default=None,
                         help="Dropout all captions every N epochs / captionを指定エポックごとにdropoutする")
-    parser.add_argument("--caption_tag_dropout_rate", type=float, default=0.0,
+    parser.add_argument("--caption_tag_dropout_rate", type=float, default=0,
                         help="Rate out dropout comma separated tokens(0.0~1.0) / カンマ区切りのタグをdropoutする割合")
 
   if support_dreambooth:
@@ -1675,256 +1504,16 @@ def add_sd_saving_arguments(parser: argparse.ArgumentParser):
 # region utils
 
 
-def get_optimizer(args, trainable_params):
-  # "Optimizer to use: AdamW, AdamW8bit, Lion, SGDNesterov, SGDNesterov8bit, DAdaptation, Adafactor"
-
-  optimizer_type = args.optimizer_type
-  if args.use_8bit_adam:
-    assert not args.use_lion_optimizer, "both option use_8bit_adam and use_lion_optimizer are specified / use_8bit_adamとuse_lion_optimizerの両方のオプションが指定されています"
-    assert optimizer_type is None or optimizer_type == "", "both option use_8bit_adam and optimizer_type are specified / use_8bit_adamとoptimizer_typeの両方のオプションが指定されています"
-    optimizer_type = "AdamW8bit"
-
-  elif args.use_lion_optimizer:
-    assert optimizer_type is None or optimizer_type == "", "both option use_lion_optimizer and optimizer_type are specified / use_lion_optimizerとoptimizer_typeの両方のオプションが指定されています"
-    optimizer_type = "Lion"
-
-  if optimizer_type is None or optimizer_type == "":
-    optimizer_type = "AdamW"
-  optimizer_type = optimizer_type.lower()
-
-  # 引数を分解する：boolとfloat、tupleのみ対応
-  optimizer_kwargs = {}
-  if args.optimizer_args is not None and len(args.optimizer_args) > 0:
-    for arg in args.optimizer_args:
-      key, value = arg.split('=')
-
-      value = value.split(",")
-      for i in range(len(value)):
-        if value[i].lower() == "true" or value[i].lower() == "false":
-          value[i] = (value[i].lower() == "true")
-        else:
-          value[i] = float(value[i])
-      if len(value) == 1:
-        value = value[0]
-      else:
-        value = tuple(value)
-
-      optimizer_kwargs[key] = value
-  # print("optkwargs:", optimizer_kwargs)
-
-  lr = args.learning_rate
-
-  if optimizer_type == "AdamW8bit".lower():
-    try:
-      import bitsandbytes as bnb
-    except ImportError:
-      raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
-    print(f"use 8-bit AdamW optimizer | {optimizer_kwargs}")
-    optimizer_class = bnb.optim.AdamW8bit
-    optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs)
-
-  elif optimizer_type == "SGDNesterov8bit".lower():
-    try:
-      import bitsandbytes as bnb
-    except ImportError:
-      raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
-    print(f"use 8-bit SGD with Nesterov optimizer | {optimizer_kwargs}")
-    if "momentum" not in optimizer_kwargs:
-      print(f"8-bit SGD with Nesterov must be with momentum, set momentum to 0.9 / 8-bit SGD with Nesterovはmomentum指定が必須のため0.9に設定します")
-      optimizer_kwargs["momentum"] = 0.9
-
-    optimizer_class = bnb.optim.SGD8bit
-    optimizer = optimizer_class(trainable_params, lr=lr, nesterov=True, **optimizer_kwargs)
-
-  elif optimizer_type == "Lion".lower():
-    try:
-      import lion_pytorch
-    except ImportError:
-      raise ImportError("No lion_pytorch / lion_pytorch がインストールされていないようです")
-    print(f"use Lion optimizer | {optimizer_kwargs}")
-    optimizer_class = lion_pytorch.Lion
-    optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs)
-
-  elif optimizer_type == "SGDNesterov".lower():
-    print(f"use SGD with Nesterov optimizer | {optimizer_kwargs}")
-    if "momentum" not in optimizer_kwargs:
-      print(f"SGD with Nesterov must be with momentum, set momentum to 0.9 / SGD with Nesterovはmomentum指定が必須のため0.9に設定します")
-      optimizer_kwargs["momentum"] = 0.9
-
-    optimizer_class = torch.optim.SGD
-    optimizer = optimizer_class(trainable_params, lr=lr,  nesterov=True, **optimizer_kwargs)
-
-  elif optimizer_type == "DAdaptation".lower():
-    try:
-      import dadaptation
-    except ImportError:
-      raise ImportError("No dadaptation / dadaptation がインストールされていないようです")
-    print(f"use D-Adaptation Adam optimizer | {optimizer_kwargs}")
-
-    actual_lr = lr
-    lr_count = 1
-    if type(trainable_params) == list and type(trainable_params[0]) == dict:
-      lrs = set()
-      actual_lr = trainable_params[0].get("lr", actual_lr)
-      for group in trainable_params:
-        lrs.add(group.get("lr", actual_lr))
-      lr_count = len(lrs)
-
-    if actual_lr <= 0.1:
-      print(
-          f'learning rate is too low. If using dadaptation, set learning rate around 1.0 / 学習率が低すぎるようです。1.0前後の値を指定してください: lr={actual_lr}')
-      print('recommend option: lr=1.0 / 推奨は1.0です')
-    if lr_count > 1:
-      print(
-          f"when multiple learning rates are specified with dadaptation (e.g. for Text Encoder and U-Net), only the first one will take effect / D-Adaptationで複数の学習率を指定した場合（Text EncoderとU-Netなど）、最初の学習率のみが有効になります: lr={actual_lr}")
-
-    optimizer_class = dadaptation.DAdaptAdam
-    optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs)
-
-  elif optimizer_type == "Adafactor".lower():
-    # 引数を確認して適宜補正する
-    if "relative_step" not in optimizer_kwargs:
-      optimizer_kwargs["relative_step"] = True                  # default
-    if not optimizer_kwargs["relative_step"] and optimizer_kwargs.get("warmup_init", False):
-      print(f"set relative_step to True because warmup_init is True / warmup_initがTrueのためrelative_stepをTrueにします")
-      optimizer_kwargs["relative_step"] = True
-    print(f"use Adafactor optimizer | {optimizer_kwargs}")
-
-    if optimizer_kwargs["relative_step"]:
-      print(f"relative_step is true / relative_stepがtrueです")
-      if lr != 0.0:
-        print(f"learning rate is used as initial_lr / 指定したlearning rateはinitial_lrとして使用されます")
-      args.learning_rate = None
-
-      # trainable_paramsがgroupだった時の処理：lrを削除する
-      if type(trainable_params) == list and type(trainable_params[0]) == dict:
-        has_group_lr = False
-        for group in trainable_params:
-          p = group.pop("lr", None)
-          has_group_lr = has_group_lr or (p is not None)
-
-        if has_group_lr:
-          # 一応argsを無効にしておく TODO 依存関係が逆転してるのであまり望ましくない
-          print(f"unet_lr and text_encoder_lr are ignored / unet_lrとtext_encoder_lrは無視されます")
-          args.unet_lr = None
-          args.text_encoder_lr = None
-
-      if args.lr_scheduler != "adafactor":
-        print(f"use adafactor_scheduler / スケジューラにadafactor_schedulerを使用します")
-      args.lr_scheduler = f"adafactor:{lr}"                               # ちょっと微妙だけど
-
-      lr = None
-    else:
-      if args.max_grad_norm != 0.0:
-        print(f"because max_grad_norm is set, clip_grad_norm is enabled. consider set to 0 / max_grad_normが設定されているためclip_grad_normが有効になります。0に設定して無効にしたほうがいいかもしれません")
-      if args.lr_scheduler != "constant_with_warmup":
-        print(f"constant_with_warmup will be good / スケジューラはconstant_with_warmupが良いかもしれません")
-      if optimizer_kwargs.get("clip_threshold", 1.0) != 1.0:
-        print(f"clip_threshold=1.0 will be good / clip_thresholdは1.0が良いかもしれません")
-
-    optimizer_class = transformers.optimization.Adafactor
-    optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs)
-
-  elif optimizer_type == "AdamW".lower():
-    print(f"use AdamW optimizer | {optimizer_kwargs}")
-    optimizer_class = torch.optim.AdamW
-    optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs)
-
-  else:
-    # 任意のoptimizerを使う
-    optimizer_type = args.optimizer_type   # lowerでないやつ（微妙）
-    print(f"use {optimizer_type} | {optimizer_kwargs}")
-    if "." not in optimizer_type:
-      optimizer_module = torch.optim
-    else:
-      values = optimizer_type.split(".")
-      optimizer_module = importlib.import_module(".".join(values[:-1]))
-      optimizer_type = values[-1]
-
-    optimizer_class = getattr(optimizer_module, optimizer_type)
-    optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs)
-
-  optimizer_name = optimizer_class.__module__ + "." + optimizer_class.__name__
-  optimizer_args = ",".join([f"{k}={v}" for k, v in optimizer_kwargs.items()])
-
-  return optimizer_name, optimizer_args, optimizer
-
-
-# Monkeypatch newer get_scheduler() function overridng current version of diffusers.optimizer.get_scheduler
-# code is taken from https://github.com/huggingface/diffusers diffusers.optimizer, commit d87cc15977b87160c30abaace3894e802ad9e1e6
-# Which is a newer release of diffusers than currently packaged with sd-scripts
-# This code can be removed when newer diffusers version (v0.12.1 or greater) is tested and implemented to sd-scripts
-
-
-def get_scheduler_fix(
-    name: Union[str, SchedulerType],
-    optimizer: Optimizer,
-    num_warmup_steps: Optional[int] = None,
-    num_training_steps: Optional[int] = None,
-    num_cycles: int = 1,
-    power: float = 1.0,
-):
-  """
-  Unified API to get any scheduler from its name.
-  Args:
-      name (`str` or `SchedulerType`):
-          The name of the scheduler to use.
-      optimizer (`torch.optim.Optimizer`):
-          The optimizer that will be used during training.
-      num_warmup_steps (`int`, *optional*):
-          The number of warmup steps to do. This is not required by all schedulers (hence the argument being
-          optional), the function will raise an error if it's unset and the scheduler type requires it.
-      num_training_steps (`int``, *optional*):
-          The number of training steps to do. This is not required by all schedulers (hence the argument being
-          optional), the function will raise an error if it's unset and the scheduler type requires it.
-      num_cycles (`int`, *optional*):
-          The number of hard restarts used in `COSINE_WITH_RESTARTS` scheduler.
-      power (`float`, *optional*, defaults to 1.0):
-          Power factor. See `POLYNOMIAL` scheduler
-      last_epoch (`int`, *optional*, defaults to -1):
-          The index of the last epoch when resuming training.
-  """
-  if name.startswith("adafactor"):
-    assert type(optimizer) == transformers.optimization.Adafactor, f"adafactor scheduler must be used with Adafactor optimizer / adafactor schedulerはAdafactorオプティマイザと同時に使ってください"
-    initial_lr = float(name.split(':')[1])
-    # print("adafactor scheduler init lr", initial_lr)
-    return transformers.optimization.AdafactorSchedule(optimizer, initial_lr)
-
-  name = SchedulerType(name)
-  schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]
-  if name == SchedulerType.CONSTANT:
-    return schedule_func(optimizer)
-
-  # All other schedulers require `num_warmup_steps`
-  if num_warmup_steps is None:
-    raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")
-
-  if name == SchedulerType.CONSTANT_WITH_WARMUP:
-    return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)
-
-  # All other schedulers require `num_training_steps`
-  if num_training_steps is None:
-    raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")
-
-  if name == SchedulerType.COSINE_WITH_RESTARTS:
-    return schedule_func(
-        optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, num_cycles=num_cycles
-    )
-
-  if name == SchedulerType.POLYNOMIAL:
-    return schedule_func(
-        optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, power=power
-    )
-
-  return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps)
-
-
 def prepare_dataset_args(args: argparse.Namespace, support_metadata: bool):
   # backward compatibility
   if args.caption_extention is not None:
     args.caption_extension = args.caption_extention
     args.caption_extention = None
 
+  if args.cache_latents:
+    assert not args.color_aug, "when caching latents, color_aug cannot be used / latentをキャッシュするときはcolor_augは使えません"
+    assert not args.random_crop, "when caching latents, random_crop cannot be used / latentをキャッシュするときはrandom_cropは使えません"
+
   # assert args.resolution is not None, f"resolution is required / resolution（解像度）を指定してください"
   if args.resolution is not None:
     args.resolution = tuple([int(r) for r in args.resolution.split(',')])
@@ -1947,28 +1536,12 @@ def prepare_dataset_args(args: argparse.Namespace, support_metadata: bool):
 
 def load_tokenizer(args: argparse.Namespace):
   print("prepare tokenizer")
-  original_path = V2_STABLE_DIFFUSION_PATH if args.v2 else TOKENIZER_PATH
-
-  tokenizer: CLIPTokenizer = None
-  if args.tokenizer_cache_dir:
-    local_tokenizer_path = os.path.join(args.tokenizer_cache_dir, original_path.replace('/', '_'))
-    if os.path.exists(local_tokenizer_path):
-      print(f"load tokenizer from cache: {local_tokenizer_path}")
-      tokenizer = CLIPTokenizer.from_pretrained(local_tokenizer_path)                   # same for v1 and v2
-
-  if tokenizer is None:
-    if args.v2:
-      tokenizer = CLIPTokenizer.from_pretrained(original_path, subfolder="tokenizer")
-    else:
-      tokenizer = CLIPTokenizer.from_pretrained(original_path)
-
-  if hasattr(args, "max_token_length") and args.max_token_length is not None:
+  if args.v2:
+    tokenizer = CLIPTokenizer.from_pretrained(V2_STABLE_DIFFUSION_PATH, subfolder="tokenizer")
+  else:
+    tokenizer = CLIPTokenizer.from_pretrained(TOKENIZER_PATH)
+  if args.max_token_length is not None:
     print(f"update token length: {args.max_token_length}")
-
-  if args.tokenizer_cache_dir and not os.path.exists(local_tokenizer_path):
-    print(f"save Tokenizer to cache: {local_tokenizer_path}")
-    tokenizer.save_pretrained(local_tokenizer_path)
-
   return tokenizer
 
 
@@ -2019,19 +1592,13 @@ def prepare_dtype(args: argparse.Namespace):
 
 
 def load_target_model(args: argparse.Namespace, weight_dtype):
-  name_or_path = args.pretrained_model_name_or_path
-  name_or_path = os.readlink(name_or_path) if os.path.islink(name_or_path) else name_or_path
-  load_stable_diffusion_format = os.path.isfile(name_or_path)           # determine SD or Diffusers
+  load_stable_diffusion_format = os.path.isfile(args.pretrained_model_name_or_path)           # determine SD or Diffusers
   if load_stable_diffusion_format:
     print("load StableDiffusion checkpoint")
-    text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, name_or_path)
+    text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.pretrained_model_name_or_path)
   else:
     print("load Diffusers pretrained models")
-    try:
-      pipe = StableDiffusionPipeline.from_pretrained(name_or_path, tokenizer=None, safety_checker=None)
-    except EnvironmentError as ex:
-      print(
-          f"model is not found as a file or in Hugging Face, perhaps file name is wrong? / 指定したモデル名のファイル、またはHugging Faceのモデルが見つかりません。ファイル名が誤っているかもしれません: {name_or_path}")
+    pipe = StableDiffusionPipeline.from_pretrained(args.pretrained_model_name_or_path, tokenizer=None, safety_checker=None)
     text_encoder = pipe.text_encoder
     vae = pipe.vae
     unet = pipe.unet
@@ -2200,197 +1767,6 @@ def save_state_on_train_end(args: argparse.Namespace, accelerator):
   model_name = DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name
   accelerator.save_state(os.path.join(args.output_dir, LAST_STATE_NAME.format(model_name)))
 
-
-# scheduler:
-SCHEDULER_LINEAR_START = 0.00085
-SCHEDULER_LINEAR_END = 0.0120
-SCHEDULER_TIMESTEPS = 1000
-SCHEDLER_SCHEDULE = 'scaled_linear'
-
-
-def sample_images(accelerator, args: argparse.Namespace, epoch, steps, device, vae, tokenizer, text_encoder, unet, prompt_replacement=None):
-  """
-  生成に使っている Diffusers の Pipeline がデフォルトなので、プロンプトの重みづけには対応していない
-  clip skipは対応した
-  """
-  if args.sample_every_n_steps is None and args.sample_every_n_epochs is None:
-    return
-  if args.sample_every_n_epochs is not None:
-    # sample_every_n_steps は無視する
-    if epoch is None or epoch % args.sample_every_n_epochs != 0:
-      return
-  else:
-    if steps % args.sample_every_n_steps != 0 or epoch is not None:       # steps is not divisible or end of epoch
-      return
-
-  print(f"generating sample images at step / サンプル画像生成 ステップ: {steps}")
-  if not os.path.isfile(args.sample_prompts):
-    print(f"No prompt file / プロンプトファイルがありません: {args.sample_prompts}")
-    return
-
-  org_vae_device = vae.device                           # CPUにいるはず
-  vae.to(device)
-
-  # clip skip 対応のための wrapper を作る
-  if args.clip_skip is None:
-    text_encoder_or_wrapper = text_encoder
-  else:
-    class Wrapper():
-      def __init__(self, tenc) -> None:
-        self.tenc = tenc
-        self.config = {}
-        super().__init__()
-
-      def __call__(self, input_ids, attention_mask):
-        enc_out = self.tenc(input_ids, output_hidden_states=True, return_dict=True)
-        encoder_hidden_states = enc_out['hidden_states'][-args.clip_skip]
-        encoder_hidden_states = self.tenc.text_model.final_layer_norm(encoder_hidden_states)
-        pooled_output = enc_out['pooler_output']
-        return encoder_hidden_states, pooled_output  # 1st output is only used
-
-    text_encoder_or_wrapper = Wrapper(text_encoder)
-
-  # read prompts
-  with open(args.sample_prompts, 'rt', encoding='utf-8') as f:
-    prompts = f.readlines()
-
-  # schedulerを用意する
-  sched_init_args = {}
-  if args.sample_sampler == "ddim":
-    scheduler_cls = DDIMScheduler
-  elif args.sample_sampler == "ddpm":                    # ddpmはおかしくなるのでoptionから外してある
-    scheduler_cls = DDPMScheduler
-  elif args.sample_sampler == "pndm":
-    scheduler_cls = PNDMScheduler
-  elif args.sample_sampler == 'lms' or args.sample_sampler == 'k_lms':
-    scheduler_cls = LMSDiscreteScheduler
-  elif args.sample_sampler == 'euler' or args.sample_sampler == 'k_euler':
-    scheduler_cls = EulerDiscreteScheduler
-  elif args.sample_sampler == 'euler_a' or args.sample_sampler == 'k_euler_a':
-    scheduler_cls = EulerAncestralDiscreteScheduler
-  elif args.sample_sampler == "dpmsolver" or args.sample_sampler == "dpmsolver++":
-    scheduler_cls = DPMSolverMultistepScheduler
-    sched_init_args['algorithm_type'] = args.sample_sampler
-  elif args.sample_sampler == "dpmsingle":
-    scheduler_cls = DPMSolverSinglestepScheduler
-  elif args.sample_sampler == "heun":
-    scheduler_cls = HeunDiscreteScheduler
-  elif args.sample_sampler == 'dpm_2' or args.sample_sampler == 'k_dpm_2':
-    scheduler_cls = KDPM2DiscreteScheduler
-  elif args.sample_sampler == 'dpm_2_a' or args.sample_sampler == 'k_dpm_2_a':
-    scheduler_cls = KDPM2AncestralDiscreteScheduler
-  else:
-    scheduler_cls = DDIMScheduler
-
-  if args.v_parameterization:
-    sched_init_args['prediction_type'] = 'v_prediction'
-
-  scheduler = scheduler_cls(num_train_timesteps=SCHEDULER_TIMESTEPS,
-                            beta_start=SCHEDULER_LINEAR_START, beta_end=SCHEDULER_LINEAR_END,
-                            beta_schedule=SCHEDLER_SCHEDULE, **sched_init_args)
-
-  # clip_sample=Trueにする
-  if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is False:
-    # print("set clip_sample to True")
-    scheduler.config.clip_sample = True
-
-  pipeline = StableDiffusionPipeline(text_encoder=text_encoder_or_wrapper, vae=vae, unet=unet, tokenizer=tokenizer,
-                                     scheduler=scheduler, safety_checker=None, feature_extractor=None, requires_safety_checker=False)
-  pipeline.to(device)
-
-  save_dir = args.output_dir + "/sample"
-  os.makedirs(save_dir, exist_ok=True)
-
-  rng_state = torch.get_rng_state()
-  cuda_rng_state = torch.cuda.get_rng_state()
-
-  with torch.no_grad():
-    with accelerator.autocast():
-      for i, prompt in enumerate(prompts):
-        if not accelerator.is_main_process:
-          continue
-        prompt = prompt.strip()
-        if len(prompt) == 0 or prompt[0] == '#':
-          continue
-
-        # subset of gen_img_diffusers
-        prompt_args = prompt.split(' --')
-        prompt = prompt_args[0]
-        negative_prompt = None
-        sample_steps = 30
-        width = height = 512
-        scale = 7.5
-        seed = None
-        for parg in prompt_args:
-          try:
-            m = re.match(r'w (\d+)', parg, re.IGNORECASE)
-            if m:
-              width = int(m.group(1))
-              continue
-
-            m = re.match(r'h (\d+)', parg, re.IGNORECASE)
-            if m:
-              height = int(m.group(1))
-              continue
-
-            m = re.match(r'd (\d+)', parg, re.IGNORECASE)
-            if m:
-              seed = int(m.group(1))
-              continue
-
-            m = re.match(r's (\d+)', parg, re.IGNORECASE)
-            if m:               # steps
-              sample_steps = max(1, min(1000, int(m.group(1))))
-              continue
-
-            m = re.match(r'l ([\d\.]+)', parg, re.IGNORECASE)
-            if m:               # scale
-              scale = float(m.group(1))
-              continue
-
-            m = re.match(r'n (.+)', parg, re.IGNORECASE)
-            if m:               # negative prompt
-              negative_prompt = m.group(1)
-              continue
-
-          except ValueError as ex:
-            print(f"Exception in parsing / 解析エラー: {parg}")
-            print(ex)
-
-        if seed is not None:
-          torch.manual_seed(seed)
-          torch.cuda.manual_seed(seed)
-
-        if prompt_replacement is not None:
-          prompt = prompt.replace(prompt_replacement[0], prompt_replacement[1])
-          if negative_prompt is not None:
-            negative_prompt = negative_prompt.replace(prompt_replacement[0], prompt_replacement[1])
-
-        height = max(64, height - height % 8)                 # round to divisible by 8
-        width = max(64, width - width % 8)                 # round to divisible by 8
-        print(f"prompt: {prompt}")
-        print(f"negative_prompt: {negative_prompt}")
-        print(f"height: {height}")
-        print(f"width: {width}")
-        print(f"sample_steps: {sample_steps}")
-        print(f"scale: {scale}")
-        image = pipeline(prompt, height, width, sample_steps, scale, negative_prompt).images[0]
-
-        ts_str = time.strftime('%Y%m%d%H%M%S', time.localtime())
-        num_suffix = f"e{epoch:06d}" if epoch is not None else f"{steps:06d}"
-        seed_suffix = "" if seed is None else f"_{seed}"
-        img_filename = f"{'' if args.output_name is None else args.output_name + '_'}{ts_str}_{num_suffix}_{i:02d}{seed_suffix}.png"
-
-        image.save(os.path.join(save_dir, img_filename))
-
-  # clear pipeline and cache to reduce vram usage
-  del pipeline
-  torch.cuda.empty_cache()
-
-  torch.set_rng_state(rng_state)
-  torch.cuda.set_rng_state(cuda_rng_state)
-  vae.to(org_vae_device)
-
 # endregion
 
 # region 前処理用

lora_train_popup.py

@@ -0,0 +1,862 @@
+import gc
+import json
+import time
+from functools import partial
+from typing import Union
+import os
+import tkinter as tk
+from tkinter import filedialog as fd, ttk
+from tkinter import simpledialog as sd
+from tkinter import messagebox as mb
+
+import torch.cuda
+import train_network
+import library.train_util as util
+import argparse
+
+
+class ArgStore:
+    # Represents the entirety of all possible inputs for sd-scripts. they are ordered from most important to least
+    def __init__(self):
+        # Important, these are the most likely things you will modify
+        self.base_model: str = r""  # example path, r"E:\sd\stable-diffusion-webui\models\Stable-diffusion\nai.ckpt"
+        self.img_folder: str = r""  # is the folder path to your img folder, make sure to follow the guide here for folder setup: https://rentry.org/2chAI_LoRA_Dreambooth_guide_english#for-kohyas-script
+        self.output_folder: str = r""  # just the folder all epochs/safetensors are output
+        self.change_output_name: Union[str, None] = None  # changes the output name of the epochs
+        self.save_json_folder: Union[str, None] = None  # OPTIONAL, saves a json folder of your config to whatever location you set here.
+        self.load_json_path: Union[str, None] = None  # OPTIONAL, loads a json file partially changes the config to match. things like folder paths do not get modified.
+        self.json_load_skip_list: Union[list[str], None] = ["save_json_folder", "reg_img_folder",
+                                                            "lora_model_for_resume", "change_output_name",
+                                                            "training_comment",
+                                                            "json_load_skip_list"]  # OPTIONAL, allows the user to define what they skip when loading a json, by default it loads everything, including all paths, set it up like this ["base_model", "img_folder", "output_folder"]
+        self.caption_dropout_rate: Union[float, None] = None  # The rate at which captions for files get dropped.
+        self.caption_dropout_every_n_epochs: Union[int, None] = None  # Defines how often an epoch will completely ignore
+        # captions, EX. 3 means it will ignore captions at epochs 3, 6, and 9
+        self.caption_tag_dropout_rate: Union[float, None] = None  # Defines the rate at which a tag would be dropped, rather than the entire caption file
+        self.noise_offset: Union[float, None] = None  # OPTIONAL, seems to help allow SD to gen better blacks and whites
+                                                      # Kohya recommends, if you have it set, to use 0.1, not sure how
+                                                      # high the value can be, I'm going to assume maximum of 1
+
+        self.net_dim: int = 128  # network dimension, 128 is the most common, however you might be able to get lesser to work
+        self.alpha: float = 128  # represents the scalar for training. the lower the alpha, the less gets learned per step. if you want the older way of training, set this to dim
+        # list of schedulers: linear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmup
+        self.scheduler: str = "cosine_with_restarts"  # the scheduler for learning rate. Each does something specific
+        self.cosine_restarts: Union[int, None] = 1  # OPTIONAL, represents the number of times it restarts. Only matters if you are using cosine_with_restarts
+        self.scheduler_power: Union[float, None] = 1  # OPTIONAL, represents the power of the polynomial. Only matters if you are using polynomial
+        self.warmup_lr_ratio: Union[float, None] = None  # OPTIONAL, Calculates the number of warmup steps based on the ratio given. Make sure to set this if you are using constant_with_warmup, None to ignore
+        self.learning_rate: Union[float, None] = 1e-4  # OPTIONAL, when not set, lr gets set to 1e-3 as per adamW. Personally, I suggest actually setting this as lower lr seems to be a small bit better.
+        self.text_encoder_lr: Union[float, None] = None  # OPTIONAL, Sets a specific lr for the text encoder, this overwrites the base lr I believe, None to ignore
+        self.unet_lr: Union[float, None] = None  # OPTIONAL, Sets a specific lr for the unet, this overwrites the base lr I believe, None to ignore
+        self.num_workers: int = 1  # The number of threads that are being used to load images, lower speeds up the start of epochs, but slows down the loading of data. The assumption here is that it increases the training time as you reduce this value
+        self.persistent_workers: bool = True  # makes workers persistent, further reduces/eliminates the lag in between epochs. however it may increase memory usage
+
+        self.batch_size: int = 1  # The number of images that get processed at one time, this is directly proportional to your vram and resolution. with 12gb of vram, at 512 reso, you can get a maximum of 6 batch size
+        self.num_epochs: int = 1  # The number of epochs, if you set max steps this value is ignored as it doesn't calculate steps.
+        self.save_every_n_epochs: Union[int, None] = None  # OPTIONAL, how often to save epochs, None to ignore
+        self.shuffle_captions: bool = False  # OPTIONAL, False to ignore
+        self.keep_tokens: Union[int, None] = None  # OPTIONAL, None to ignore
+        self.max_steps: Union[int, None] = None  # OPTIONAL, if you have specific steps you want to hit, this allows you to set it directly. None to ignore
+        self.tag_occurrence_txt_file: bool = False  # OPTIONAL, creates a txt file that has the entire occurrence of all tags in your dataset
+                                                    # the metadata will also have this so long as you have metadata on, so no reason to have this on by default
+                                                    # will automatically output to the same folder as your output checkpoints
+        self.sort_tag_occurrence_alphabetically: bool = False  # OPTIONAL, only applies if tag_occurrence_txt_file is also true
+                                                               # Will change the output to be alphabetically vs being occurrence based
+
+        # These are the second most likely things you will modify
+        self.train_resolution: int = 512
+        self.min_bucket_resolution: int = 320
+        self.max_bucket_resolution: int = 960
+        self.lora_model_for_resume: Union[str, None] = None  # OPTIONAL, takes an input lora to continue training from, not exactly the way it *should* be, but it works, None to ignore
+        self.save_state: bool = False  # OPTIONAL, is the intended way to save a training state to use for continuing training, False to ignore
+        self.load_previous_save_state: Union[str, None] = None  # OPTIONAL, is the intended way to load a training state to use for continuing training, None to ignore
+        self.training_comment: Union[str, None] = None  # OPTIONAL, great way to put in things like activation tokens right into the metadata. seems to not work at this point and time
+        self.unet_only: bool = False  # OPTIONAL, set it to only train the unet
+        self.text_only: bool = False  # OPTIONAL, set it to only train the text encoder
+
+        # These are the least likely things you will modify
+        self.reg_img_folder: Union[str, None] = None  # OPTIONAL, None to ignore
+        self.clip_skip: int = 2  # If you are training on a model that is anime based, keep this at 2 as most models are designed for that
+        self.test_seed: int = 23  # this is the "reproducable seed", basically if you set the seed to this, you should be able to input a prompt from one of your training images and get a close representation of it
+        self.prior_loss_weight: float = 1  # is the loss weight much like Dreambooth, is required for LoRA training
+        self.gradient_checkpointing: bool = False  # OPTIONAL, enables gradient checkpointing
+        self.gradient_acc_steps: Union[int, None] = None  # OPTIONAL, not sure exactly what this means
+        self.mixed_precision: str = "fp16"  # If you have the ability to use bf16, do it, it's better
+        self.save_precision: str = "fp16"  # You can also save in bf16, but because it's not universally supported, I suggest you keep saving at fp16
+        self.save_as: str = "safetensors"  # list is pt, ckpt, safetensors
+        self.caption_extension: str = ".txt"  # the other option is .captions, but since wd1.4 tagger outputs as txt files, this is the default
+        self.max_clip_token_length = 150  # can be 75, 150, or 225 I believe, there is no reason to go higher than 150 though
+        self.buckets: bool = True
+        self.xformers: bool = True
+        self.use_8bit_adam: bool = True
+        self.cache_latents: bool = True
+        self.color_aug: bool = False  # IMPORTANT: Clashes with cache_latents, only have one of the two on!
+        self.flip_aug: bool = False
+        self.vae: Union[str, None] = None  # Seems to only make results worse when not using that specific vae, should probably not use
+        self.no_meta: bool = False  # This removes the metadata that now gets saved into safetensors, (you should keep this on)
+        self.log_dir: Union[str, None] = None  # output of logs, not useful to most people.
+        self.v2: bool = False  # Sets up training for SD2.1
+        self.v_parameterization: bool = False  # Only is used when v2 is also set and you are using the 768x version of v2
+
+    # Creates the dict that is used for the rest of the code, to facilitate easier json saving and loading
+    @staticmethod
+    def convert_args_to_dict():
+        return ArgStore().__dict__
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    setup_args(parser)
+    pre_args = parser.parse_args()
+    queues = 0
+    args_queue = []
+    cont = True
+    while cont:
+        arg_dict = ArgStore.convert_args_to_dict()
+        ret = mb.askyesno(message="Do you want to load a json config file?")
+        if ret:
+            load_json(ask_file("select json to load from", {"json"}), arg_dict)
+            arg_dict = ask_elements_trunc(arg_dict)
+        else:
+            arg_dict = ask_elements(arg_dict)
+        if pre_args.save_json_path or arg_dict["save_json_folder"]:
+            save_json(pre_args.save_json_path if pre_args.save_json_path else arg_dict['save_json_folder'], arg_dict)
+        args = create_arg_space(arg_dict)
+        args = parser.parse_args(args)
+        queues += 1
+        args_queue.append(args)
+        if arg_dict['tag_occurrence_txt_file']:
+            get_occurrence_of_tags(arg_dict)
+        ret = mb.askyesno(message="Do you want to queue another training?")
+        if not ret:
+            cont = False
+    for args in args_queue:
+        try:
+            train_network.train(args)
+        except Exception as e:
+            print(f"Failed to train this set of args.\nSkipping this training session.\nError is: {e}")
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+def create_arg_space(args: dict) -> [str]:
+    # This is the list of args that are to be used regardless of setup
+    output = ["--network_module=networks.lora", f"--pretrained_model_name_or_path={args['base_model']}",
+              f"--train_data_dir={args['img_folder']}", f"--output_dir={args['output_folder']}",
+              f"--prior_loss_weight={args['prior_loss_weight']}", f"--caption_extension=" + args['caption_extension'],
+              f"--resolution={args['train_resolution']}", f"--train_batch_size={args['batch_size']}",
+              f"--mixed_precision={args['mixed_precision']}", f"--save_precision={args['save_precision']}",
+              f"--network_dim={args['net_dim']}", f"--save_model_as={args['save_as']}",
+              f"--clip_skip={args['clip_skip']}", f"--seed={args['test_seed']}",
+              f"--max_token_length={args['max_clip_token_length']}", f"--lr_scheduler={args['scheduler']}",
+              f"--network_alpha={args['alpha']}", f"--max_data_loader_n_workers={args['num_workers']}"]
+    if not args['max_steps']:
+        output.append(f"--max_train_epochs={args['num_epochs']}")
+        output += create_optional_args(args, find_max_steps(args))
+    else:
+        output.append(f"--max_train_steps={args['max_steps']}")
+        output += create_optional_args(args, args['max_steps'])
+    return output
+
+
+def create_optional_args(args: dict, steps):
+    output = []
+    if args["reg_img_folder"]:
+        output.append(f"--reg_data_dir={args['reg_img_folder']}")
+
+    if args['lora_model_for_resume']:
+        output.append(f"--network_weights={args['lora_model_for_resume']}")
+
+    if args['save_every_n_epochs']:
+        output.append(f"--save_every_n_epochs={args['save_every_n_epochs']}")
+    else:
+        output.append("--save_every_n_epochs=999999")
+
+    if args['shuffle_captions']:
+        output.append("--shuffle_caption")
+
+    if args['keep_tokens'] and args['keep_tokens'] > 0:
+        output.append(f"--keep_tokens={args['keep_tokens']}")
+
+    if args['buckets']:
+        output.append("--enable_bucket")
+        output.append(f"--min_bucket_reso={args['min_bucket_resolution']}")
+        output.append(f"--max_bucket_reso={args['max_bucket_resolution']}")
+
+    if args['use_8bit_adam']:
+        output.append("--use_8bit_adam")
+
+    if args['xformers']:
+        output.append("--xformers")
+
+    if args['color_aug']:
+        if args['cache_latents']:
+            print("color_aug and cache_latents conflict with one another. Please select only one")
+            quit(1)
+        output.append("--color_aug")
+
+    if args['flip_aug']:
+        output.append("--flip_aug")
+
+    if args['cache_latents']:
+        output.append("--cache_latents")
+
+    if args['warmup_lr_ratio'] and args['warmup_lr_ratio'] > 0:
+        warmup_steps = int(steps * args['warmup_lr_ratio'])
+        output.append(f"--lr_warmup_steps={warmup_steps}")
+
+    if args['gradient_checkpointing']:
+        output.append("--gradient_checkpointing")
+
+    if args['gradient_acc_steps'] and args['gradient_acc_steps'] > 0 and args['gradient_checkpointing']:
+        output.append(f"--gradient_accumulation_steps={args['gradient_acc_steps']}")
+
+    if args['learning_rate'] and args['learning_rate'] > 0:
+        output.append(f"--learning_rate={args['learning_rate']}")
+
+    if args['text_encoder_lr'] and args['text_encoder_lr'] > 0:
+        output.append(f"--text_encoder_lr={args['text_encoder_lr']}")
+
+    if args['unet_lr'] and args['unet_lr'] > 0:
+        output.append(f"--unet_lr={args['unet_lr']}")
+
+    if args['vae']:
+        output.append(f"--vae={args['vae']}")
+
+    if args['no_meta']:
+        output.append("--no_metadata")
+
+    if args['save_state']:
+        output.append("--save_state")
+
+    if args['load_previous_save_state']:
+        output.append(f"--resume={args['load_previous_save_state']}")
+
+    if args['change_output_name']:
+        output.append(f"--output_name={args['change_output_name']}")
+
+    if args['training_comment']:
+        output.append(f"--training_comment={args['training_comment']}")
+
+    if args['cosine_restarts'] and args['scheduler'] == "cosine_with_restarts":
+        output.append(f"--lr_scheduler_num_cycles={args['cosine_restarts']}")
+
+    if args['scheduler_power'] and args['scheduler'] == "polynomial":
+        output.append(f"--lr_scheduler_power={args['scheduler_power']}")
+
+    if args['persistent_workers']:
+        output.append(f"--persistent_data_loader_workers")
+
+    if args['unet_only']:
+        output.append("--network_train_unet_only")
+
+    if args['text_only'] and not args['unet_only']:
+        output.append("--network_train_text_encoder_only")
+
+    if args["log_dir"]:
+        output.append(f"--logging_dir={args['log_dir']}")
+
+    if args['caption_dropout_rate']:
+        output.append(f"--caption_dropout_rate={args['caption_dropout_rate']}")
+
+    if args['caption_dropout_every_n_epochs']:
+        output.append(f"--caption_dropout_every_n_epochs={args['caption_dropout_every_n_epochs']}")
+
+    if args['caption_tag_dropout_rate']:
+        output.append(f"--caption_tag_dropout_rate={args['caption_tag_dropout_rate']}")
+
+    if args['v2']:
+        output.append("--v2")
+
+    if args['v2'] and args['v_parameterization']:
+        output.append("--v_parameterization")
+
+    if args['noise_offset']:
+        output.append(f"--noise_offset={args['noise_offset']}")
+    return output
+
+
+def find_max_steps(args: dict) -> int:
+    total_steps = 0
+    folders = os.listdir(args["img_folder"])
+    for folder in folders:
+        if not os.path.isdir(os.path.join(args["img_folder"], folder)):
+            continue
+        num_repeats = folder.split("_")
+        if len(num_repeats) < 2:
+            print(f"folder {folder} is not in the correct format. Format is x_name. skipping")
+            continue
+        try:
+            num_repeats = int(num_repeats[0])
+        except ValueError:
+            print(f"folder {folder} is not in the correct format. Format is x_name. skipping")
+            continue
+        imgs = 0
+        for file in os.listdir(os.path.join(args["img_folder"], folder)):
+            if os.path.isdir(file):
+                continue
+            ext = file.split(".")
+            if ext[-1].lower() in {"png", "bmp", "gif", "jpeg", "jpg", "webp"}:
+                imgs += 1
+        total_steps += (num_repeats * imgs)
+    total_steps = int((total_steps / args["batch_size"]) * args["num_epochs"])
+    return total_steps
+
+
+def add_misc_args(parser):
+    parser.add_argument("--save_json_path", type=str, default=None,
+                        help="Path to save a configuration json file to")
+    parser.add_argument("--load_json_path", type=str, default=None,
+                        help="Path to a json file to configure things from")
+    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"],
+                        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("--lr_scheduler_num_cycles", type=int, default=1,
+                        help="Number of restarts for cosine scheduler with restarts / cosine with restartsスケジューラでのリスタート回数")
+    parser.add_argument("--lr_scheduler_power", type=float, default=1,
+                        help="Polynomial power for polynomial scheduler / polynomialスケジューラでのpolynomial power")
+
+    parser.add_argument("--network_weights", type=str, default=None,
+                        help="pretrained weights for network / 学習するネットワークの初期重み")
+    parser.add_argument("--network_module", type=str, default=None,
+                        help='network module to train / 学習対象のネットワークのモジュール')
+    parser.add_argument("--network_dim", type=int, default=None,
+                        help='network dimensions (depends on each network) / モジュールの次元数（ネットワークにより定義は異なります）')
+    parser.add_argument("--network_alpha", type=float, 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",
+                        help="only training Text Encoder part / Text Encoder関連部分のみ学習する")
+    parser.add_argument("--training_comment", type=str, default=None,
+                        help="arbitrary comment string stored in metadata / メタデータに記録する任意のコメント文字列")
+
+
+def setup_args(parser):
+    util.add_sd_models_arguments(parser)
+    util.add_dataset_arguments(parser, True, True, True)
+    util.add_training_arguments(parser, True)
+    add_misc_args(parser)
+
+
+def get_occurrence_of_tags(args):
+    extension = args['caption_extension']
+    img_folder = args['img_folder']
+    output_folder = args['output_folder']
+    occurrence_dict = {}
+    print(img_folder)
+    for folder in os.listdir(img_folder):
+        print(folder)
+        if not os.path.isdir(os.path.join(img_folder, folder)):
+            continue
+        for file in os.listdir(os.path.join(img_folder, folder)):
+            if not os.path.isfile(os.path.join(img_folder, folder, file)):
+                continue
+            ext = os.path.splitext(file)[1]
+            if ext != extension:
+                continue
+            get_tags_from_file(os.path.join(img_folder, folder, file), occurrence_dict)
+    if not args['sort_tag_occurrence_alphabetically']:
+        output_list = {k: v for k, v in sorted(occurrence_dict.items(), key=lambda item: item[1], reverse=True)}
+    else:
+        output_list = {k: v for k, v in sorted(occurrence_dict.items(), key=lambda item: item[0])}
+    name = args['change_output_name'] if args['change_output_name'] else "last"
+    with open(os.path.join(output_folder, f"{name}.txt"), "w") as f:
+        f.write(f"Below is a list of keywords used during the training of {args['change_output_name']}:\n")
+        for k, v in output_list.items():
+            f.write(f"[{v}] {k}\n")
+    print(f"Created a txt file named {name}.txt in the output folder")
+
+
+def get_tags_from_file(file, occurrence_dict):
+    f = open(file)
+    temp = f.read().replace(", ", ",").split(",")
+    f.close()
+    for tag in temp:
+        if tag in occurrence_dict:
+            occurrence_dict[tag] += 1
+        else:
+            occurrence_dict[tag] = 1
+
+
+def ask_file(message, accepted_ext_list, file_path=None):
+    mb.showinfo(message=message)
+    res = ""
+    _initialdir = ""
+    _initialfile = ""
+    if file_path != None:
+        _initialdir = os.path.dirname(file_path) if os.path.exists(file_path) else ""
+        _initialfile = os.path.basename(file_path) if os.path.exists(file_path) else ""
+
+    while res == "":
+        res = fd.askopenfilename(title=message, initialdir=_initialdir, initialfile=_initialfile)
+        if res == "" or type(res) == tuple:
+            ret = mb.askretrycancel(message="Do you want to to cancel training?")
+            if not ret:
+                exit()
+            continue
+        elif not os.path.exists(res):
+            res = ""
+            continue
+        _, name = os.path.split(res)
+        split_name = name.split(".")
+        if split_name[-1] not in accepted_ext_list:
+            res = ""
+    return res
+
+
+def ask_dir(message, dir_path=None):
+    mb.showinfo(message=message)
+    res = ""
+    _initialdir = ""
+    if dir_path != None:
+        _initialdir = dir_path if os.path.exists(dir_path) else ""
+    while res == "":
+        res = fd.askdirectory(title=message, initialdir=_initialdir)
+        if res == "" or type(res) == tuple:
+            ret = mb.askretrycancel(message="Do you want to to cancel training?")
+            if not ret:
+                exit()
+            continue
+        if not os.path.exists(res):
+            res = ""
+    return res
+
+
+def ask_elements_trunc(args: dict):
+    args['base_model'] = ask_file("Select your base model", {"ckpt", "safetensors"}, args['base_model'])
+    args['img_folder'] = ask_dir("Select your image folder", args['img_folder'])
+    args['output_folder'] = ask_dir("Select your output folder", args['output_folder'])
+
+    ret = mb.askyesno(message="Do you want to save a json of your configuration?")
+    if ret:
+        args['save_json_folder'] = ask_dir("Select the folder to save json files to", args['save_json_folder'])
+    else:
+        args['save_json_folder'] = None
+
+    ret = mb.askyesno(message="Are you training on a SD2 based model?")
+    if ret:
+        args['v2'] = True
+
+    ret = mb.askyesno(message="Are you training on an realistic model?")
+    if ret:
+        args['clip_skip'] = 1
+
+    if args['v2']:
+        ret = mb.askyesno(message="Are you training on a model based on the 768x version of SD2?")
+        if ret:
+            args['v_parameterization'] = True
+
+    ret = mb.askyesno(message="Do you want to use regularization images?")
+    if ret:
+        args['reg_img_folder'] = ask_dir("Select your regularization folder", args['reg_img_folder'])
+    else:
+        args['reg_img_folder'] = None
+
+    ret = mb.askyesno(message="Do you want to continue from an earlier version?")
+    if ret:
+        args['lora_model_for_resume'] = ask_file("Select your lora model", {"ckpt", "pt", "safetensors"},
+                                                 args['lora_model_for_resume'])
+    else:
+        args['lora_model_for_resume'] = None
+
+    ret = mb.askyesno(message="Do you want to flip all of your images? It is supposed to reduce biases\n"
+                              "within your dataset but it can also ruin learning an asymmetrical element\n")
+    if ret:
+        args['flip_aug'] = True
+
+    ret = mb.askyesno(message="Do you want to change the name of output checkpoints?")
+    if ret:
+        ret = sd.askstring(title="output_name", prompt="What do you want your output name to be?\n"
+                                                       "Cancel keeps outputs the original")
+        if ret:
+            args['change_output_name'] = ret
+        else:
+            args['change_output_name'] = None
+
+    ret = sd.askstring(title="comment",
+                       prompt="Do you want to set a comment that gets put into the metadata?\nA good use of this would "
+                              "be to include how to use, such as activation keywords.\nCancel will leave empty")
+    if ret is None:
+        args['training_comment'] = ret
+    else:
+        args['training_comment'] = None
+
+    ret = mb.askyesno(message="Do you want to train only one of unet and text encoder?")
+    if ret:
+        button = ButtonBox("Which do you want to train with?", ["unet_only", "text_only"])
+        button.window.mainloop()
+        if button.current_value != "":
+            args[button.current_value] = True
+
+    ret = mb.askyesno(message="Do you want to save a txt file that contains a list\n"
+                              "of all tags that you have used in your training data?\n")
+    if ret:
+        args['tag_occurrence_txt_file'] = True
+        button = ButtonBox("How do you want tags to be ordered?", ["alphabetically", "occurrence-ly"])
+        button.window.mainloop()
+        if button.current_value == "alphabetically":
+            args['sort_tag_occurrence_alphabetically'] = True
+
+    ret = mb.askyesno(message="Do you want to use caption dropout?")
+    if ret:
+        ret = mb.askyesno(message="Do you want full caption files to dropout randomly?")
+        if ret:
+            ret = sd.askinteger(title="Caption_File_Dropout",
+                                prompt="How often do you want caption files to drop out?\n"
+                                       "enter a number from 0 to 100 that is the percentage chance of dropout\n"
+                                       "Cancel sets to 0")
+            if ret and 0 <= ret <= 100:
+                args['caption_dropout_rate'] = ret / 100.0
+
+        ret = mb.askyesno(message="Do you want to have full epochs have no captions?")
+        if ret:
+            ret = sd.askinteger(title="Caption_epoch_dropout", prompt="The number set here is how often you will have an"
+                                                                      "epoch with no captions\nSo if you set 3, then every"
+                                                                      "three epochs will not have captions (3, 6, 9)\n"
+                                                                      "Cancel will set to None")
+            if ret:
+                args['caption_dropout_every_n_epochs'] = ret
+
+        ret = mb.askyesno(message="Do you want to have tags to randomly drop?")
+        if ret:
+            ret = sd.askinteger(title="Caption_tag_dropout", prompt="How often do you want tags to randomly drop out?\n"
+                                                                    "Enter a number between 0 and 100, that is the percentage"
+                                                                    "chance of dropout.\nCancel sets to 0")
+            if ret and 0 <= ret <= 100:
+                args['caption_tag_dropout_rate'] = ret / 100.0
+
+    ret = mb.askyesno(message="Do you want to use noise offset? Noise offset seems to allow for SD to better generate\n"
+                              "darker or lighter images using this than normal.")
+    if ret:
+        ret = sd.askfloat(title="noise_offset", prompt="What value do you want to set? recommended value is 0.1,\n"
+                                                       "but it can go higher. Cancel defaults to 0.1")
+        if ret:
+            args['noise_offset'] = ret
+        else:
+            args['noise_offset'] = 0.1
+    return args
+
+
+def ask_elements(args: dict):
+    # start with file dialog
+    args['base_model'] = ask_file("Select your base model", {"ckpt", "safetensors"}, args['base_model'])
+    args['img_folder'] = ask_dir("Select your image folder", args['img_folder'])
+    args['output_folder'] = ask_dir("Select your output folder", args['output_folder'])
+
+    # optional file dialog
+    ret = mb.askyesno(message="Do you want to save a json of your configuration?")
+    if ret:
+        args['save_json_folder'] = ask_dir("Select the folder to save json files to", args['save_json_folder'])
+    else:
+        args['save_json_folder'] = None
+
+    ret = mb.askyesno(message="Are you training on a SD2 based model?")
+    if ret:
+        args['v2'] = True
+
+    ret = mb.askyesno(message="Are you training on an realistic model?")
+    if ret:
+        args['clip_skip'] = 1
+
+    if args['v2']:
+        ret = mb.askyesno(message="Are you training on a model based on the 768x version of SD2?")
+        if ret:
+            args['v_parameterization'] = True
+
+    ret = mb.askyesno(message="Do you want to use regularization images?")
+    if ret:
+        args['reg_img_folder'] = ask_dir("Select your regularization folder", args['reg_img_folder'])
+    else:
+        args['reg_img_folder'] = None
+
+    ret = mb.askyesno(message="Do you want to continue from an earlier version?")
+    if ret:
+        args['lora_model_for_resume'] = ask_file("Select your lora model", {"ckpt", "pt", "safetensors"},
+                                                 args['lora_model_for_resume'])
+    else:
+        args['lora_model_for_resume'] = None
+
+    ret = mb.askyesno(message="Do you want to flip all of your images? It is supposed to reduce biases\n"
+                              "within your dataset but it can also ruin learning an asymmetrical element\n")
+    if ret:
+        args['flip_aug'] = True
+
+    # text based required elements
+    ret = sd.askinteger(title="batch_size",
+                        prompt="The number of images that get processed at one time, this is directly proportional to "
+                               "your vram and resolution. with 12gb of vram, at 512 reso, you can get a maximum of 6 "
+                               "batch size\nHow large is your batch size going to be?\nCancel will default to 1")
+    if ret is None:
+        args['batch_size'] = 1
+    else:
+        args['batch_size'] = ret
+
+    ret = sd.askinteger(title="num_epochs", prompt="How many epochs do you want?\nCancel will default to 1")
+    if ret is None:
+        args['num_epochs'] = 1
+    else:
+        args['num_epochs'] = ret
+
+    ret = sd.askinteger(title="network_dim", prompt="What is the dim size you want to use?\nCancel will default to 128")
+    if ret is None:
+        args['net_dim'] = 128
+    else:
+        args['net_dim'] = ret
+
+    ret = sd.askfloat(title="alpha", prompt="Alpha is the scalar of the training, generally a good starting point is "
+                                            "0.5x dim size\nWhat Alpha do you want?\nCancel will default to equal to "
+                                            "0.5 x network_dim")
+    if ret is None:
+        args['alpha'] = args['net_dim'] / 2
+    else:
+        args['alpha'] = ret
+
+    ret = sd.askinteger(title="resolution", prompt="How large of a resolution do you want to train at?\n"
+                                                   "Cancel will default to 512")
+    if ret is None:
+        args['train_resolution'] = 512
+    else:
+        args['train_resolution'] = ret
+
+    ret = sd.askfloat(title="learning_rate", prompt="What learning rate do you want to use?\n"
+                                                    "Cancel will default to 1e-4")
+    if ret is None:
+        args['learning_rate'] = 1e-4
+    else:
+        args['learning_rate'] = ret
+
+    ret = sd.askfloat(title="text_encoder_lr", prompt="Do you want to set the text_encoder_lr?\n"
+                                                      "Cancel will default to None")
+    if ret is None:
+        args['text_encoder_lr'] = None
+    else:
+        args['text_encoder_lr'] = ret
+
+    ret = sd.askfloat(title="unet_lr", prompt="Do you want to set the unet_lr?\nCancel will default to None")
+    if ret is None:
+        args['unet_lr'] = None
+    else:
+        args['unet_lr'] = ret
+
+    button = ButtonBox("Which scheduler do you want?", ["cosine_with_restarts", "cosine", "polynomial",
+                                                        "constant", "constant_with_warmup", "linear"])
+    button.window.mainloop()
+    args['scheduler'] = button.current_value if button.current_value != "" else "cosine_with_restarts"
+
+    if args['scheduler'] == "cosine_with_restarts":
+        ret = sd.askinteger(title="Cycle Count",
+                            prompt="How many times do you want cosine to restart?\nThis is the entire amount of times "
+                                   "it will restart for the entire training\nCancel will default to 1")
+        if ret is None:
+            args['cosine_restarts'] = 1
+        else:
+            args['cosine_restarts'] = ret
+
+    if args['scheduler'] == "polynomial":
+        ret = sd.askfloat(title="Poly Strength",
+                          prompt="What power do you want to set your polynomial to?\nhigher power means that the "
+                                 "model reduces the learning more more aggressively from initial training.\n1 = "
+                                 "linear\nCancel sets to 1")
+        if ret is None:
+            args['scheduler_power'] = 1
+        else:
+            args['scheduler_power'] = ret
+
+    ret = mb.askyesno(message="Do you want to save epochs as it trains?")
+    if ret:
+        ret = sd.askinteger(title="save_epoch",
+                            prompt="How often do you want to save epochs?\nCancel will default to 1")
+        if ret is None:
+            args['save_every_n_epochs'] = 1
+        else:
+            args['save_every_n_epochs'] = ret
+
+    ret = mb.askyesno(message="Do you want to shuffle captions?")
+    if ret:
+        args['shuffle_captions'] = True
+    else:
+        args['shuffle_captions'] = False
+
+    ret = mb.askyesno(message="Do you want to keep some tokens at the front of your captions?")
+    if ret:
+        ret = sd.askinteger(title="keep_tokens", prompt="How many do you want to keep at the front?"
+                                                        "\nCancel will default to 1")
+        if ret is None:
+            args['keep_tokens'] = 1
+        else:
+            args['keep_tokens'] = ret
+
+    ret = mb.askyesno(message="Do you want to have a warmup ratio?")
+    if ret:
+        ret = sd.askfloat(title="warmup_ratio", prompt="What is the ratio of steps to use as warmup "
+                                                       "steps?\nCancel will default to None")
+        if ret is None:
+            args['warmup_lr_ratio'] = None
+        else:
+            args['warmup_lr_ratio'] = ret
+
+    ret = mb.askyesno(message="Do you want to change the name of output checkpoints?")
+    if ret:
+        ret = sd.askstring(title="output_name", prompt="What do you want your output name to be?\n"
+                                                       "Cancel keeps outputs the original")
+        if ret:
+            args['change_output_name'] = ret
+        else:
+            args['change_output_name'] = None
+
+    ret = sd.askstring(title="comment",
+                       prompt="Do you want to set a comment that gets put into the metadata?\nA good use of this would "
+                              "be to include how to use, such as activation keywords.\nCancel will leave empty")
+    if ret is None:
+        args['training_comment'] = ret
+    else:
+        args['training_comment'] = None
+
+    ret = mb.askyesno(message="Do you want to train only one of unet and text encoder?")
+    if ret:
+        if ret:
+            button = ButtonBox("Which do you want to train with?", ["unet_only", "text_only"])
+            button.window.mainloop()
+            if button.current_value != "":
+                args[button.current_value] = True
+
+    ret = mb.askyesno(message="Do you want to save a txt file that contains a list\n"
+                              "of all tags that you have used in your training data?\n")
+    if ret:
+        args['tag_occurrence_txt_file'] = True
+        button = ButtonBox("How do you want tags to be ordered?", ["alphabetically", "occurrence-ly"])
+        button.window.mainloop()
+        if button.current_value == "alphabetically":
+            args['sort_tag_occurrence_alphabetically'] = True
+
+    ret = mb.askyesno(message="Do you want to use caption dropout?")
+    if ret:
+        ret = mb.askyesno(message="Do you want full caption files to dropout randomly?")
+        if ret:
+            ret = sd.askinteger(title="Caption_File_Dropout",
+                                prompt="How often do you want caption files to drop out?\n"
+                                       "enter a number from 0 to 100 that is the percentage chance of dropout\n"
+                                       "Cancel sets to 0")
+            if ret and 0 <= ret <= 100:
+                args['caption_dropout_rate'] = ret / 100.0
+
+        ret = mb.askyesno(message="Do you want to have full epochs have no captions?")
+        if ret:
+            ret = sd.askinteger(title="Caption_epoch_dropout", prompt="The number set here is how often you will have an"
+                                                                      "epoch with no captions\nSo if you set 3, then every"
+                                                                      "three epochs will not have captions (3, 6, 9)\n"
+                                                                      "Cancel will set to None")
+            if ret:
+                args['caption_dropout_every_n_epochs'] = ret
+
+        ret = mb.askyesno(message="Do you want to have tags to randomly drop?")
+        if ret:
+            ret = sd.askinteger(title="Caption_tag_dropout", prompt="How often do you want tags to randomly drop out?\n"
+                                                                    "Enter a number between 0 and 100, that is the percentage"
+                                                                    "chance of dropout.\nCancel sets to 0")
+            if ret and 0 <= ret <= 100:
+                args['caption_tag_dropout_rate'] = ret / 100.0
+
+    ret = mb.askyesno(message="Do you want to use noise offset? Noise offset seems to allow for SD to better generate\n"
+                              "darker or lighter images using this than normal.")
+    if ret:
+        ret = sd.askfloat(title="noise_offset", prompt="What value do you want to set? recommended value is 0.1,\n"
+                                                       "but it can go higher. Cancel defaults to 0.1")
+        if ret:
+            args['noise_offset'] = ret
+        else:
+            args['noise_offset'] = 0.1
+    return args
+
+
+def save_json(path, obj: dict) -> None:
+    fp = open(os.path.join(path, f"config-{time.time()}.json"), "w")
+    json.dump(obj, fp=fp, indent=4)
+    fp.close()
+
+
+def load_json(path, obj: dict) -> dict:
+    with open(path) as f:
+        json_obj = json.loads(f.read())
+    print("loaded json, setting variables...")
+    ui_name_scheme = {"pretrained_model_name_or_path": "base_model", "logging_dir": "log_dir",
+                      "train_data_dir": "img_folder", "reg_data_dir": "reg_img_folder",
+                      "output_dir": "output_folder", "max_resolution": "train_resolution",
+                      "lr_scheduler": "scheduler", "lr_warmup": "warmup_lr_ratio",
+                      "train_batch_size": "batch_size", "epoch": "num_epochs",
+                      "save_at_n_epochs": "save_every_n_epochs", "num_cpu_threads_per_process": "num_workers",
+                      "enable_bucket": "buckets", "save_model_as": "save_as", "shuffle_caption": "shuffle_captions",
+                      "resume": "load_previous_save_state", "network_dim": "net_dim",
+                      "gradient_accumulation_steps": "gradient_acc_steps", "output_name": "change_output_name",
+                      "network_alpha": "alpha", "lr_scheduler_num_cycles": "cosine_restarts",
+                      "lr_scheduler_power": "scheduler_power"}
+
+    for key in list(json_obj):
+        if key in ui_name_scheme:
+            json_obj[ui_name_scheme[key]] = json_obj[key]
+            if ui_name_scheme[key] in {"batch_size", "num_epochs"}:
+                try:
+                    json_obj[ui_name_scheme[key]] = int(json_obj[ui_name_scheme[key]])
+                except ValueError:
+                    print(f"attempting to load {key} from json failed as input isn't an integer")
+                    quit(1)
+
+    for key in list(json_obj):
+        if obj["json_load_skip_list"] and key in obj["json_load_skip_list"]:
+            continue
+        if key in obj:
+            if key in {"keep_tokens", "warmup_lr_ratio"}:
+                json_obj[key] = int(json_obj[key]) if json_obj[key] is not None else None
+            if key in {"learning_rate", "unet_lr", "text_encoder_lr"}:
+                json_obj[key] = float(json_obj[key]) if json_obj[key] is not None else None
+            if obj[key] != json_obj[key]:
+                print_change(key, obj[key], json_obj[key])
+                obj[key] = json_obj[key]
+    print("completed changing variables.")
+    return obj
+
+
+def print_change(value, old, new):
+    print(f"{value} changed from {old} to {new}")
+
+
+class ButtonBox:
+    def __init__(self, label: str, button_name_list: list[str]) -> None:
+        self.window = tk.Tk()
+        self.button_list = []
+        self.current_value = ""
+
+        self.window.attributes("-topmost", True)
+        self.window.resizable(False, False)
+        self.window.eval('tk::PlaceWindow . center')
+
+        def del_window():
+            self.window.quit()
+            self.window.destroy()
+
+        self.window.protocol("WM_DELETE_WINDOW", del_window)
+        tk.Label(text=label, master=self.window).pack()
+        for button in button_name_list:
+            self.button_list.append(ttk.Button(text=button, master=self.window,
+                                               command=partial(self.set_current_value, button)))
+            self.button_list[-1].pack()
+
+    def set_current_value(self, value):
+        self.current_value = value
+        self.window.quit()
+        self.window.destroy()
+
+
+root = tk.Tk()
+root.attributes('-topmost', True)
+root.withdraw()
+
+if __name__ == "__main__":
+    main()

lycoris/kohya.py

@@ -5,7 +5,6 @@
 # https://github.com/kohya-ss/sd-scripts/blob/main/networks/lora.py
 
 import math
-from warnings import warn
 import os
 from typing import List
 import torch
@@ -28,22 +27,6 @@ def create_network(multiplier, network_dim, network_alpha, vae, text_encoder, un
     }[algo]
     
     print(f'Using rank adaptation algo: {algo}')
-    
-    if (algo == 'loha' 
-        and not kwargs.get('no_dim_warn', False) 
-        and (network_dim>64 or conv_dim>64)):
-        print('='*20 + 'WARNING' + '='*20)
-        warn(
-            (
-                "You are not supposed to use dim>64 (64*64 = 4096, it already has enough rank)"
-                "in Hadamard Product representation!\n"
-                "Please consider use lower dim or disable this warning with --network_args no_dim_warn=True\n"
-                "If you just want to use high dim loha, please consider use lower lr."
-            ),
-            stacklevel=2,
-        )
-        print('='*20 + 'WARNING' + '='*20)
-    
     network = LoRANetwork(
         text_encoder, unet, 
         multiplier=multiplier, 

lycoris/loha.py

@@ -36,12 +36,7 @@ class LohaModule(nn.Module):
     Hadamard product Implementaion for Low Rank Adaptation
     """
 
-    def __init__(
-        self, 
-        lora_name, 
-        org_module: nn.Module, 
-        multiplier=1.0, lora_dim=4, alpha=1, dropout=0.,
-    ):
+    def __init__(self, lora_name, org_module: nn.Module, multiplier=1.0, lora_dim=4, alpha=1, dropout=0.):
         """ if alpha == 0 or None, alpha is rank (no scaling). """
         super().__init__()
         self.lora_name = lora_name

lycoris/utils.py

@@ -28,13 +28,11 @@ def extract_conv(
         assert 1>=mode_param>=0
         min_s = torch.max(S)*mode_param
         lora_rank = torch.sum(S>min_s)
-    elif mode=='quantile' or mode=='percentile':
+    elif mode=='percentile':
         assert 1>=mode_param>=0
         s_cum = torch.cumsum(S, dim=0)
         min_cum_sum = mode_param * torch.sum(S)
         lora_rank = torch.sum(s_cum<min_cum_sum)
-    else:
-        raise NotImplementedError('Extract mode should be "fixed", "threshold", "ratio" or "quantile"')
     lora_rank = max(1, lora_rank)
     lora_rank = min(out_ch, in_ch, lora_rank)
     
@@ -90,13 +88,11 @@ def extract_linear(
         assert 1>=mode_param>=0
         min_s = torch.max(S)*mode_param
         lora_rank = torch.sum(S>min_s)
-    elif mode=='quantile' or mode=='percentile':
+    elif mode=='percentile':
         assert 1>=mode_param>=0
         s_cum = torch.cumsum(S, dim=0)
         min_cum_sum = mode_param * torch.sum(S)
         lora_rank = torch.sum(s_cum<min_cum_sum)
-    else:
-        raise NotImplementedError('Extract mode should be "fixed", "threshold", "ratio" or "quantile"')
     lora_rank = max(1, lora_rank)
     lora_rank = min(out_ch, in_ch, lora_rank)
     
@@ -263,69 +259,6 @@ def merge_locon(
                         child_module.weight += (alpha.to(device)/rank * scale * delta).cpu()
                         del delta
     
-    merge(
-        LORA_PREFIX_TEXT_ENCODER, 
-        base_model[0], 
-        TEXT_ENCODER_TARGET_REPLACE_MODULE
-    )
-    merge(
-        LORA_PREFIX_UNET,
-        base_model[2], 
-        UNET_TARGET_REPLACE_MODULE
-    )
-
-
-def merge_loha(
-    base_model,
-    loha_state_dict: Dict[str, torch.TensorType],
-    scale: float = 1.0,
-    device = 'cpu'
-):
-    UNET_TARGET_REPLACE_MODULE = [
-        "Transformer2DModel", 
-        "Attention", 
-        "ResnetBlock2D", 
-        "Downsample2D", 
-        "Upsample2D"
-    ]
-    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
-    LORA_PREFIX_UNET = 'lora_unet'
-    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
-    def merge(
-        prefix, 
-        root_module: torch.nn.Module,
-        target_replace_modules
-    ):
-        temp = {}
-        
-        for name, module in tqdm(list(root_module.named_modules())):
-            if module.__class__.__name__ in target_replace_modules:
-                temp[name] = {}
-                for child_name, child_module in module.named_modules():
-                    layer = child_module.__class__.__name__
-                    if layer not in {'Linear', 'Conv2d'}:
-                        continue
-                    lora_name = prefix + '.' + name + '.' + child_name
-                    lora_name = lora_name.replace('.', '_')
-                    
-                    w1a = loha_state_dict[f'{lora_name}.hada_w1_a'].float().to(device)
-                    w1b = loha_state_dict[f'{lora_name}.hada_w1_b'].float().to(device)
-                    w2a = loha_state_dict[f'{lora_name}.hada_w2_a'].float().to(device)
-                    w2b = loha_state_dict[f'{lora_name}.hada_w2_b'].float().to(device)
-                    alpha = loha_state_dict[f'{lora_name}.alpha'].float().to(device)
-                    dim = w1b.shape[0]
-                    
-                    delta = (w1a @ w1b) * (w2a @ w2b)
-                    delta = delta.reshape(child_module.weight.shape)
-                    
-                    if layer == 'Conv2d':
-                        child_module.weight.requires_grad_(False)
-                        child_module.weight += (alpha.to(device)/dim * scale * delta).cpu()
-                    elif layer == 'Linear':
-                        child_module.weight.requires_grad_(False)
-                        child_module.weight += (alpha.to(device)/dim * scale * delta).cpu()
-                    del delta
-    
     merge(
         LORA_PREFIX_TEXT_ENCODER, 
         base_model[0], 

networks/check_lora_weights.py

@@ -21,7 +21,7 @@ def main(file):
 
   for key, value in values:
     value = value.to(torch.float32)
-    print(f"{key},{str(tuple(value.size())).replace(', ', '-')},{torch.mean(torch.abs(value))},{torch.min(torch.abs(value))}")
+    print(f"{key},{torch.mean(torch.abs(value))},{torch.min(torch.abs(value))}")
 
 
 if __name__ == '__main__':

networks/extract_lora_from_models.py

@@ -45,13 +45,8 @@ def svd(args):
   text_encoder_t, _, unet_t = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.model_tuned)
 
   # create LoRA network to extract weights: Use dim (rank) as alpha
-  if args.conv_dim is None:
-    kwargs = {}
-  else:
-    kwargs = {"conv_dim": args.conv_dim, "conv_alpha": args.conv_dim}
-
-  lora_network_o = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_o, unet_o, **kwargs)
-  lora_network_t = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_t, unet_t, **kwargs)
+  lora_network_o = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_o, unet_o)
+  lora_network_t = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_t, unet_t)
   assert len(lora_network_o.text_encoder_loras) == len(
       lora_network_t.text_encoder_loras), f"model version is different (SD1.x vs SD2.x) / それぞれのモデルのバージョンが違います（SD1.xベースとSD2.xベース） "
 
@@ -90,28 +85,13 @@ def svd(args):
 
   # make LoRA with svd
   print("calculating by svd")
+  rank = args.dim
   lora_weights = {}
   with torch.no_grad():
     for lora_name, mat in tqdm(list(diffs.items())):
-      # if args.conv_dim is None, diffs do not include LoRAs for conv2d-3x3
       conv2d = (len(mat.size()) == 4)
-      kernel_size = None if not conv2d else mat.size()[2:4]
-      conv2d_3x3 = conv2d and kernel_size != (1, 1)
-
-      rank = args.dim if not conv2d_3x3 or args.conv_dim is None else args.conv_dim
-      out_dim, in_dim = mat.size()[0:2]
-
-      if args.device:
-        mat = mat.to(args.device)
-
-      # print(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:
-        if conv2d_3x3:
-          mat = mat.flatten(start_dim=1)
-        else:
-          mat = mat.squeeze()
+        mat = mat.squeeze()
 
       U, S, Vh = torch.linalg.svd(mat)
 
@@ -128,27 +108,30 @@ def svd(args):
       U = U.clamp(low_val, hi_val)
       Vh = Vh.clamp(low_val, hi_val)
 
-      if conv2d:
-        U = U.reshape(out_dim, rank, 1, 1)
-        Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
-
-      U = U.to("cpu").contiguous()
-      Vh = Vh.to("cpu").contiguous()
-
       lora_weights[lora_name] = (U, Vh)
 
   # make state dict for LoRA
-  lora_sd = {}
-  for lora_name, (up_weight, down_weight) in lora_weights.items():
-    lora_sd[lora_name + '.lora_up.weight'] = up_weight
-    lora_sd[lora_name + '.lora_down.weight'] = down_weight
-    lora_sd[lora_name + '.alpha'] = torch.tensor(down_weight.size()[0])
+  lora_network_o.apply_to(text_encoder_o, unet_o, text_encoder_different, True)   # to make state dict
+  lora_sd = lora_network_o.state_dict()
+  print(f"LoRA has {len(lora_sd)} weights.")
 
-  # load state dict to LoRA and save it
-  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
+  for key in list(lora_sd.keys()):
+    if "alpha" in key:
+      continue
+
+    lora_name = key.split('.')[0]
+    i = 0 if "lora_up" in key else 1
 
-  info = lora_network_save.load_state_dict(lora_sd)
+    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)
   print(f"Loading extracted LoRA weights: {info}")
 
   dir_name = os.path.dirname(args.save_to)
@@ -156,9 +139,9 @@ def svd(args):
     os.makedirs(dir_name, exist_ok=True)
 
   # minimum metadata
-  metadata = {"ss_network_module": "networks.lora", "ss_network_dim": str(args.dim), "ss_network_alpha": str(args.dim)}
+  metadata = {"ss_network_dim": str(args.dim), "ss_network_alpha": str(args.dim)}
 
-  lora_network_save.save_weights(args.save_to, save_dtype, metadata)
+  lora_network_o.save_weights(args.save_to, save_dtype, metadata)
   print(f"LoRA weights are saved to: {args.save_to}")
 
 
@@ -175,8 +158,6 @@ if __name__ == '__main__':
   parser.add_argument("--save_to", type=str, default=None,
                       help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors")
   parser.add_argument("--dim", type=int, default=4, help="dimension (rank) of LoRA (default 4) / LoRAの次元数（rank）（デフォルト4）")
-  parser.add_argument("--conv_dim", type=int, default=None,
-                      help="dimension (rank) of LoRA for Conv2d-3x3 (default None, disabled) / LoRAのConv2d-3x3の次元数（rank）（デフォルトNone、適用なし）")
   parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
 
   args = parser.parse_args()

networks/lora.py

@@ -6,7 +6,6 @@
 import math
 import os
 from typing import List
-import numpy as np
 import torch
 
 from library import train_util
@@ -21,34 +20,22 @@ 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
+      self.lora_down = torch.nn.Conv2d(in_dim, lora_dim, (1, 1), bias=False)
+      self.lora_up = torch.nn.Conv2d(lora_dim, out_dim, (1, 1), bias=False)
     else:
       in_dim = org_module.in_features
       out_dim = org_module.out_features
-
-    # if limit_rank:
-    #   self.lora_dim = min(lora_dim, in_dim, out_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:
-      self.lora_down = torch.nn.Linear(in_dim, self.lora_dim, bias=False)
-      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 = self.lora_dim if alpha is None or alpha == 0 else alpha
+    alpha = lora_dim if alpha is None or alpha == 0 else alpha
     self.scale = alpha / self.lora_dim
     self.register_buffer('alpha', torch.tensor(alpha))                    # 定数として扱える
 
@@ -58,192 +45,69 @@ class LoRAModule(torch.nn.Module):
 
     self.multiplier = multiplier
     self.org_module = org_module                  # remove in applying
-    self.region = None
-    self.region_mask = None
 
   def apply_to(self):
     self.org_forward = self.org_module.forward
     self.org_module.forward = self.forward
     del self.org_module
 
-  def set_region(self, region):
-    self.region = region
-    self.region_mask = None
-
   def forward(self, x):
-    if self.region is None:
-      return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
-
-    # regional LoRA   FIXME same as additional-network extension
-    if x.size()[1] % 77 == 0:
-      # print(f"LoRA for context: {self.lora_name}")
-      self.region = None
-      return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
-
-    # calculate region mask first time
-    if self.region_mask is None:
-      if len(x.size()) == 4:
-        h, w = x.size()[2:4]
-      else:
-        seq_len = x.size()[1]
-        ratio = math.sqrt((self.region.size()[0] * self.region.size()[1]) / seq_len)
-        h = int(self.region.size()[0] / ratio + .5)
-        w = seq_len // h
-
-      r = self.region.to(x.device)
-      if r.dtype == torch.bfloat16:
-        r = r.to(torch.float)
-      r = r.unsqueeze(0).unsqueeze(1)
-      # print(self.lora_name, self.region.size(), x.size(), r.size(), h, w)
-      r = torch.nn.functional.interpolate(r, (h, w), mode='bilinear')
-      r = r.to(x.dtype)
-
-      if len(x.size()) == 3:
-        r = torch.reshape(r, (1, x.size()[1], -1))
-
-      self.region_mask = r
-
-    return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale * self.region_mask
+    return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
 
 
 def create_network(multiplier, network_dim, network_alpha, vae, text_encoder, unet, **kwargs):
   if network_dim is None:
     network_dim = 4                     # default
-
-  # extract dim/alpha for conv2d, and block dim
-  conv_dim = kwargs.get('conv_dim', None)
-  conv_alpha = kwargs.get('conv_alpha', None)
-  if conv_dim is not None:
-    conv_dim = int(conv_dim)
-    if conv_alpha is None:
-      conv_alpha = 1.0
-    else:
-      conv_alpha = float(conv_alpha)
-
-  """
-  block_dims = kwargs.get("block_dims")
-  block_alphas = None
-
-  if block_dims is not None:
-    block_dims = [int(d) for d in block_dims.split(',')]
-    assert len(block_dims) == NUM_BLOCKS, f"Number of block dimensions is not same to {NUM_BLOCKS}"
-    block_alphas = kwargs.get("block_alphas")
-    if block_alphas is None:
-      block_alphas = [1] * len(block_dims)
-    else:
-      block_alphas = [int(a) for a in block_alphas(',')]
-    assert len(block_alphas) == NUM_BLOCKS, f"Number of block alphas is not same to {NUM_BLOCKS}"
-
-  conv_block_dims = kwargs.get("conv_block_dims")
-  conv_block_alphas = None
-
-  if conv_block_dims is not None:
-    conv_block_dims = [int(d) for d in conv_block_dims.split(',')]
-    assert len(conv_block_dims) == NUM_BLOCKS, f"Number of block dimensions is not same to {NUM_BLOCKS}"
-    conv_block_alphas = kwargs.get("conv_block_alphas")
-    if conv_block_alphas is None:
-      conv_block_alphas = [1] * len(conv_block_dims)
-    else:
-      conv_block_alphas = [int(a) for a in conv_block_alphas(',')]
-    assert len(conv_block_alphas) == NUM_BLOCKS, f"Number of block alphas is not same to {NUM_BLOCKS}"
-  """
-
-  network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim,
-                        alpha=network_alpha, conv_lora_dim=conv_dim, conv_alpha=conv_alpha)
+  network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim, alpha=network_alpha)
   return network
 
 
-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)
-    else:
-      weights_sd = torch.load(file, map_location='cpu')
+def create_network_from_weights(multiplier, file, vae, text_encoder, unet, **kwargs):
+  if os.path.splitext(file)[1] == '.safetensors':
+    from safetensors.torch import load_file, safe_open
+    weights_sd = load_file(file)
+  else:
+    weights_sd = torch.load(file, map_location='cpu')
 
-  # get dim/alpha mapping
-  modules_dim = {}
-  modules_alpha = {}
+  # get dim (rank)
+  network_alpha = None
+  network_dim = None
   for key, value in weights_sd.items():
-    if '.' not in key:
-      continue
-
-    lora_name = key.split('.')[0]
-    if 'alpha' in key:
-      modules_alpha[lora_name] = value
-    elif 'lora_down' in key:
-      dim = value.size()[0]
-      modules_dim[lora_name] = dim
-      # print(lora_name, value.size(), dim)
-
-  # support old LoRA without alpha
-  for key in modules_dim.keys():
-    if key not in modules_alpha:
-      modules_alpha = modules_dim[key]
-
-  network = LoRANetwork(text_encoder, unet, multiplier=multiplier, modules_dim=modules_dim, modules_alpha=modules_alpha)
+    if network_alpha is None and 'alpha' in key:
+      network_alpha = value
+    if network_dim is None and 'lora_down' in key and len(value.size()) == 2:
+      network_dim = value.size()[0]
+
+  if network_alpha is None:
+    network_alpha = network_dim
+
+  network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim, alpha=network_alpha)
   network.weights_sd = weights_sd
   return network
 
 
 class LoRANetwork(torch.nn.Module):
-  # is it possible to apply conv_in and conv_out?
   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'
 
-  def __init__(self, text_encoder, unet, multiplier=1.0, lora_dim=4, alpha=1, conv_lora_dim=None, conv_alpha=None, modules_dim=None, modules_alpha=None) -> None:
+  def __init__(self, text_encoder, unet, multiplier=1.0, lora_dim=4, alpha=1) -> None:
     super().__init__()
     self.multiplier = multiplier
-
     self.lora_dim = lora_dim
     self.alpha = alpha
-    self.conv_lora_dim = conv_lora_dim
-    self.conv_alpha = conv_alpha
-
-    if modules_dim is not None:
-      print(f"create LoRA network from weights")
-    else:
-      print(f"create LoRA network. base dim (rank): {lora_dim}, alpha: {alpha}")
-
-    self.apply_to_conv2d_3x3 = self.conv_lora_dim is not None
-    if self.apply_to_conv2d_3x3:
-      if self.conv_alpha is None:
-        self.conv_alpha = self.alpha
-      print(f"apply LoRA to Conv2d with kernel size (3,3). dim (rank): {self.conv_lora_dim}, alpha: {self.conv_alpha}")
 
     # create module instances
     def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> List[LoRAModule]:
       loras = []
       for name, module in root_module.named_modules():
         if module.__class__.__name__ in target_replace_modules:
-          # TODO get block index here
           for child_name, child_module in module.named_modules():
-            is_linear = child_module.__class__.__name__ == "Linear"
-            is_conv2d = child_module.__class__.__name__ == "Conv2d"
-            is_conv2d_1x1 = is_conv2d and child_module.kernel_size == (1, 1)
-            if is_linear or is_conv2d:
+            if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)):
               lora_name = prefix + '.' + name + '.' + child_name
               lora_name = lora_name.replace('.', '_')
-
-              if modules_dim is not None:
-                if lora_name not in modules_dim:
-                  continue                                      # no LoRA module in this weights file
-                dim = modules_dim[lora_name]
-                alpha = modules_alpha[lora_name]
-              else:
-                if is_linear or is_conv2d_1x1:
-                  dim = self.lora_dim
-                  alpha = self.alpha
-                elif self.apply_to_conv2d_3x3:
-                  dim = self.conv_lora_dim
-                  alpha = self.conv_alpha
-                else:
-                  continue
-
-              lora = LoRAModule(lora_name, child_module, self.multiplier, dim, alpha)
+              lora = LoRAModule(lora_name, child_module, self.multiplier, self.lora_dim, self.alpha)
               loras.append(lora)
       return loras
 
@@ -251,12 +115,7 @@ 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.")
 
-    # 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)
+    self.unet_loras = create_modules(LoRANetwork.LORA_PREFIX_UNET, unet, LoRANetwork.UNET_TARGET_REPLACE_MODULE)
     print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
 
     self.weights_sd = None
@@ -267,11 +126,6 @@ class LoRANetwork(torch.nn.Module):
       assert lora.lora_name not in names, f"duplicated lora name: {lora.lora_name}"
       names.add(lora.lora_name)
 
-  def set_multiplier(self, multiplier):
-    self.multiplier = multiplier
-    for lora in self.text_encoder_loras + self.unet_loras:
-      lora.multiplier = self.multiplier
-
   def load_weights(self, file):
     if os.path.splitext(file)[1] == '.safetensors':
       from safetensors.torch import load_file, safe_open
@@ -381,18 +235,3 @@ class LoRANetwork(torch.nn.Module):
       save_file(state_dict, file, metadata)
     else:
       torch.save(state_dict, file)
-
-  @ staticmethod
-  def set_regions(networks, image):
-    image = image.astype(np.float32) / 255.0
-    for i, network in enumerate(networks[:3]):
-      # NOTE: consider averaging overwrapping area
-      region = image[:, :, i]
-      if region.max() == 0:
-        continue
-      region = torch.tensor(region)
-      network.set_region(region)
-
-  def set_region(self, region):
-    for lora in self.unet_loras:
-      lora.set_region(region)

networks/merge_lora.py

@@ -48,7 +48,7 @@ def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):
     for name, module in root_module.named_modules():
       if module.__class__.__name__ in target_replace_modules:
         for child_name, child_module in module.named_modules():
-          if child_module.__class__.__name__ == "Linear" or child_module.__class__.__name__ == "Conv2d":
+          if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)):
             lora_name = prefix + '.' + name + '.' + child_name
             lora_name = lora_name.replace('.', '_')
             name_to_module[lora_name] = child_module
@@ -80,19 +80,13 @@ def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):
 
         # W <- W + U * D
         weight = module.weight
-        # print(module_name, down_weight.size(), up_weight.size())
         if len(weight.size()) == 2:
           # linear
           weight = weight + ratio * (up_weight @ down_weight) * scale
-        elif down_weight.size()[2:4] == (1, 1):
-          # conv2d 1x1
+        else:
+          # conv2d
           weight = weight + ratio * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)
                                      ).unsqueeze(2).unsqueeze(3) * scale
-        else:
-          # conv2d 3x3
-          conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
-          # print(conved.size(), weight.size(), module.stride, module.padding)
-          weight = weight + ratio * conved * scale
 
         module.weight = torch.nn.Parameter(weight)
 
@@ -129,7 +123,7 @@ def merge_lora_models(models, ratios, merge_dtype):
         alphas[lora_module_name] = alpha
         if lora_module_name not in base_alphas:
           base_alphas[lora_module_name] = alpha
-
+    
     print(f"dim: {list(set(dims.values()))}, alpha: {list(set(alphas.values()))}")
 
     # merge
@@ -151,7 +145,7 @@ def merge_lora_models(models, ratios, merge_dtype):
         merged_sd[key] = merged_sd[key] + lora_sd[key] * scale
       else:
         merged_sd[key] = lora_sd[key] * scale
-
+  
   # set alpha to sd
   for lora_module_name, alpha in base_alphas.items():
     key = lora_module_name + ".alpha"

networks/resize_lora.py

@@ -1,15 +1,14 @@
 # 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
+# Thanks to cloneofsimo and kohya
 
 import argparse
+import os
 import torch
 from safetensors.torch import load_file, save_file, safe_open
 from tqdm import tqdm
 from library import train_util, model_util
-import numpy as np
 
-MIN_SV = 1e-6
 
 def load_state_dict(file_name, dtype):
   if model_util.is_safetensors(file_name):
@@ -39,149 +38,12 @@ def save_to_file(file_name, model, state_dict, dtype, metadata):
     torch.save(model, file_name)
 
 
-def index_sv_cumulative(S, target):
-  original_sum = float(torch.sum(S))
-  cumulative_sums = torch.cumsum(S, dim=0)/original_sum
-  index = int(torch.searchsorted(cumulative_sums, target)) + 1
-  if index >= len(S):
-    index = len(S) - 1
-
-  return index
-
-
-def index_sv_fro(S, target):
-  S_squared = S.pow(2)
-  s_fro_sq = float(torch.sum(S_squared))
-  sum_S_squared = torch.cumsum(S_squared, dim=0)/s_fro_sq
-  index = int(torch.searchsorted(sum_S_squared, target**2)) + 1
-  if index >= len(S):
-    index = len(S) - 1
-
-  return index
-
-
-# Modified from Kohaku-blueleaf's extract/merge functions
-def extract_conv(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1):
-    out_size, in_size, kernel_size, _ = weight.size()
-    U, S, Vh = torch.linalg.svd(weight.reshape(out_size, -1).to(device))
-    
-    param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
-    lora_rank = param_dict["new_rank"]
-
-    U = U[:, :lora_rank]
-    S = S[:lora_rank]
-    U = U @ torch.diag(S)
-    Vh = Vh[:lora_rank, :]
-
-    param_dict["lora_down"] = Vh.reshape(lora_rank, in_size, kernel_size, kernel_size).cpu()
-    param_dict["lora_up"] = U.reshape(out_size, lora_rank, 1, 1).cpu()
-    del U, S, Vh, weight
-    return param_dict
-
-
-def extract_linear(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1):
-    out_size, in_size = weight.size()
-    
-    U, S, Vh = torch.linalg.svd(weight.to(device))
-    
-    param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
-    lora_rank = param_dict["new_rank"]
-    
-    U = U[:, :lora_rank]
-    S = S[:lora_rank]
-    U = U @ torch.diag(S)
-    Vh = Vh[:lora_rank, :]
-    
-    param_dict["lora_down"] = Vh.reshape(lora_rank, in_size).cpu()
-    param_dict["lora_up"] = U.reshape(out_size, lora_rank).cpu()
-    del U, S, Vh, weight
-    return param_dict
-
-
-def merge_conv(lora_down, lora_up, device):
-    in_rank, in_size, kernel_size, k_ = lora_down.shape
-    out_size, out_rank, _, _ = lora_up.shape
-    assert in_rank == out_rank and kernel_size == k_, f"rank {in_rank} {out_rank} or kernel {kernel_size} {k_} mismatch"
-    
-    lora_down = lora_down.to(device)
-    lora_up = lora_up.to(device)
-
-    merged = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1)
-    weight = merged.reshape(out_size, in_size, kernel_size, kernel_size)
-    del lora_up, lora_down
-    return weight
-
-
-def merge_linear(lora_down, lora_up, device):
-    in_rank, in_size = lora_down.shape
-    out_size, out_rank = lora_up.shape
-    assert in_rank == out_rank, f"rank {in_rank} {out_rank} mismatch"
-    
-    lora_down = lora_down.to(device)
-    lora_up = lora_up.to(device)
-    
-    weight = lora_up @ lora_down
-    del lora_up, lora_down
-    return weight
-  
-
-def rank_resize(S, rank, dynamic_method, dynamic_param, scale=1):
-    param_dict = {}
-
-    if dynamic_method=="sv_ratio":
-        # Calculate new dim and alpha based off ratio
-        max_sv = S[0]
-        min_sv = max_sv/dynamic_param
-        new_rank = max(torch.sum(S > min_sv).item(),1)
-        new_alpha = float(scale*new_rank)
-
-    elif dynamic_method=="sv_cumulative":
-        # Calculate new dim and alpha based off cumulative sum
-        new_rank = index_sv_cumulative(S, dynamic_param)
-        new_rank = max(new_rank, 1)
-        new_alpha = float(scale*new_rank)
-
-    elif dynamic_method=="sv_fro":
-        # Calculate new dim and alpha based off sqrt sum of squares
-        new_rank = index_sv_fro(S, dynamic_param)
-        new_rank = min(max(new_rank, 1), len(S)-1)
-        new_alpha = float(scale*new_rank)
-    else:
-        new_rank = rank
-        new_alpha = float(scale*new_rank)
-
-    
-    if S[0] <= MIN_SV: # Zero matrix, set dim to 1
-        new_rank = 1
-        new_alpha = float(scale*new_rank)
-    elif new_rank > rank: # cap max rank at rank
-        new_rank = rank
-        new_alpha = float(scale*new_rank)
-
-
-    # Calculate resize info
-    s_sum = torch.sum(torch.abs(S))
-    s_rank = torch.sum(torch.abs(S[:new_rank]))
-    
-    S_squared = S.pow(2)
-    s_fro = torch.sqrt(torch.sum(S_squared))
-    s_red_fro = torch.sqrt(torch.sum(S_squared[:new_rank]))
-    fro_percent = float(s_red_fro/s_fro)
-
-    param_dict["new_rank"] = new_rank
-    param_dict["new_alpha"] = new_alpha
-    param_dict["sum_retained"] = (s_rank)/s_sum
-    param_dict["fro_retained"] = fro_percent
-    param_dict["max_ratio"] = S[0]/S[new_rank]
-
-    return param_dict
-
-
-def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dynamic_param, verbose):
+def resize_lora_model(lora_sd, new_rank, save_dtype, device, verbose):
   network_alpha = None
   network_dim = None
   verbose_str = "\n"
-  fro_list = []
+
+  CLAMP_QUANTILE = 0.99
 
   # Extract loaded lora dim and alpha
   for key, value in lora_sd.items():
@@ -195,9 +57,9 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
       network_alpha = network_dim
 
   scale = network_alpha/network_dim
+  new_alpha = float(scale*new_rank)  # calculate new alpha from scale
 
-  if dynamic_method:
-    print(f"Dynamically determining new alphas and dims based off {dynamic_method}: {dynamic_param}, max rank is {new_rank}")
+  print(f"old dimension: {network_dim}, old alpha: {network_alpha}, new alpha: {new_alpha}")
 
   lora_down_weight = None
   lora_up_weight = None
@@ -206,6 +68,7 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
   block_down_name = None
   block_up_name = None
 
+  print("resizing lora...")
   with torch.no_grad():
     for key, value in tqdm(lora_sd.items()):
       if 'lora_down' in key:
@@ -222,43 +85,57 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
         conv2d = (len(lora_down_weight.size()) == 4)
 
         if conv2d:
-          full_weight_matrix = merge_conv(lora_down_weight, lora_up_weight, device)
-          param_dict = extract_conv(full_weight_matrix, new_rank, dynamic_method, dynamic_param, device, scale)
-        else:
-          full_weight_matrix = merge_linear(lora_down_weight, lora_up_weight, device)
-          param_dict = extract_linear(full_weight_matrix, new_rank, dynamic_method, dynamic_param, device, scale)
+          lora_down_weight = lora_down_weight.squeeze()
+          lora_up_weight = lora_up_weight.squeeze()
+
+        if device:
+          org_device = lora_up_weight.device
+          lora_up_weight = lora_up_weight.to(args.device)
+          lora_down_weight = lora_down_weight.to(args.device)
+
+        full_weight_matrix = torch.matmul(lora_up_weight, lora_down_weight)
+
+        U, S, Vh = torch.linalg.svd(full_weight_matrix)
 
         if verbose:
-          max_ratio = param_dict['max_ratio']
-          sum_retained = param_dict['sum_retained']
-          fro_retained = param_dict['fro_retained']
-          if not np.isnan(fro_retained):
-            fro_list.append(float(fro_retained))
+          s_sum = torch.sum(torch.abs(S))
+          s_rank = torch.sum(torch.abs(S[:new_rank]))
+          verbose_str+=f"{block_down_name:76} | "
+          verbose_str+=f"sum(S) retained: {(s_rank)/s_sum:.1%}, max(S) ratio: {S[0]/S[new_rank]:0.1f}\n"
 
-          verbose_str+=f"{block_down_name:75} | "
-          verbose_str+=f"sum(S) retained: {sum_retained:.1%}, fro retained: {fro_retained:.1%}, max(S) ratio: {max_ratio:0.1f}"
+        U = U[:, :new_rank]
+        S = S[:new_rank]
+        U = U @ torch.diag(S)
 
-        if verbose and dynamic_method:
-          verbose_str+=f", dynamic | dim: {param_dict['new_rank']}, alpha: {param_dict['new_alpha']}\n"
-        else:
-          verbose_str+=f"\n"
+        Vh = Vh[:new_rank, :]
 
-        new_alpha = param_dict['new_alpha']
-        o_lora_sd[block_down_name + "." + "lora_down.weight"] = param_dict["lora_down"].to(save_dtype).contiguous()
-        o_lora_sd[block_up_name + "." + "lora_up.weight"] = param_dict["lora_up"].to(save_dtype).contiguous()
-        o_lora_sd[block_up_name + "." "alpha"] = torch.tensor(param_dict['new_alpha']).to(save_dtype)
+        dist = torch.cat([U.flatten(), Vh.flatten()])
+        hi_val = torch.quantile(dist, CLAMP_QUANTILE)
+        low_val = -hi_val
+
+        U = U.clamp(low_val, hi_val)
+        Vh = Vh.clamp(low_val, hi_val)
+
+        if conv2d:
+          U = U.unsqueeze(2).unsqueeze(3)
+          Vh = Vh.unsqueeze(2).unsqueeze(3)
+
+        if device:
+          U = U.to(org_device)
+          Vh = Vh.to(org_device)
+
+        o_lora_sd[block_down_name + "." + "lora_down.weight"] = Vh.to(save_dtype).contiguous()
+        o_lora_sd[block_up_name + "." + "lora_up.weight"] = U.to(save_dtype).contiguous()
+        o_lora_sd[block_up_name + "." "alpha"] = torch.tensor(new_alpha).to(save_dtype)
 
         block_down_name = None
         block_up_name = None
         lora_down_weight = None
         lora_up_weight = None
         weights_loaded = False
-        del param_dict
 
   if verbose:
     print(verbose_str)
-
-    print(f"Average Frobenius norm retention: {np.mean(fro_list):.2%} | std: {np.std(fro_list):0.3f}")
   print("resizing complete")
   return o_lora_sd, network_dim, new_alpha
 
@@ -274,9 +151,6 @@ def resize(args):
       return torch.bfloat16
     return None
 
-  if args.dynamic_method and not args.dynamic_param:
-    raise Exception("If using dynamic_method, then dynamic_param is required")
-
   merge_dtype = str_to_dtype('float')  # matmul method above only seems to work in float32
   save_dtype = str_to_dtype(args.save_precision)
   if save_dtype is None:
@@ -285,23 +159,17 @@ def resize(args):
   print("loading Model...")
   lora_sd, metadata = load_state_dict(args.model, merge_dtype)
 
-  print("Resizing Lora...")
-  state_dict, old_dim, new_alpha = resize_lora_model(lora_sd, args.new_rank, save_dtype, args.device, args.dynamic_method, args.dynamic_param, args.verbose)
+  print("resizing rank...")
+  state_dict, old_dim, new_alpha = resize_lora_model(lora_sd, args.new_rank, save_dtype, args.device, args.verbose)
 
   # update metadata
   if metadata is None:
     metadata = {}
 
   comment = metadata.get("ss_training_comment", "")
-
-  if not args.dynamic_method:
-    metadata["ss_training_comment"] = f"dimension is resized from {old_dim} to {args.new_rank}; {comment}"
-    metadata["ss_network_dim"] = str(args.new_rank)
-    metadata["ss_network_alpha"] = str(new_alpha)
-  else:
-    metadata["ss_training_comment"] = f"Dynamic resize with {args.dynamic_method}: {args.dynamic_param} from {old_dim}; {comment}"
-    metadata["ss_network_dim"] = 'Dynamic'
-    metadata["ss_network_alpha"] = 'Dynamic'
+  metadata["ss_training_comment"] = f"dimension is resized from {old_dim} to {args.new_rank}; {comment}"
+  metadata["ss_network_dim"] = str(args.new_rank)
+  metadata["ss_network_alpha"] = str(new_alpha)
 
   model_hash, legacy_hash = train_util.precalculate_safetensors_hashes(state_dict, metadata)
   metadata["sshs_model_hash"] = model_hash
@@ -325,11 +193,6 @@ if __name__ == '__main__':
   parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
   parser.add_argument("--verbose", action="store_true", 
                       help="Display verbose resizing information / rank変更時の詳細情報を出力する")
-  parser.add_argument("--dynamic_method", type=str, default=None, choices=[None, "sv_ratio", "sv_fro", "sv_cumulative"],
-                      help="Specify dynamic resizing method, --new_rank is used as a hard limit for max rank")
-  parser.add_argument("--dynamic_param", type=float, default=None,
-                      help="Specify target for dynamic reduction")
-                                           
 
   args = parser.parse_args()
   resize(args)

networks/svd_merge_lora.py

@@ -23,20 +23,19 @@ def load_state_dict(file_name, dtype):
   return sd
 
 
-def save_to_file(file_name, state_dict, dtype):
+def save_to_file(file_name, model, 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(state_dict, file_name)
+    save_file(model, file_name)
   else:
-    torch.save(state_dict, file_name)
+    torch.save(model, file_name)
 
 
-def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dtype):
-  print(f"new rank: {new_rank}, new conv rank: {new_conv_rank}")
+def merge_lora_models(models, ratios, new_rank, device,  merge_dtype):
   merged_sd = {}
   for model, ratio in zip(models, ratios):
     print(f"loading: {model}")
@@ -59,12 +58,11 @@ def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dty
       in_dim = down_weight.size()[1]
       out_dim = up_weight.size()[0]
       conv2d = len(down_weight.size()) == 4
-      kernel_size = None if not conv2d else down_weight.size()[2:4]
-      # print(lora_module_name, network_dim, alpha, in_dim, out_dim, kernel_size)
+      print(lora_module_name, network_dim, alpha, in_dim, out_dim)
 
       # make original weight if not exist
       if lora_module_name not in merged_sd:
-        weight = torch.zeros((out_dim, in_dim, *kernel_size) if conv2d else (out_dim, in_dim), dtype=merge_dtype)
+        weight = torch.zeros((out_dim, in_dim, 1, 1) if conv2d else (out_dim, in_dim), dtype=merge_dtype)
         if device:
           weight = weight.to(device)
       else:
@@ -77,18 +75,11 @@ 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):
+      else:
         weight = weight + ratio * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)
                                    ).unsqueeze(2).unsqueeze(3) * scale
-      else:
-        conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
-        weight = weight + ratio * conved * scale
 
       merged_sd[lora_module_name] = weight
 
@@ -98,26 +89,16 @@ def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dty
   with torch.no_grad():
     for lora_module_name, mat in tqdm(list(merged_sd.items())):
       conv2d = (len(mat.size()) == 4)
-      kernel_size = None if not conv2d else mat.size()[2:4]
-      conv2d_3x3 = conv2d and kernel_size != (1, 1)
-      out_dim, in_dim = mat.size()[0:2]
-
       if conv2d:
-        if conv2d_3x3:
-          mat = mat.flatten(start_dim=1)
-        else:
-          mat = mat.squeeze()
-
-      module_new_rank = new_conv_rank if conv2d_3x3 else new_rank
-      module_new_rank = min(module_new_rank, in_dim, out_dim)                           # LoRA rank cannot exceed the original dim
+        mat = mat.squeeze()
 
       U, S, Vh = torch.linalg.svd(mat)
 
-      U = U[:, :module_new_rank]
-      S = S[:module_new_rank]
+      U = U[:, :new_rank]
+      S = S[:new_rank]
       U = U @ torch.diag(S)
 
-      Vh = Vh[:module_new_rank, :]
+      Vh = Vh[:new_rank, :]
 
       dist = torch.cat([U.flatten(), Vh.flatten()])
       hi_val = torch.quantile(dist, CLAMP_QUANTILE)
@@ -126,16 +107,16 @@ def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dty
       U = U.clamp(low_val, hi_val)
       Vh = Vh.clamp(low_val, hi_val)
 
-      if conv2d:
-        U = U.reshape(out_dim, module_new_rank, 1, 1)
-        Vh = Vh.reshape(module_new_rank, in_dim, kernel_size[0], kernel_size[1])
-
       up_weight = U
       down_weight = Vh
 
+      if conv2d:
+        up_weight = up_weight.unsqueeze(2).unsqueeze(3)
+        down_weight = down_weight.unsqueeze(2).unsqueeze(3)
+
       merged_lora_sd[lora_module_name + '.lora_up.weight'] = up_weight.to("cpu").contiguous()
       merged_lora_sd[lora_module_name + '.lora_down.weight'] = down_weight.to("cpu").contiguous()
-      merged_lora_sd[lora_module_name + '.alpha'] = torch.tensor(module_new_rank)
+      merged_lora_sd[lora_module_name + '.alpha'] = torch.tensor(new_rank)
 
   return merged_lora_sd
 
@@ -157,11 +138,10 @@ def merge(args):
   if save_dtype is None:
     save_dtype = merge_dtype
 
-  new_conv_rank = args.new_conv_rank if args.new_conv_rank is not None else args.new_rank
-  state_dict = merge_lora_models(args.models, args.ratios, args.new_rank, new_conv_rank, args.device, merge_dtype)
+  state_dict = merge_lora_models(args.models, args.ratios, args.new_rank, args.device, merge_dtype)
 
   print(f"saving model to: {args.save_to}")
-  save_to_file(args.save_to, state_dict, save_dtype)
+  save_to_file(args.save_to, state_dict, state_dict, save_dtype)
 
 
 if __name__ == '__main__':
@@ -178,8 +158,6 @@ if __name__ == '__main__':
                       help="ratios for each model / それぞれのLoRAモデルの比率")
   parser.add_argument("--new_rank", type=int, default=4,
                       help="Specify rank of output LoRA / 出力するLoRAのrank (dim)")
-  parser.add_argument("--new_conv_rank", type=int, default=None,
-                      help="Specify rank of output LoRA for Conv2d 3x3, None for same as new_rank / 出力するConv2D 3x3 LoRAのrank (dim)、Noneでnew_rankと同じ")
   parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
 
   args = parser.parse_args()

requirements.txt

@@ -12,8 +12,6 @@ safetensors==0.2.6
 gradio==3.16.2
 altair==4.2.2
 easygui==0.98.3
-toml==0.10.2
-voluptuous==0.13.1
 # for BLIP captioning
 requests==2.28.2
 timm==0.6.12
@@ -23,4 +21,5 @@ fairscale==0.4.13
 tensorflow==2.10.1
 huggingface-hub==0.12.0
 # for kohya_ss library
+#locon.locon_kohya
 .

requirements_startup.txt

@@ -0,0 +1,23 @@
+accelerate==0.15.0
+transformers==4.26.0
+ftfy==6.1.1
+albumentations==1.3.0
+opencv-python==4.7.0.68
+einops==0.6.0
+diffusers[torch]==0.10.2
+pytorch-lightning==1.9.0
+bitsandbytes==0.35.0
+tensorboard==2.10.1
+safetensors==0.2.6
+gradio==3.18.0
+altair==4.2.2
+easygui==0.98.3
+# for BLIP captioning
+requests==2.28.2
+timm==0.4.12
+fairscale==0.4.4
+# for WD14 captioning
+tensorflow==2.10.1
+huggingface-hub==0.12.0
+# for kohya_ss library
+.

train_db.py

@@ -15,11 +15,7 @@ import diffusers
 from diffusers import DDPMScheduler
 
 import library.train_util as train_util
-import library.config_util as config_util
-from library.config_util import (
-  ConfigSanitizer,
-  BlueprintGenerator,
-)
+from library.train_util import DreamBoothDataset
 
 
 def collate_fn(examples):
@@ -37,33 +33,24 @@ def train(args):
 
   tokenizer = train_util.load_tokenizer(args)
 
-  blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, False, True))
-  if args.dataset_config is not None:
-    print(f"Load dataset config from {args.dataset_config}")
-    user_config = config_util.load_user_config(args.dataset_config)
-    ignored = ["train_data_dir", "reg_data_dir"]
-    if any(getattr(args, attr) is not None for attr in ignored):
-      print("ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(', '.join(ignored)))
-  else:
-    user_config = {
-      "datasets": [{
-        "subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)
-      }]
-    }
-
-  blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
-  train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
+  train_dataset = DreamBoothDataset(args.train_batch_size, args.train_data_dir, args.reg_data_dir,
+                                    tokenizer, args.max_token_length, args.caption_extension, args.shuffle_caption, args.keep_tokens,
+                                    args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                    args.bucket_reso_steps, args.bucket_no_upscale,
+                                    args.prior_loss_weight, args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop, args.debug_dataset)
 
   if args.no_token_padding:
-    train_dataset_group.disable_token_padding()
+    train_dataset.disable_token_padding()
+
+  # 学習データのdropout率を設定する
+  train_dataset.set_caption_dropout(args.caption_dropout_rate, args.caption_dropout_every_n_epochs, args.caption_tag_dropout_rate)
+
+  train_dataset.make_buckets()
 
   if args.debug_dataset:
-    train_util.debug_dataset(train_dataset_group)
+    train_util.debug_dataset(train_dataset)
     return
 
-  if cache_latents:
-    assert train_dataset_group.is_latent_cacheable(), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
-
   # acceleratorを準備する
   print("prepare accelerator")
 
@@ -104,7 +91,7 @@ def train(args):
     vae.requires_grad_(False)
     vae.eval()
     with torch.no_grad():
-      train_dataset_group.cache_latents(vae)
+      train_dataset.cache_latents(vae)
     vae.to("cpu")
     if torch.cuda.is_available():
       torch.cuda.empty_cache()
@@ -128,18 +115,38 @@ def train(args):
 
   # 学習に必要なクラスを準備する
   print("prepare optimizer, data loader etc.")
+
+  # 8-bit Adamを使う
+  if args.use_8bit_adam:
+    try:
+      import bitsandbytes as bnb
+    except ImportError:
+      raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
+    print("use 8-bit Adam optimizer")
+    optimizer_class = bnb.optim.AdamW8bit
+  elif args.use_lion_optimizer:
+    try:
+      import lion_pytorch
+    except ImportError:
+      raise ImportError("No lion_pytorch / lion_pytorch がインストールされていないようです")
+    print("use Lion optimizer")
+    optimizer_class = lion_pytorch.Lion
+  else:
+    optimizer_class = torch.optim.AdamW
+
   if train_text_encoder:
     trainable_params = (itertools.chain(unet.parameters(), text_encoder.parameters()))
   else:
     trainable_params = unet.parameters()
 
-  _, _, optimizer = train_util.get_optimizer(args, trainable_params)
+  # betaやweight decayはdiffusers DreamBoothもDreamBooth SDもデフォルト値のようなのでオプションはとりあえず省略
+  optimizer = optimizer_class(trainable_params, lr=args.learning_rate)
 
   # dataloaderを準備する
   # 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, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
 
   # 学習ステップ数を計算する
   if args.max_train_epochs is not None:
@@ -149,10 +156,9 @@ def train(args):
   if args.stop_text_encoder_training is None:
     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,
-                                              num_training_steps=args.max_train_steps,
-                                              num_cycles=args.lr_scheduler_num_cycles, power=args.lr_scheduler_power)
+  # lr schedulerを用意する
+  lr_scheduler = diffusers.optimization.get_scheduler(
+      args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps, num_training_steps=args.max_train_steps)
 
   # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
   if args.full_fp16:
@@ -189,8 +195,8 @@ def train(args):
   # 学習する
   total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
   print("running training / 学習開始")
-  print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
-  print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
+  print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset.num_train_images}")
+  print(f"  num reg images / 正則化画像の数: {train_dataset.num_reg_images}")
   print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
   print(f"  num epochs / epoch数: {num_train_epochs}")
   print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
@@ -211,7 +217,7 @@ def train(args):
   loss_total = 0.0
   for epoch in range(num_train_epochs):
     print(f"epoch {epoch+1}/{num_train_epochs}")
-    train_dataset_group.set_current_epoch(epoch + 1)
+    train_dataset.set_current_epoch(epoch + 1)
 
     # 指定したステップ数までText Encoderを学習する：epoch最初の状態
     unet.train()
@@ -275,12 +281,12 @@ def train(args):
         loss = loss.mean()                # 平均なのでbatch_sizeで割る必要なし
 
         accelerator.backward(loss)
-        if accelerator.sync_gradients and args.max_grad_norm != 0.0:
+        if accelerator.sync_gradients:
           if train_text_encoder:
             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)
+          accelerator.clip_grad_norm_(params_to_clip, 1.0)  # args.max_grad_norm)
 
         optimizer.step()
         lr_scheduler.step()
@@ -291,13 +297,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)
-
       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 = {"loss": current_loss, "lr": lr_scheduler.get_last_lr()[0]}
         accelerator.log(logs, step=global_step)
 
       if epoch == 0:
@@ -324,8 +326,6 @@ def train(args):
       train_util.save_sd_model_on_epoch_end(args, accelerator, src_path, save_stable_diffusion_format, use_safetensors,
                                             save_dtype, epoch, num_train_epochs, global_step,  unwrap_model(text_encoder), unwrap_model(unet), vae)
 
-    train_util.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
-
   is_main_process = accelerator.is_main_process
   if is_main_process:
     unet = unwrap_model(unet)
@@ -352,8 +352,6 @@ if __name__ == '__main__':
   train_util.add_dataset_arguments(parser, True, False, True)
   train_util.add_training_arguments(parser, True)
   train_util.add_sd_saving_arguments(parser)
-  train_util.add_optimizer_arguments(parser)
-  config_util.add_config_arguments(parser)
 
   parser.add_argument("--no_token_padding", action="store_true",
                       help="disable token padding (same as Diffuser's DreamBooth) / トークンのpaddingを無効にする（Diffusers版DreamBoothと同じ動作）")

train_network.py

@@ -1,4 +1,8 @@
+from diffusers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION
+from torch.optim import Optimizer
+from torch.cuda.amp import autocast
 from torch.nn.parallel import DistributedDataParallel as DDP
+from typing import Optional, Union
 import importlib
 import argparse
 import gc
@@ -11,41 +15,94 @@ import json
 from tqdm import tqdm
 import torch
 from accelerate.utils import set_seed
+import diffusers
 from diffusers import DDPMScheduler
 
 import library.train_util as train_util
-from library.train_util import (
-    DreamBoothDataset,
-)
-import library.config_util as config_util
-from library.config_util import (
-    ConfigSanitizer,
-    BlueprintGenerator,
-)
+from library.train_util import DreamBoothDataset, FineTuningDataset
 
 
 def collate_fn(examples):
   return examples[0]
 
 
-# TODO 他のスクリプトと共通化する
 def generate_step_logs(args: argparse.Namespace, current_loss, avr_loss, lr_scheduler):
   logs = {"loss/current": current_loss, "loss/average": avr_loss}
 
   if args.network_train_unet_only:
-    logs["lr/unet"] = float(lr_scheduler.get_last_lr()[0])
+    logs["lr/unet"] = lr_scheduler.get_last_lr()[0]
   elif args.network_train_text_encoder_only:
-    logs["lr/textencoder"] = float(lr_scheduler.get_last_lr()[0])
+    logs["lr/textencoder"] = lr_scheduler.get_last_lr()[0]
   else:
-    logs["lr/textencoder"] = float(lr_scheduler.get_last_lr()[0])
-    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/textencoder"] = lr_scheduler.get_last_lr()[0]
+    logs["lr/unet"] = lr_scheduler.get_last_lr()[-1]          # may be same to textencoder
 
   return logs
 
 
+# Monkeypatch newer get_scheduler() function overridng current version of diffusers.optimizer.get_scheduler
+# code is taken from https://github.com/huggingface/diffusers diffusers.optimizer, commit d87cc15977b87160c30abaace3894e802ad9e1e6
+# Which is a newer release of diffusers than currently packaged with sd-scripts
+# This code can be removed when newer diffusers version (v0.12.1 or greater) is tested and implemented to sd-scripts
+
+
+def get_scheduler_fix(
+    name: Union[str, SchedulerType],
+    optimizer: Optimizer,
+    num_warmup_steps: Optional[int] = None,
+    num_training_steps: Optional[int] = None,
+    num_cycles: int = 1,
+    power: float = 1.0,
+):
+  """
+  Unified API to get any scheduler from its name.
+  Args:
+      name (`str` or `SchedulerType`):
+          The name of the scheduler to use.
+      optimizer (`torch.optim.Optimizer`):
+          The optimizer that will be used during training.
+      num_warmup_steps (`int`, *optional*):
+          The number of warmup steps to do. This is not required by all schedulers (hence the argument being
+          optional), the function will raise an error if it's unset and the scheduler type requires it.
+      num_training_steps (`int``, *optional*):
+          The number of training steps to do. This is not required by all schedulers (hence the argument being
+          optional), the function will raise an error if it's unset and the scheduler type requires it.
+      num_cycles (`int`, *optional*):
+          The number of hard restarts used in `COSINE_WITH_RESTARTS` scheduler.
+      power (`float`, *optional*, defaults to 1.0):
+          Power factor. See `POLYNOMIAL` scheduler
+      last_epoch (`int`, *optional*, defaults to -1):
+          The index of the last epoch when resuming training.
+  """
+  name = SchedulerType(name)
+  schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]
+  if name == SchedulerType.CONSTANT:
+    return schedule_func(optimizer)
+
+  # All other schedulers require `num_warmup_steps`
+  if num_warmup_steps is None:
+    raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")
+
+  if name == SchedulerType.CONSTANT_WITH_WARMUP:
+    return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)
+
+  # All other schedulers require `num_training_steps`
+  if num_training_steps is None:
+    raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")
+
+  if name == SchedulerType.COSINE_WITH_RESTARTS:
+    return schedule_func(
+        optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, num_cycles=num_cycles
+    )
+
+  if name == SchedulerType.POLYNOMIAL:
+    return schedule_func(
+        optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, power=power
+    )
+
+  return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps)
+
+
 def train(args):
   session_id = random.randint(0, 2**32)
   training_started_at = time.time()
@@ -54,7 +111,6 @@ def train(args):
 
   cache_latents = args.cache_latents
   use_dreambooth_method = args.in_json is None
-  use_user_config = args.dataset_config is not None
 
   if args.seed is not None:
     set_seed(args.seed)
@@ -62,51 +118,38 @@ def train(args):
   tokenizer = train_util.load_tokenizer(args)
 
   # データセットを準備する
-  blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, True))
-  if use_user_config:
-    print(f"Load dataset config from {args.dataset_config}")
-    user_config = config_util.load_user_config(args.dataset_config)
-    ignored = ["train_data_dir", "reg_data_dir", "in_json"]
-    if any(getattr(args, attr) is not None for attr in ignored):
-      print(
-          "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(', '.join(ignored)))
+  if use_dreambooth_method:
+    print("Use DreamBooth method.")
+    train_dataset = DreamBoothDataset(args.train_batch_size, args.train_data_dir, args.reg_data_dir,
+                                      tokenizer, args.max_token_length, args.caption_extension, args.shuffle_caption, args.keep_tokens,
+                                      args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                      args.bucket_reso_steps, args.bucket_no_upscale,
+                                      args.prior_loss_weight, args.flip_aug, args.color_aug, args.face_crop_aug_range,
+                                      args.random_crop, args.debug_dataset)
   else:
-    if use_dreambooth_method:
-      print("Use DreamBooth method.")
-      user_config = {
-          "datasets": [{
-              "subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)
-          }]
-      }
-    else:
-      print("Train with captions.")
-      user_config = {
-          "datasets": [{
-              "subsets": [{
-                  "image_dir": args.train_data_dir,
-                  "metadata_file": args.in_json,
-              }]
-          }]
-      }
-
-  blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
-  train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
+    print("Train with captions.")
+    train_dataset = FineTuningDataset(args.in_json, args.train_batch_size, args.train_data_dir,
+                                      tokenizer, args.max_token_length, args.shuffle_caption, args.keep_tokens,
+                                      args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                      args.bucket_reso_steps, args.bucket_no_upscale,
+                                      args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop,
+                                      args.dataset_repeats, args.debug_dataset)
+
+  # 学習データのdropout率を設定する
+  train_dataset.set_caption_dropout(args.caption_dropout_rate, args.caption_dropout_every_n_epochs, args.caption_tag_dropout_rate)
+
+  train_dataset.make_buckets()
 
   if args.debug_dataset:
-    train_util.debug_dataset(train_dataset_group)
+    train_util.debug_dataset(train_dataset)
     return
-  if len(train_dataset_group) == 0:
+  if len(train_dataset) == 0:
     print("No data found. Please verify arguments (train_data_dir must be the parent of folders with images) / 画像がありません。引数指定を確認してください（train_data_dirには画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）")
     return
 
-  if cache_latents:
-    assert train_dataset_group.is_latent_cacheable(
-    ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
-
   # acceleratorを準備する
   print("prepare accelerator")
   accelerator, unwrap_model = train_util.prepare_accelerator(args)
-  is_main_process = accelerator.is_main_process
 
   # mixed precisionに対応した型を用意しておき適宜castする
   weight_dtype, save_dtype = train_util.prepare_dtype(args)
@@ -118,7 +161,7 @@ def train(args):
   if args.lowram:
     text_encoder.to("cuda")
     unet.to("cuda")
-
+  
   # モデルに xformers とか memory efficient attention を組み込む
   train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
 
@@ -128,15 +171,13 @@ def train(args):
     vae.requires_grad_(False)
     vae.eval()
     with torch.no_grad():
-      train_dataset_group.cache_latents(vae)
+      train_dataset.cache_latents(vae)
     vae.to("cpu")
     if torch.cuda.is_available():
       torch.cuda.empty_cache()
     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)
 
@@ -167,25 +208,48 @@ def train(args):
   # 学習に必要なクラスを準備する
   print("prepare optimizer, data loader etc.")
 
+  # 8-bit Adamを使う
+  if args.use_8bit_adam:
+    try:
+      import bitsandbytes as bnb
+    except ImportError:
+      raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
+    print("use 8-bit Adam optimizer")
+    optimizer_class = bnb.optim.AdamW8bit
+  elif args.use_lion_optimizer:
+    try:
+      import lion_pytorch
+    except ImportError:
+      raise ImportError("No lion_pytorch / lion_pytorch がインストールされていないようです")
+    print("use Lion optimizer")
+    optimizer_class = lion_pytorch.Lion
+  else:
+    optimizer_class = torch.optim.AdamW
+
+  optimizer_name = optimizer_class.__module__ + "." + optimizer_class.__name__
+
   trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr)
-  optimizer_name, optimizer_args, optimizer = train_util.get_optimizer(args, trainable_params)
+
+  # betaやweight decayはdiffusers DreamBoothもDreamBooth SDもデフォルト値のようなのでオプションはとりあえず省略
+  optimizer = optimizer_class(trainable_params, lr=args.learning_rate)
 
   # dataloaderを準備する
   # 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, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
 
   # 学習ステップ数を計算する
   if args.max_train_epochs is not None:
-    args.max_train_steps = args.max_train_epochs * math.ceil(len(train_dataloader) / accelerator.num_processes)
-    if is_main_process:
-      print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
+    args.max_train_steps = args.max_train_epochs * len(train_dataloader)
+    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,
-                                              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)
+  # lr_scheduler = diffusers.optimization.get_scheduler(
+  lr_scheduler = 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)
 
   # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
   if args.full_fp16:
@@ -253,21 +317,17 @@ def train(args):
     args.save_every_n_epochs = math.floor(num_train_epochs / args.save_n_epoch_ratio) or 1
 
   # 学習する
-  # TODO: find a way to handle total batch size when there are multiple datasets
   total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
-  
-  if is_main_process:
-    print("running training / 学習開始")
-    print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
-    print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
-    print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
-    print(f"  num epochs / epoch数: {num_train_epochs}")
-    print(f"  batch size per device / バッチサイズ: {', '.join([str(d.batch_size) for d in train_dataset_group.datasets])}")
-    # print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
-    print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
-    print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
-
-  # TODO refactor metadata creation and move to util
+  print("running training / 学習開始")
+  print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset.num_train_images}")
+  print(f"  num reg images / 正則化画像の数: {train_dataset.num_reg_images}")
+  print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
+  print(f"  num epochs / epoch数: {num_train_epochs}")
+  print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
+  print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
+  print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
+  print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
+
   metadata = {
       "ss_session_id": session_id,            # random integer indicating which group of epochs the model came from
       "ss_training_started_at": training_started_at,          # unix timestamp
@@ -275,10 +335,12 @@ def train(args):
       "ss_learning_rate": args.learning_rate,
       "ss_text_encoder_lr": args.text_encoder_lr,
       "ss_unet_lr": args.unet_lr,
-      "ss_num_train_images": train_dataset_group.num_train_images,
-      "ss_num_reg_images": train_dataset_group.num_reg_images,
+      "ss_num_train_images": train_dataset.num_train_images,          # includes repeating
+      "ss_num_reg_images": train_dataset.num_reg_images,
       "ss_num_batches_per_epoch": len(train_dataloader),
       "ss_num_epochs": num_train_epochs,
+      "ss_batch_size_per_device": args.train_batch_size,
+      "ss_total_batch_size": total_batch_size,
       "ss_gradient_checkpointing": args.gradient_checkpointing,
       "ss_gradient_accumulation_steps": args.gradient_accumulation_steps,
       "ss_max_train_steps": args.max_train_steps,
@@ -290,156 +352,33 @@ def train(args):
       "ss_mixed_precision": args.mixed_precision,
       "ss_full_fp16": bool(args.full_fp16),
       "ss_v2": bool(args.v2),
+      "ss_resolution": args.resolution,
       "ss_clip_skip": args.clip_skip,
       "ss_max_token_length": args.max_token_length,
+      "ss_color_aug": bool(args.color_aug),
+      "ss_flip_aug": bool(args.flip_aug),
+      "ss_random_crop": bool(args.random_crop),
+      "ss_shuffle_caption": bool(args.shuffle_caption),
       "ss_cache_latents": bool(args.cache_latents),
+      "ss_enable_bucket": bool(train_dataset.enable_bucket),
+      "ss_min_bucket_reso": train_dataset.min_bucket_reso,
+      "ss_max_bucket_reso": train_dataset.max_bucket_reso,
       "ss_seed": args.seed,
-      "ss_lowram": args.lowram,
+      "ss_keep_tokens": args.keep_tokens,
       "ss_noise_offset": args.noise_offset,
+      "ss_dataset_dirs": json.dumps(train_dataset.dataset_dirs_info),
+      "ss_reg_dataset_dirs": json.dumps(train_dataset.reg_dataset_dirs_info),
+      "ss_tag_frequency": json.dumps(train_dataset.tag_frequency),
+      "ss_bucket_info": json.dumps(train_dataset.bucket_info),
       "ss_training_comment": args.training_comment,       # will not be updated after training
       "ss_sd_scripts_commit_hash": train_util.get_git_revision_hash(),
-      "ss_optimizer": optimizer_name + (f"({optimizer_args})" if len(optimizer_args) > 0 else ""),
-      "ss_max_grad_norm": args.max_grad_norm,
-      "ss_caption_dropout_rate": args.caption_dropout_rate,
-      "ss_caption_dropout_every_n_epochs": args.caption_dropout_every_n_epochs,
-      "ss_caption_tag_dropout_rate": args.caption_tag_dropout_rate,
-      "ss_face_crop_aug_range": args.face_crop_aug_range,
-      "ss_prior_loss_weight": args.prior_loss_weight,
+      "ss_optimizer": optimizer_name
   }
 
-  if use_user_config:
-    # save metadata of multiple datasets
-    # NOTE: pack "ss_datasets" value as json one time
-    #   or should also pack nested collections as json?
-    datasets_metadata = []
-    tag_frequency = {}                    # merge tag frequency for metadata editor
-    dataset_dirs_info = {}                # merge subset dirs for metadata editor
-
-    for dataset in train_dataset_group.datasets:
-      is_dreambooth_dataset = isinstance(dataset, DreamBoothDataset)
-      dataset_metadata = {
-          "is_dreambooth": is_dreambooth_dataset,
-          "batch_size_per_device": dataset.batch_size,
-          "num_train_images": dataset.num_train_images,          # includes repeating
-          "num_reg_images": dataset.num_reg_images,
-          "resolution": (dataset.width, dataset.height),
-          "enable_bucket": bool(dataset.enable_bucket),
-          "min_bucket_reso": dataset.min_bucket_reso,
-          "max_bucket_reso": dataset.max_bucket_reso,
-          "tag_frequency": dataset.tag_frequency,
-          "bucket_info": dataset.bucket_info,
-      }
-
-      subsets_metadata = []
-      for subset in dataset.subsets:
-        subset_metadata = {
-            "img_count": subset.img_count,
-            "num_repeats": subset.num_repeats,
-            "color_aug": bool(subset.color_aug),
-            "flip_aug": bool(subset.flip_aug),
-            "random_crop": bool(subset.random_crop),
-            "shuffle_caption": bool(subset.shuffle_caption),
-            "keep_tokens": subset.keep_tokens,
-        }
-
-        image_dir_or_metadata_file = None
-        if subset.image_dir:
-          image_dir = os.path.basename(subset.image_dir)
-          subset_metadata["image_dir"] = image_dir
-          image_dir_or_metadata_file = image_dir
-
-        if is_dreambooth_dataset:
-          subset_metadata["class_tokens"] = subset.class_tokens
-          subset_metadata["is_reg"] = subset.is_reg
-          if subset.is_reg:
-            image_dir_or_metadata_file = None                    # not merging reg dataset
-        else:
-          metadata_file = os.path.basename(subset.metadata_file)
-          subset_metadata["metadata_file"] = metadata_file
-          image_dir_or_metadata_file = metadata_file           # may overwrite
-
-        subsets_metadata.append(subset_metadata)
-
-        # merge dataset dir: not reg subset only
-        # TODO update additional-network extension to show detailed dataset config from metadata
-        if image_dir_or_metadata_file is not None:
-          # datasets may have a certain dir multiple times
-          v = image_dir_or_metadata_file
-          i = 2
-          while v in dataset_dirs_info:
-            v = image_dir_or_metadata_file + f" ({i})"
-            i += 1
-          image_dir_or_metadata_file = v
-
-          dataset_dirs_info[image_dir_or_metadata_file] = {
-              "n_repeats": subset.num_repeats,
-              "img_count": subset.img_count
-          }
-
-      dataset_metadata["subsets"] = subsets_metadata
-      datasets_metadata.append(dataset_metadata)
-
-      # merge tag frequency:
-      for ds_dir_name, ds_freq_for_dir in dataset.tag_frequency.items():
-        # あるディレクトリが複数のdatasetで使用されている場合、一度だけ数える
-        # もともと繰り返し回数を指定しているので、キャプション内でのタグの出現回数と、それが学習で何度使われるかは一致しない
-        # なので、ここで複数datasetの回数を合算してもあまり意味はない
-        if ds_dir_name in tag_frequency:
-          continue
-        tag_frequency[ds_dir_name] = ds_freq_for_dir
-
-    metadata["ss_datasets"] = json.dumps(datasets_metadata)
-    metadata["ss_tag_frequency"] = json.dumps(tag_frequency)
-    metadata["ss_dataset_dirs"] = json.dumps(dataset_dirs_info)
-  else:
-    # conserving backward compatibility when using train_dataset_dir and reg_dataset_dir
-    assert len(
-        train_dataset_group.datasets) == 1, f"There should be a single dataset but {len(train_dataset_group.datasets)} found. This seems to be a bug. / データセットは1個だけ存在するはずですが、実際には{len(train_dataset_group.datasets)}個でした。プログラムのバグかもしれません。"
-
-    dataset = train_dataset_group.datasets[0]
-
-    dataset_dirs_info = {}
-    reg_dataset_dirs_info = {}
-    if use_dreambooth_method:
-      for subset in dataset.subsets:
-        info = reg_dataset_dirs_info if subset.is_reg else dataset_dirs_info
-        info[os.path.basename(subset.image_dir)] = {
-            "n_repeats": subset.num_repeats,
-            "img_count": subset.img_count
-        }
-    else:
-      for subset in dataset.subsets:
-        dataset_dirs_info[os.path.basename(subset.metadata_file)] = {
-            "n_repeats": subset.num_repeats,
-            "img_count": subset.img_count
-        }
-
-    metadata.update({
-        "ss_batch_size_per_device": args.train_batch_size,
-        "ss_total_batch_size": total_batch_size,
-        "ss_resolution": args.resolution,
-        "ss_color_aug": bool(args.color_aug),
-        "ss_flip_aug": bool(args.flip_aug),
-        "ss_random_crop": bool(args.random_crop),
-        "ss_shuffle_caption": bool(args.shuffle_caption),
-        "ss_enable_bucket": bool(dataset.enable_bucket),
-        "ss_bucket_no_upscale": bool(dataset.bucket_no_upscale),
-        "ss_min_bucket_reso": dataset.min_bucket_reso,
-        "ss_max_bucket_reso": dataset.max_bucket_reso,
-        "ss_keep_tokens": args.keep_tokens,
-        "ss_dataset_dirs": json.dumps(dataset_dirs_info),
-        "ss_reg_dataset_dirs": json.dumps(reg_dataset_dirs_info),
-        "ss_tag_frequency": json.dumps(dataset.tag_frequency),
-        "ss_bucket_info": json.dumps(dataset.bucket_info),
-    })
-
-  # add extra args
-  if args.network_args:
-    metadata["ss_network_args"] = json.dumps(net_kwargs)
-    # for key, value in net_kwargs.items():
-    #   metadata["ss_arg_" + key] = value
-
-  # model name and hash
+  # uncomment if another network is added
+  # for key, value in net_kwargs.items():
+  #   metadata["ss_arg_" + key] = value
+
   if args.pretrained_model_name_or_path is not None:
     sd_model_name = args.pretrained_model_name_or_path
     if os.path.exists(sd_model_name):
@@ -458,13 +397,6 @@ def train(args):
 
   metadata = {k: str(v) for k, v in metadata.items()}
 
-  # make minimum metadata for filtering
-  minimum_keys = ["ss_network_module", "ss_network_dim", "ss_network_alpha", "ss_network_args"]
-  minimum_metadata = {}
-  for key in minimum_keys:
-    if key in metadata:
-      minimum_metadata[key] = metadata[key]
-
   progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
   global_step = 0
 
@@ -477,9 +409,8 @@ def train(args):
   loss_list = []
   loss_total = 0.0
   for epoch in range(num_train_epochs):
-    if is_main_process:
-      print(f"epoch {epoch+1}/{num_train_epochs}")
-    train_dataset_group.set_current_epoch(epoch + 1)
+    print(f"epoch {epoch+1}/{num_train_epochs}")
+    train_dataset.set_current_epoch(epoch + 1)
 
     metadata["ss_epoch"] = str(epoch+1)
 
@@ -516,7 +447,7 @@ def train(args):
         noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
 
         # Predict the noise residual
-        with accelerator.autocast():
+        with autocast():
           noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
 
         if args.v_parameterization:
@@ -534,9 +465,9 @@ def train(args):
         loss = loss.mean()                # 平均なのでbatch_sizeで割る必要なし
 
         accelerator.backward(loss)
-        if accelerator.sync_gradients and args.max_grad_norm != 0.0:
+        if accelerator.sync_gradients:
           params_to_clip = network.get_trainable_params()
-          accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+          accelerator.clip_grad_norm_(params_to_clip, 1.0)  # args.max_grad_norm)
 
         optimizer.step()
         lr_scheduler.step()
@@ -547,8 +478,6 @@ def train(args):
         progress_bar.update(1)
         global_step += 1
 
-        train_util.sample_images(accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
-
       current_loss = loss.detach().item()
       if epoch == 0:
         loss_list.append(current_loss)
@@ -579,9 +508,8 @@ def train(args):
       def save_func():
         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)
+        unwrap_model(network).save_weights(ckpt_file, save_dtype, None 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
@@ -590,18 +518,15 @@ def train(args):
           print(f"removing old checkpoint: {old_ckpt_file}")
           os.remove(old_ckpt_file)
 
-      if is_main_process:
-        saving = train_util.save_on_epoch_end(args, save_func, remove_old_func, epoch + 1, num_train_epochs)
-        if saving and args.save_state:
-          train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)
-
-    train_util.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
+      saving = train_util.save_on_epoch_end(args, save_func, remove_old_func, epoch + 1, num_train_epochs)
+      if saving and args.save_state:
+        train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)
 
     # end of epoch
 
   metadata["ss_epoch"] = str(num_train_epochs)
-  metadata["ss_training_finished_at"] = str(time.time())
 
+  is_main_process = accelerator.is_main_process
   if is_main_process:
     network = unwrap_model(network)
 
@@ -620,7 +545,7 @@ def train(args):
     ckpt_file = os.path.join(args.output_dir, ckpt_name)
 
     print(f"save trained model to {ckpt_file}")
-    network.save_weights(ckpt_file, save_dtype, minimum_metadata if args.no_metadata else metadata)
+    network.save_weights(ckpt_file, save_dtype, None if args.no_metadata else metadata)
     print("model saved.")
 
 
@@ -630,8 +555,6 @@ if __name__ == '__main__':
   train_util.add_sd_models_arguments(parser)
   train_util.add_dataset_arguments(parser, True, True, True)
   train_util.add_training_arguments(parser, True)
-  train_util.add_optimizer_arguments(parser)
-  config_util.add_config_arguments(parser)
 
   parser.add_argument("--no_metadata", action='store_true', help="do not save metadata in output model / メタデータを出力先モデルに保存しない")
   parser.add_argument("--save_model_as", type=str, default="safetensors", choices=[None, "ckpt", "pt", "safetensors"],
@@ -639,6 +562,10 @@ if __name__ == '__main__':
 
   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("--lr_scheduler_num_cycles", type=int, default=1,
+                      help="Number of restarts for cosine scheduler with restarts / cosine with restartsスケジューラでのリスタート回数")
+  parser.add_argument("--lr_scheduler_power", type=float, default=1,
+                      help="Polynomial power for polynomial scheduler / polynomialスケジューラでのpolynomial power")
 
   parser.add_argument("--network_weights", type=str, default=None,
                       help="pretrained weights for network / 学習するネットワークの初期重み")

train_network_opt.py

@@ -1,5 +1,8 @@
+from diffusers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION
+from torch.optim import Optimizer
 from torch.cuda.amp import autocast
 from torch.nn.parallel import DistributedDataParallel as DDP
+from typing import Optional, Union
 import importlib
 import argparse
 import gc
@@ -12,49 +15,138 @@ import json
 from tqdm import tqdm
 import torch
 from accelerate.utils import set_seed
-#import diffusers
+import diffusers
 from diffusers import DDPMScheduler
-
+print("**********************************")
+#先に
+#pip install torch_optimizer
+#が必要
+try:
+  import torch_optimizer as optim
+except:
+  print("torch_optimizerがインストールされていないためAdafactorとAdastand以外の追加optimzierは使えません。\noptimizerの変更をしたい場合先にpip install torch_optimizerでライブラリを追加してください")
+try:
+  import adastand
+except:
+  print("※Adastandが使えません")
+
+from transformers.optimization import Adafactor, AdafactorSchedule
+print("**********************************")
 ##### バケット拡張のためのモジュール
 import append_module
 ######
 import library.train_util as train_util
-from library.train_util import (
-    DreamBoothDataset,
-)
-import library.config_util as config_util
-from library.config_util import (
-    ConfigSanitizer,
-    BlueprintGenerator,
-)
+from library.train_util import DreamBoothDataset, FineTuningDataset
 
 
 def collate_fn(examples):
   return examples[0]
 
 
-# TODO 他のスクリプトと共通化する
-def generate_step_logs(args: argparse.Namespace, current_loss, avr_loss, lr_scheduler, split_names=None):
+def generate_step_logs(args: argparse.Namespace, current_loss, avr_loss, lr_scheduler):
   logs = {"loss/current": current_loss, "loss/average": avr_loss}
-  if not args.split_lora_networks:
-    if args.network_train_unet_only:
-      logs["lr/unet"] = float(lr_scheduler.get_last_lr()[0])
-    elif args.network_train_text_encoder_only:
-      logs["lr/textencoder"] = float(lr_scheduler.get_last_lr()[0])
-    else:
-      logs["lr/textencoder"] = float(lr_scheduler.get_last_lr()[0])
-      logs["lr/unet"] = float(lr_scheduler.get_last_lr()[-1])          # may be same to textencoder
+
+  if args.network_train_unet_only:
+    logs["lr/unet"] = lr_scheduler.get_last_lr()[0]
+  elif args.network_train_text_encoder_only:
+    logs["lr/textencoder"] = lr_scheduler.get_last_lr()[0]
   else:
     last_lrs = lr_scheduler.get_last_lr()
-    for last_lr, t_name in zip(last_lrs, split_names):
-      logs[f"lr/{t_name}"] = float(last_lr)
-  #D-Adaptationの仕様ちゃんと見てないからたぶん分割したのをちゃんと表示するならそれに合わせた記述が必要　でも多分D-Adaptationの挙動的に全部同一の形になるのでいらない
-  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']
+    if len(last_lrs) == 2:
+      logs["lr/textencoder"] = float(last_lrs[0])
+      logs["lr/unet"] = float(last_lrs[-1])          # may be same to textencoder
+    else:
+      if len(last_lrs) == 4:
+        logs_names = ["textencoder", "lora_unet_mid_block", "unet_down_blocks", "unet_up_blocks"]
+      elif len(last_lrs) == 8:
+        logs_names = ["textencoder", "unet_midblock"]
+        for i in range(3):
+          logs_names.append(f"unet_down_blocks_{i}")
+          logs_names.append(f"unet_up_blocks_{i+1}")
+      else:
+        logs_names = []
+        for i in range(12):
+          logs_names.append(f"text_model_encoder_layers_{i}_")
+        logs_names.append("unet_midblock")
+        for i in range(3):
+          logs_names.append(f"unet_down_blocks_{i}")
+          logs_names.append(f"unet_up_blocks_{i+1}")
+      
+      for last_lr, logs_name in zip(last_lrs, logs_names):
+        logs[f"lr/{logs_name}"] = float(last_lr)
 
   return logs
 
 
+# Monkeypatch newer get_scheduler() function overridng current version of diffusers.optimizer.get_scheduler
+# code is taken from https://github.com/huggingface/diffusers diffusers.optimizer, commit d87cc15977b87160c30abaace3894e802ad9e1e6
+# Which is a newer release of diffusers than currently packaged with sd-scripts
+# This code can be removed when newer diffusers version (v0.12.1 or greater) is tested and implemented to sd-scripts
+
+
+def get_scheduler_fix(
+    name: Union[str, SchedulerType],
+    optimizer: Optimizer,
+    num_warmup_steps: Optional[int] = None,
+    num_training_steps: Optional[int] = None,
+    num_cycles: float = 1.,
+    power: float = 1.0,
+):
+  """
+  Unified API to get any scheduler from its name.
+  Args:
+      name (`str` or `SchedulerType`):
+          The name of the scheduler to use.
+      optimizer (`torch.optim.Optimizer`):
+          The optimizer that will be used during training.
+      num_warmup_steps (`int`, *optional*):
+          The number of warmup steps to do. This is not required by all schedulers (hence the argument being
+          optional), the function will raise an error if it's unset and the scheduler type requires it.
+      num_training_steps (`int``, *optional*):
+          The number of training steps to do. This is not required by all schedulers (hence the argument being
+          optional), the function will raise an error if it's unset and the scheduler type requires it.
+      num_cycles (`int`, *optional*):
+          The number of hard restarts used in `COSINE_WITH_RESTARTS` scheduler.
+      power (`float`, *optional*, defaults to 1.0):
+          Power factor. See `POLYNOMIAL` scheduler
+      last_epoch (`int`, *optional*, defaults to -1):
+          The index of the last epoch when resuming training.
+  """
+  name = SchedulerType(name)
+  schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]
+  if name == SchedulerType.CONSTANT:
+    return schedule_func(optimizer)
+
+  # All other schedulers require `num_warmup_steps`
+  if num_warmup_steps is None:
+    raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")
+
+  if name == SchedulerType.CONSTANT_WITH_WARMUP:
+    return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)
+
+  # All other schedulers require `num_training_steps`
+  if num_training_steps is None:
+    raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")
+
+  if name == SchedulerType.COSINE:
+      print(f"{name} num_cycles: {num_cycles}")
+      return schedule_func(
+          optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, num_cycles=num_cycles
+      )
+  if name == SchedulerType.COSINE_WITH_RESTARTS:
+      print(f"{name} num_cycles: {int(num_cycles)}")
+      return schedule_func(
+          optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, num_cycles=int(num_cycles)
+      )
+
+  if name == SchedulerType.POLYNOMIAL:
+    return schedule_func(
+        optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, power=power
+    )
+
+  return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps)
+
+
 def train(args):
   session_id = random.randint(0, 2**32)
   training_started_at = time.time()
@@ -63,7 +155,6 @@ def train(args):
 
   cache_latents = args.cache_latents
   use_dreambooth_method = args.in_json is None
-  use_user_config = args.dataset_config is not None
 
   if args.seed is not None:
     set_seed(args.seed)
@@ -71,72 +162,52 @@ def train(args):
   tokenizer = train_util.load_tokenizer(args)
 
   # データセットを準備する
-  if args.min_resolution:
-    args.min_resolution = tuple([int(r) for r in args.min_resolution.split(',')])
-    if len(args.min_resolution) == 1:
-      args.min_resolution = (args.min_resolution[0], args.min_resolution[0])
-    blueprint_generator = append_module.BlueprintGenerator(append_module.ConfigSanitizer(True, True, True))
-  else:
-    blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, True))
-  if use_user_config:
-    print(f"Load dataset config from {args.dataset_config}")
-    user_config = config_util.load_user_config(args.dataset_config)
-    ignored = ["train_data_dir", "reg_data_dir", "in_json"]
-    if any(getattr(args, attr) is not None for attr in ignored):
-      print(
-          "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(', '.join(ignored)))
-  else:
-    if use_dreambooth_method:
-      print("Use DreamBooth method.")
-      user_config = {
-          "datasets": [{
-              "subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)
-          }]
-      }
-    else:
-      print("Train with captions.")
-      user_config = {
-          "datasets": [{
-              "subsets": [{
-                  "image_dir": args.train_data_dir,
-                  "metadata_file": args.in_json,
-              }]
-          }]
-      }
-
-  blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
-  if args.min_resolution:
-    train_dataset_group = append_module.generate_dataset_group_by_blueprint(blueprint.dataset_group)
+  if use_dreambooth_method:
+    if args.min_resolution:
+      args.min_resolution = tuple([int(r) for r in args.min_resolution.split(',')])
+      if len(args.min_resolution) == 1:
+        args.min_resolution = (args.min_resolution[0], args.min_resolution[0])
+
+    print("Use DreamBooth method.")
+    train_dataset = append_module.DreamBoothDataset(args.train_batch_size, args.train_data_dir, args.reg_data_dir,
+                                      tokenizer, args.max_token_length, args.caption_extension, args.shuffle_caption, args.keep_tokens,
+                                      args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso, 
+                                      args.bucket_reso_steps, args.bucket_no_upscale, 
+                                      args.prior_loss_weight, args.flip_aug, args.color_aug, args.face_crop_aug_range, 
+                                      args.random_crop, args.debug_dataset, args.min_resolution, args.area_step)
   else:
-    train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
+    print("Train with captions.")
+    train_dataset = FineTuningDataset(args.in_json, args.train_batch_size, args.train_data_dir,
+                                      tokenizer, args.max_token_length, args.shuffle_caption, args.keep_tokens,
+                                      args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                      args.bucket_reso_steps, args.bucket_no_upscale, 
+                                      args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop,
+                                      args.dataset_repeats, args.debug_dataset)
+
+  # 学習データのdropout率を設定する
+  train_dataset.set_caption_dropout(args.caption_dropout_rate, args.caption_dropout_every_n_epochs, args.caption_tag_dropout_rate)
+
+  train_dataset.make_buckets()
 
   if args.debug_dataset:
-    train_util.debug_dataset(train_dataset_group)
+    train_util.debug_dataset(train_dataset)
     return
-  if len(train_dataset_group) == 0:
+  if len(train_dataset) == 0:
     print("No data found. Please verify arguments (train_data_dir must be the parent of folders with images) / 画像がありません。引数指定を確認してください（train_data_dirには画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）")
     return
 
-  if cache_latents:
-    assert train_dataset_group.is_latent_cacheable(
-    ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
-
   # acceleratorを準備する
   print("prepare accelerator")
   accelerator, unwrap_model = train_util.prepare_accelerator(args)
-  is_main_process = accelerator.is_main_process
 
   # mixed precisionに対応した型を用意しておき適宜castする
   weight_dtype, save_dtype = train_util.prepare_dtype(args)
 
   # モデルを読み込む
   text_encoder, vae, unet, _ = train_util.load_target_model(args, weight_dtype)
-
-  # work on low-ram device
-  if args.lowram:
-    text_encoder.to("cuda")
-    unet.to("cuda")
-
+  # unnecessary, but work on low-ram device
+  text_encoder.to("cuda")
+  unet.to("cuda")
   # モデルに xformers とか memory efficient attention を組み込む
   train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
 
@@ -146,15 +217,13 @@ def train(args):
     vae.requires_grad_(False)
     vae.eval()
     with torch.no_grad():
-      train_dataset_group.cache_latents(vae)
+      train_dataset.cache_latents(vae)
     vae.to("cpu")
     if torch.cuda.is_available():
       torch.cuda.empty_cache()
     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)
 
@@ -184,65 +253,188 @@ def train(args):
 
   # 学習に必要なクラスを準備する
   print("prepare optimizer, data loader etc.")
-  split_flag = (args.split_lora_networks) or ((not args.network_train_text_encoder_only) and (not args.network_train_unet_only))
-
-  used_names = None
+  try:
+    print(f"torch_optimzier version is {optim.__version__}")
+    not_torch_optimizer_flag = False
+  except:
+    not_torch_optimizer_flag = True
+  try:
+    print(f"adastand version is {adastand.__version__()}")
+    not_adasatand_optimzier_flag = False
+  except:
+    not_adasatand_optimzier_flag = True
+
+  # 8-bit Adamを使う
+  if args.optimizer=="Adafactor" or args.optimizer=="Adastand" or args.optimizer=="Adastand_belief":
+    not_torch_optimizer_flag = False
+    if args.optimizer=="Adafactor":
+      not_adasatand_optimzier_flag = False
+  if not_torch_optimizer_flag or not_adasatand_optimzier_flag:
+    print(f"==========================\n必要なライブラリがないため {args.optimizer} の使用ができません。optimizerを AdamW に変更して実行します\n==========================")
+    args.optimizer="AdamW"
+  if args.use_8bit_adam:
+    if not args.optimizer=="AdamW" and not args.optimizer=="Lamb":
+      print(f"\n==========================\n{args.optimizer} は8bitAdamに実装されていないので8bitAdamをオフにします\n==========================\n")
+      args.use_8bit_adam=False
+  if args.use_8bit_adam:
+    try:
+      import bitsandbytes as bnb
+    except ImportError:
+      raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
+    print("use 8-bit Adam optimizer")
+    args.training_comment=f"{args.training_comment} use_8bit_adam=True"
+    if args.optimizer=="Lamb":
+      optimizer_class = bnb.optim.LAMB8bit
+    else:
+      args.optimizer="AdamW"
+      optimizer_class = bnb.optim.AdamW8bit
+  else:
+    print(f"use {args.optimizer}")
+    if args.optimizer=="RAdam":
+      optimizer_class = torch.optim.RAdam
+    elif args.optimizer=="AdaBound":
+      optimizer_class = optim.AdaBound
+    elif args.optimizer=="AdaBelief":
+      optimizer_class = optim.AdaBelief
+    elif args.optimizer=="AdamP":
+      optimizer_class = optim.AdamP
+    elif args.optimizer=="Adafactor":
+      optimizer_class = Adafactor
+    elif args.optimizer=="Adastand":
+      optimizer_class = adastand.Adastand
+    elif args.optimizer=="Adastand_belief":
+      optimizer_class = adastand.Adastand_b
+    elif args.optimizer=="AggMo":
+      optimizer_class = optim.AggMo
+    elif args.optimizer=="Apollo":
+      optimizer_class = optim.Apollo
+    elif args.optimizer=="Lamb":
+      optimizer_class = optim.Lamb
+    elif args.optimizer=="Ranger":
+      optimizer_class = optim.Ranger
+    elif args.optimizer=="RangerVA":
+      optimizer_class = optim.RangerVA
+    elif args.optimizer=="Yogi":
+      optimizer_class = optim.Yogi
+    elif args.optimizer=="Shampoo":
+      optimizer_class = optim.Shampoo
+    elif args.optimizer=="NovoGrad":
+      optimizer_class = optim.NovoGrad
+    elif args.optimizer=="QHAdam":
+      optimizer_class = optim.QHAdam
+    elif args.optimizer=="DiffGrad" or args.optimizer=="Lookahead_DiffGrad":
+      optimizer_class = optim.DiffGrad
+    elif args.optimizer=="MADGRAD":
+      optimizer_class = optim.MADGRAD
+    else:
+      optimizer_class = torch.optim.AdamW
+
+  # betaやweight decayはdiffusers DreamBoothもDreamBooth SDもデフォルト値のようなのでオプションはとりあえず省略
+  #optimizerデフォ設定
+  if args.optimizer_arg==None:
+    if args.optimizer=="AdaBelief":
+      args.optimizer_arg = ["eps=1e-16","betas=0.9,0.999","weight_decouple=True","rectify=False","fixed_decay=False"]
+    elif args.optimizer=="DiffGrad":
+      args.optimizer_arg = ["eps=1e-16"]
+  optimizer_arg = {}
+  lookahed_arg = {"k": 5, "alpha": 0.5}
+  adafactor_scheduler_arg = {"initial_lr": 0.}
+  int_args = ["k","n_sma_threshold","warmup"]
+  str_args = ["transformer","grad_transformer"]
+  if not args.optimizer_arg==None and len(args.optimizer_arg)>0:
+    for _opt_arg in args.optimizer_arg:
+      key, value = _opt_arg.split("=")
+      if value=="True" or value=="False":
+        optimizer_arg[key]=bool((value=="True"))
+      elif key=="betas" or key=="nus" or key=="eps2" or (key=="eps" and "," in value):
+        _value = value.split(",")
+        optimizer_arg[key] = (float(_value[0]),float(_value[1]))
+        del _value
+      elif key in int_args:
+        if "Lookahead" in args.optimizer:
+          lookahed_arg[key] = int(value)
+        else:
+          optimizer_arg[key] = int(value)
+      elif key in str_args:
+        optimizer_arg[key] = value
+      else:
+        if key=="alpha" and "Lookahead" in args.optimizer:
+          lookahed_arg[key] = int(value)
+        elif key=="initial_lr" and args.optimizer == "Adafactor":
+          adafactor_scheduler_arg[key] = float(value)
+        else:
+          optimizer_arg[key] = float(value)
+    del _opt_arg
+  AdafactorScheduler_Flag = False
+  list_of_init_lr = []
+  if args.optimizer=="Adafactor":
+    if not "relative_step" in optimizer_arg:
+      optimizer_arg["relative_step"] = True
+    if "warmup_init" in optimizer_arg:
+      if optimizer_arg["warmup_init"]==True and optimizer_arg["relative_step"]==False:
+        print("**************\nwarmup_initはrelative_stepがオンである必要があるためrelative_stepをオンにします\n**************")
+        optimizer_arg["relative_step"] = True
+    if optimizer_arg["relative_step"] == True:
+      AdafactorScheduler_Flag = True
+      list_of_init_lr = [0.,0.]
+      if args.text_encoder_lr is not None: list_of_init_lr[0] = float(args.text_encoder_lr)
+      if args.unet_lr is not None: list_of_init_lr[1] = float(args.unet_lr)
+      #if not "initial_lr" in adafactor_scheduler_arg:
+      #  adafactor_scheduler_arg = args.learning_rate
+      args.learning_rate = None
+      args.text_encoder_lr = None
+      args.unet_lr = None
+  print(f"optimizer arg: {optimizer_arg}")
+  print("=-----------------------------------=")
+  if not AdafactorScheduler_Flag: args.split_lora_networks = False
   if args.split_lora_networks:
-    lr_dic, block_args_dic = append_module.create_lr_blocks(args.blocks_lr_setting, args.block_optim_args)
     lora_names = append_module.create_split_names(args.split_lora_networks, args.split_lora_level)
-    append_module.replace_prepare_optimizer_params(network, network_module)
-    trainable_params, adafactor_scheduler_arg, used_names = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr, lora_names, lr_dic, block_args_dic)
+    append_module.replace_prepare_optimizer_params(network)
+    trainable_params, _list_of_init_lr = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr, list_of_init_lr, lora_names)
   else:
     trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr)
-    if split_flag:
-      _t_lr = 0.
-      _u_lr = 0.
-      if args.text_encoder_lr:
-        _t_lr = args.text_encoder_lr
-      if args.unet_lr:
-        _u_lr = args.unet_lr
-      adafactor_scheduler_arg = {"initial_lr": [_t_lr, _u_lr]}
-  
-  optimizer_name, optimizer_args, optimizer = train_util.get_optimizer(args, trainable_params)
-  if args.use_lookahead:
-    try:
-      import torch_optimizer
-      lookahed_arg = {"k": 5, "alpha": 0.5}
-      if args.lookahead_arg is not None:
-        for _arg in args.lookahead_arg:
-          k, v = _arg.split("=")
-          if k == "k":
-            lookahed_arg[k] = int(v)
-          else:
-            lookahed_arg[k] = float(v)
-      optimizer = torch_optimizer.Lookahead(optimizer, **lookahed_arg)
-    except:
-      print("\n============\ntorch_optimizerのimportに失敗しました Lookaheadを無効化して処理を続けます\n============\n")
+    _list_of_init_lr = []
+  print(f"trainable_params_len: {len(trainable_params)}")
+  if len(_list_of_init_lr)>0:
+    list_of_init_lr = _list_of_init_lr
+    print(f"split loras network is {len(list_of_init_lr)}")
+  if len(list_of_init_lr) > 0:
+    adafactor_scheduler_arg["initial_lr"] = list_of_init_lr
+
+  optimizer = optimizer_class(trainable_params, lr=args.learning_rate, **optimizer_arg)
+
+  if args.optimizer=="Lookahead_DiffGrad" or args.optimizer=="Lookahedad_Adam":
+    optimizer = optim.Lookahead(optimizer, **lookahed_arg)
+    print(f"lookahed_arg: {lookahed_arg}")
+
   # dataloaderを準備する
   # 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, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
 
   # 学習ステップ数を計算する
   if args.max_train_epochs is not None:
-    args.max_train_steps = args.max_train_epochs * math.ceil(len(train_dataloader) / accelerator.num_processes)
-    if is_main_process:
-      print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
+    args.max_train_steps = args.max_train_epochs * len(train_dataloader)
+    print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
 
   # lr schedulerを用意する
-  if args.lr_scheduler.startswith("adafactor") and split_flag:
-    lr_scheduler = append_module.get_scheduler_Adafactor(args.lr_scheduler, optimizer, adafactor_scheduler_arg)
+  # lr_scheduler = diffusers.optimization.get_scheduler(
+  if AdafactorScheduler_Flag:
+    print("===================================\nAdafactorはデフォルトでrelative_stepがオンになっているので lrは自動算出されるためLrScheculerの指定も無効になります\nもし任意のLrやLr_Schedulerを使いたい場合は --optimizer_arg relative_ste=False を指定してください\nまた任意のLrを使う場合は scale_parameter=False も併せて指定するのが推奨です\n===================================")
+    lr_scheduler = append_module.AdafactorSchedule_append(optimizer, **adafactor_scheduler_arg)
+    print(f"AdafactorSchedule initial lrs: {lr_scheduler.get_lr()}")
+    del list_of_init_lr
   else:
-    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)
-    
+    lr_scheduler = 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)
+  
   #追加機能の設定をコメントに追記して残す
-  if args.use_lookahead:
-    args.training_comment=f"{args.training_comment} use Lookahead: True Lookahead args: {lookahed_arg}"
-  if args.split_lora_networks:
-    args.training_comment=f"{args.training_comment} split_lora_networks: {args.split_lora_networks} split_level: {args.split_lora_level}"
+  args.training_comment=f"{args.training_comment} optimizer: {args.optimizer} / optimizer_arg: {args.optimizer_arg}"
+  if AdafactorScheduler_Flag:
+    args.training_comment=f"{args.training_comment} split_lora_networks: {args.split_lora_networks}"
   if args.min_resolution:
     args.training_comment=f"{args.training_comment} min_resolution: {args.min_resolution} area_step: {args.area_step}"
 
@@ -312,21 +504,17 @@ def train(args):
     args.save_every_n_epochs = math.floor(num_train_epochs / args.save_n_epoch_ratio) or 1
 
   # 学習する
-  # TODO: find a way to handle total batch size when there are multiple datasets
   total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
-  
-  if is_main_process:
-    print("running training / 学習開始")
-    print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
-    print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
-    print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
-    print(f"  num epochs / epoch数: {num_train_epochs}")
-    print(f"  batch size per device / バッチサイズ: {', '.join([str(d.batch_size) for d in train_dataset_group.datasets])}")
-    # print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
-    print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
-    print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
-
-  # TODO refactor metadata creation and move to util
+  print("running training / 学習開始")
+  print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset.num_train_images}")
+  print(f"  num reg images / 正則化画像の数: {train_dataset.num_reg_images}")
+  print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
+  print(f"  num epochs / epoch数: {num_train_epochs}")
+  print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
+  print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
+  print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
+  print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
+
   metadata = {
       "ss_session_id": session_id,            # random integer indicating which group of epochs the model came from
       "ss_training_started_at": training_started_at,          # unix timestamp
@@ -334,10 +522,12 @@ def train(args):
       "ss_learning_rate": args.learning_rate,
       "ss_text_encoder_lr": args.text_encoder_lr,
       "ss_unet_lr": args.unet_lr,
-      "ss_num_train_images": train_dataset_group.num_train_images,
-      "ss_num_reg_images": train_dataset_group.num_reg_images,
+      "ss_num_train_images": train_dataset.num_train_images,          # includes repeating
+      "ss_num_reg_images": train_dataset.num_reg_images,
       "ss_num_batches_per_epoch": len(train_dataloader),
       "ss_num_epochs": num_train_epochs,
+      "ss_batch_size_per_device": args.train_batch_size,
+      "ss_total_batch_size": total_batch_size,
       "ss_gradient_checkpointing": args.gradient_checkpointing,
       "ss_gradient_accumulation_steps": args.gradient_accumulation_steps,
       "ss_max_train_steps": args.max_train_steps,
@@ -349,156 +539,32 @@ def train(args):
       "ss_mixed_precision": args.mixed_precision,
       "ss_full_fp16": bool(args.full_fp16),
       "ss_v2": bool(args.v2),
+      "ss_resolution": args.resolution,
       "ss_clip_skip": args.clip_skip,
       "ss_max_token_length": args.max_token_length,
+      "ss_color_aug": bool(args.color_aug),
+      "ss_flip_aug": bool(args.flip_aug),
+      "ss_random_crop": bool(args.random_crop),
+      "ss_shuffle_caption": bool(args.shuffle_caption),
       "ss_cache_latents": bool(args.cache_latents),
+      "ss_enable_bucket": bool(train_dataset.enable_bucket),
+      "ss_min_bucket_reso": train_dataset.min_bucket_reso,
+      "ss_max_bucket_reso": train_dataset.max_bucket_reso,
       "ss_seed": args.seed,
-      "ss_lowram": args.lowram,
+      "ss_keep_tokens": args.keep_tokens,
       "ss_noise_offset": args.noise_offset,
+      "ss_dataset_dirs": json.dumps(train_dataset.dataset_dirs_info),
+      "ss_reg_dataset_dirs": json.dumps(train_dataset.reg_dataset_dirs_info),
+      "ss_tag_frequency": json.dumps(train_dataset.tag_frequency),
+      "ss_bucket_info": json.dumps(train_dataset.bucket_info),
       "ss_training_comment": args.training_comment,       # will not be updated after training
-      "ss_sd_scripts_commit_hash": train_util.get_git_revision_hash(),
-      "ss_optimizer": optimizer_name + (f"({optimizer_args})" if len(optimizer_args) > 0 else ""),
-      "ss_max_grad_norm": args.max_grad_norm,
-      "ss_caption_dropout_rate": args.caption_dropout_rate,
-      "ss_caption_dropout_every_n_epochs": args.caption_dropout_every_n_epochs,
-      "ss_caption_tag_dropout_rate": args.caption_tag_dropout_rate,
-      "ss_face_crop_aug_range": args.face_crop_aug_range,
-      "ss_prior_loss_weight": args.prior_loss_weight,
+      "ss_sd_scripts_commit_hash": train_util.get_git_revision_hash()
   }
 
-  if use_user_config:
-    # save metadata of multiple datasets
-    # NOTE: pack "ss_datasets" value as json one time
-    #   or should also pack nested collections as json?
-    datasets_metadata = []
-    tag_frequency = {}                    # merge tag frequency for metadata editor
-    dataset_dirs_info = {}                # merge subset dirs for metadata editor
-
-    for dataset in train_dataset_group.datasets:
-      is_dreambooth_dataset = isinstance(dataset, DreamBoothDataset)
-      dataset_metadata = {
-          "is_dreambooth": is_dreambooth_dataset,
-          "batch_size_per_device": dataset.batch_size,
-          "num_train_images": dataset.num_train_images,          # includes repeating
-          "num_reg_images": dataset.num_reg_images,
-          "resolution": (dataset.width, dataset.height),
-          "enable_bucket": bool(dataset.enable_bucket),
-          "min_bucket_reso": dataset.min_bucket_reso,
-          "max_bucket_reso": dataset.max_bucket_reso,
-          "tag_frequency": dataset.tag_frequency,
-          "bucket_info": dataset.bucket_info,
-      }
-
-      subsets_metadata = []
-      for subset in dataset.subsets:
-        subset_metadata = {
-            "img_count": subset.img_count,
-            "num_repeats": subset.num_repeats,
-            "color_aug": bool(subset.color_aug),
-            "flip_aug": bool(subset.flip_aug),
-            "random_crop": bool(subset.random_crop),
-            "shuffle_caption": bool(subset.shuffle_caption),
-            "keep_tokens": subset.keep_tokens,
-        }
-
-        image_dir_or_metadata_file = None
-        if subset.image_dir:
-          image_dir = os.path.basename(subset.image_dir)
-          subset_metadata["image_dir"] = image_dir
-          image_dir_or_metadata_file = image_dir
-
-        if is_dreambooth_dataset:
-          subset_metadata["class_tokens"] = subset.class_tokens
-          subset_metadata["is_reg"] = subset.is_reg
-          if subset.is_reg:
-            image_dir_or_metadata_file = None                    # not merging reg dataset
-        else:
-          metadata_file = os.path.basename(subset.metadata_file)
-          subset_metadata["metadata_file"] = metadata_file
-          image_dir_or_metadata_file = metadata_file           # may overwrite
-
-        subsets_metadata.append(subset_metadata)
-
-        # merge dataset dir: not reg subset only
-        # TODO update additional-network extension to show detailed dataset config from metadata
-        if image_dir_or_metadata_file is not None:
-          # datasets may have a certain dir multiple times
-          v = image_dir_or_metadata_file
-          i = 2
-          while v in dataset_dirs_info:
-            v = image_dir_or_metadata_file + f" ({i})"
-            i += 1
-          image_dir_or_metadata_file = v
-
-          dataset_dirs_info[image_dir_or_metadata_file] = {
-              "n_repeats": subset.num_repeats,
-              "img_count": subset.img_count
-          }
-
-      dataset_metadata["subsets"] = subsets_metadata
-      datasets_metadata.append(dataset_metadata)
-
-      # merge tag frequency:
-      for ds_dir_name, ds_freq_for_dir in dataset.tag_frequency.items():
-        # あるディレクトリが複数のdatasetで使用されている場合、一度だけ数える
-        # もともと繰り返し回数を指定しているので、キャプション内でのタグの出現回数と、それが学習で何度使われるかは一致しない
-        # なので、ここで複数datasetの回数を合算してもあまり意味はない
-        if ds_dir_name in tag_frequency:
-          continue
-        tag_frequency[ds_dir_name] = ds_freq_for_dir
-
-    metadata["ss_datasets"] = json.dumps(datasets_metadata)
-    metadata["ss_tag_frequency"] = json.dumps(tag_frequency)
-    metadata["ss_dataset_dirs"] = json.dumps(dataset_dirs_info)
-  else:
-    # conserving backward compatibility when using train_dataset_dir and reg_dataset_dir
-    assert len(
-        train_dataset_group.datasets) == 1, f"There should be a single dataset but {len(train_dataset_group.datasets)} found. This seems to be a bug. / データセットは1個だけ存在するはずですが、実際には{len(train_dataset_group.datasets)}個でした。プログラムのバグかもしれません。"
-
-    dataset = train_dataset_group.datasets[0]
-
-    dataset_dirs_info = {}
-    reg_dataset_dirs_info = {}
-    if use_dreambooth_method:
-      for subset in dataset.subsets:
-        info = reg_dataset_dirs_info if subset.is_reg else dataset_dirs_info
-        info[os.path.basename(subset.image_dir)] = {
-            "n_repeats": subset.num_repeats,
-            "img_count": subset.img_count
-        }
-    else:
-      for subset in dataset.subsets:
-        dataset_dirs_info[os.path.basename(subset.metadata_file)] = {
-            "n_repeats": subset.num_repeats,
-            "img_count": subset.img_count
-        }
-
-    metadata.update({
-        "ss_batch_size_per_device": args.train_batch_size,
-        "ss_total_batch_size": total_batch_size,
-        "ss_resolution": args.resolution,
-        "ss_color_aug": bool(args.color_aug),
-        "ss_flip_aug": bool(args.flip_aug),
-        "ss_random_crop": bool(args.random_crop),
-        "ss_shuffle_caption": bool(args.shuffle_caption),
-        "ss_enable_bucket": bool(dataset.enable_bucket),
-        "ss_bucket_no_upscale": bool(dataset.bucket_no_upscale),
-        "ss_min_bucket_reso": dataset.min_bucket_reso,
-        "ss_max_bucket_reso": dataset.max_bucket_reso,
-        "ss_keep_tokens": args.keep_tokens,
-        "ss_dataset_dirs": json.dumps(dataset_dirs_info),
-        "ss_reg_dataset_dirs": json.dumps(reg_dataset_dirs_info),
-        "ss_tag_frequency": json.dumps(dataset.tag_frequency),
-        "ss_bucket_info": json.dumps(dataset.bucket_info),
-    })
-
-  # add extra args
-  if args.network_args:
-    metadata["ss_network_args"] = json.dumps(net_kwargs)
+  # uncomment if another network is added
   # for key, value in net_kwargs.items():
   #   metadata["ss_arg_" + key] = value
 
-  # model name and hash
   if args.pretrained_model_name_or_path is not None:
     sd_model_name = args.pretrained_model_name_or_path
     if os.path.exists(sd_model_name):
@@ -517,13 +583,6 @@ def train(args):
 
   metadata = {k: str(v) for k, v in metadata.items()}
 
-  # make minimum metadata for filtering
-  minimum_keys = ["ss_network_module", "ss_network_dim", "ss_network_alpha", "ss_network_args"]
-  minimum_metadata = {}
-  for key in minimum_keys:
-    if key in metadata:
-      minimum_metadata[key] = metadata[key]
-
   progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
   global_step = 0
 
@@ -536,9 +595,8 @@ def train(args):
   loss_list = []
   loss_total = 0.0
   for epoch in range(num_train_epochs):
-    if is_main_process:
-      print(f"epoch {epoch+1}/{num_train_epochs}")
-    train_dataset_group.set_current_epoch(epoch + 1)
+    print(f"epoch {epoch+1}/{num_train_epochs}")
+    train_dataset.set_current_epoch(epoch + 1)
 
     metadata["ss_epoch"] = str(epoch+1)
 
@@ -575,7 +633,7 @@ def train(args):
         noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
 
         # Predict the noise residual
-        with accelerator.autocast():
+        with autocast():
           noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
 
         if args.v_parameterization:
@@ -593,13 +651,12 @@ def train(args):
         loss = loss.mean()                # 平均なのでbatch_sizeで割る必要なし
 
         accelerator.backward(loss)
-        if accelerator.sync_gradients and args.max_grad_norm != 0.0:
+        if accelerator.sync_gradients:
           params_to_clip = network.get_trainable_params()
-          accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+          accelerator.clip_grad_norm_(params_to_clip, 1.0)  # args.max_grad_norm)
 
         optimizer.step()
-        if accelerator.sync_gradients:
-          lr_scheduler.step()
+        lr_scheduler.step()
         optimizer.zero_grad(set_to_none=True)
 
       # Checks if the accelerator has performed an optimization step behind the scenes
@@ -607,8 +664,6 @@ def train(args):
         progress_bar.update(1)
         global_step += 1
 
-        train_util.sample_images(accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
-
       current_loss = loss.detach().item()
       if epoch == 0:
         loss_list.append(current_loss)
@@ -621,7 +676,7 @@ def train(args):
       progress_bar.set_postfix(**logs)
 
       if args.logging_dir is not None:
-        logs = generate_step_logs(args, current_loss, avr_loss, lr_scheduler, used_names)
+        logs = generate_step_logs(args, current_loss, avr_loss, lr_scheduler)
         accelerator.log(logs, step=global_step)
 
       if global_step >= args.max_train_steps:
@@ -639,9 +694,8 @@ def train(args):
       def save_func():
         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)
+        unwrap_model(network).save_weights(ckpt_file, save_dtype, None 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
@@ -650,18 +704,15 @@ def train(args):
           print(f"removing old checkpoint: {old_ckpt_file}")
           os.remove(old_ckpt_file)
 
-      if is_main_process:
-        saving = train_util.save_on_epoch_end(args, save_func, remove_old_func, epoch + 1, num_train_epochs)
-        if saving and args.save_state:
-          train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)
-
-    train_util.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
+      saving = train_util.save_on_epoch_end(args, save_func, remove_old_func, epoch + 1, num_train_epochs)
+      if saving and args.save_state:
+        train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)
 
     # end of epoch
 
   metadata["ss_epoch"] = str(num_train_epochs)
-  metadata["ss_training_finished_at"] = str(time.time())
 
+  is_main_process = accelerator.is_main_process
   if is_main_process:
     network = unwrap_model(network)
 
@@ -680,7 +731,7 @@ def train(args):
     ckpt_file = os.path.join(args.output_dir, ckpt_name)
 
     print(f"save trained model to {ckpt_file}")
-    network.save_weights(ckpt_file, save_dtype, minimum_metadata if args.no_metadata else metadata)
+    network.save_weights(ckpt_file, save_dtype, None if args.no_metadata else metadata)
     print("model saved.")
 
 
@@ -690,8 +741,6 @@ if __name__ == '__main__':
   train_util.add_sd_models_arguments(parser)
   train_util.add_dataset_arguments(parser, True, True, True)
   train_util.add_training_arguments(parser, True)
-  train_util.add_optimizer_arguments(parser)
-  config_util.add_config_arguments(parser)
 
   parser.add_argument("--no_metadata", action='store_true', help="do not save metadata in output model / メタデータを出力先モデルに保存しない")
   parser.add_argument("--save_model_as", type=str, default="safetensors", choices=[None, "ckpt", "pt", "safetensors"],
@@ -699,6 +748,10 @@ if __name__ == '__main__':
 
   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("--lr_scheduler_num_cycles", type=int, default=1,
+                      help="Number of restarts for cosine scheduler with restarts / cosine with restartsスケジューラでのリスタート回数")
+  parser.add_argument("--lr_scheduler_power", type=float, default=1,
+                      help="Polynomial power for polynomial scheduler / polynomialスケジューラでのpolynomial power")
 
   parser.add_argument("--network_weights", type=str, default=None,
                       help="pretrained weights for network / 学習するネットワークの初期重み")
@@ -718,30 +771,27 @@ if __name__ == '__main__':
   #Optimizer変更関連のオプション追加
   append_module.add_append_arguments(parser)
   args = append_module.get_config(parser)
-  if not args.not_output_config:
-    #argsを保存する
-    import yaml
-    import datetime
-    _t = datetime.datetime.today().strftime('%Y%m%d_%H%M')
-    if args.output_name==None:
-      config_name = f"train_network_config_{_t}.yaml"
-    else:
-      config_name = f"train_network_config_{os.path.basename(args.output_name)}_{_t}.yaml"
-    print(f"{config_name} に設定を書き出し中...")
-    with open(config_name, mode="w") as f:
-        yaml.dump(args.__dict__, f, indent=4)
 
   if args.resolution==args.min_resolution:
     args.min_resolution=None
 
   train(args)
-  print("done!")
 
+  #学習が終わったら現在のargsを保存する
+#  import yaml
+#  import datetime
+#  _t = datetime.datetime.today().strftime('%Y%m%d_%H%M')
+#  if args.output_name==None:
+#    config_name = f"train_network_config_{_t}.yaml"
+#  else:
+#    config_name = f"train_network_config_{os.path.basename(args.output_name)}_{_t}.yaml"
+#  print(f"{config_name} に設定を書き出し中...")
+#  with open(config_name, mode="w") as f:
+#      yaml.dump(args.__dict__, f, indent=4)
+#  print("done!")
 
 '''
 optimizer設定メモ
-torch_optimizer.AdaBelief
-adastand.Adastand
 (optimizer_argから設定できるように変更するためのメモ)
 
 AdamWのweight_decay初期値は1e-2
@@ -771,7 +821,6 @@ Adafactor
 transformerベースのT5学習において最強とかいう噂のoptimizer
 huggingfaceのサンプルパラ
 eps=1e-30,1e-3 clip_threshold=1.0 decay_rate=-0.8 relative_step=False scale_parameter=False warmup_init=False
-epsの二つ目の値1e-3が学習率に影響大きい
 
 AggMo
 

train_textual_inversion.py

@@ -11,11 +11,7 @@ import diffusers
 from diffusers import DDPMScheduler
 
 import library.train_util as train_util
-import library.config_util as config_util
-from library.config_util import (
-  ConfigSanitizer,
-  BlueprintGenerator,
-)
+from library.train_util import DreamBoothDataset, FineTuningDataset
 
 imagenet_templates_small = [
     "a photo of a {}",
@@ -83,6 +79,7 @@ def train(args):
   train_util.prepare_dataset_args(args, True)
 
   cache_latents = args.cache_latents
+  use_dreambooth_method = args.in_json is None
 
   if args.seed is not None:
     set_seed(args.seed)
@@ -142,35 +139,21 @@ def train(args):
   print(f"create embeddings for {args.num_vectors_per_token} tokens, for {args.token_string}")
 
   # データセットを準備する
-  blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False))
-  if args.dataset_config is not None:
-    print(f"Load dataset config from {args.dataset_config}")
-    user_config = config_util.load_user_config(args.dataset_config)
-    ignored = ["train_data_dir", "reg_data_dir", "in_json"]
-    if any(getattr(args, attr) is not None for attr in ignored):
-      print("ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(', '.join(ignored)))
+  if use_dreambooth_method:
+    print("Use DreamBooth method.")
+    train_dataset = DreamBoothDataset(args.train_batch_size, args.train_data_dir, args.reg_data_dir,
+                                      tokenizer, args.max_token_length, args.caption_extension, args.shuffle_caption, args.keep_tokens,
+                                      args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                      args.bucket_reso_steps, args.bucket_no_upscale,
+                                      args.prior_loss_weight, args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop, args.debug_dataset)
   else:
-    use_dreambooth_method = args.in_json is None
-    if use_dreambooth_method:
-      print("Use DreamBooth method.")
-      user_config = {
-        "datasets": [{
-          "subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)
-        }]
-      }
-    else:
-      print("Train with captions.")
-      user_config = {
-        "datasets": [{
-          "subsets": [{
-            "image_dir": args.train_data_dir,
-            "metadata_file": args.in_json,
-          }]
-        }]
-      }
-
-  blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
-  train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
+    print("Train with captions.")
+    train_dataset = FineTuningDataset(args.in_json, args.train_batch_size, args.train_data_dir,
+                                      tokenizer, args.max_token_length, args.shuffle_caption, args.keep_tokens,
+                                      args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                      args.bucket_reso_steps, args.bucket_no_upscale,
+                                      args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop,
+                                      args.dataset_repeats, args.debug_dataset)
 
   # make captions: tokenstring tokenstring1 tokenstring2 ...tokenstringn という文字列に書き換える超乱暴な実装
   if use_template:
@@ -180,30 +163,20 @@ def train(args):
     captions = []
     for tmpl in templates:
       captions.append(tmpl.format(replace_to))
-    train_dataset_group.add_replacement("", captions)
+    train_dataset.add_replacement("", captions)
+  elif args.num_vectors_per_token > 1:
+    replace_to = " ".join(token_strings)
+    train_dataset.add_replacement(args.token_string, replace_to)
 
-    if args.num_vectors_per_token > 1:
-      prompt_replacement = (args.token_string, replace_to)
-    else:
-      prompt_replacement = None
-  else:
-    if args.num_vectors_per_token > 1:
-      replace_to = " ".join(token_strings)
-      train_dataset_group.add_replacement(args.token_string, replace_to)
-      prompt_replacement = (args.token_string, replace_to)
-    else:
-      prompt_replacement = None
+  train_dataset.make_buckets()
 
   if args.debug_dataset:
-    train_util.debug_dataset(train_dataset_group, show_input_ids=True)
+    train_util.debug_dataset(train_dataset, show_input_ids=True)
     return
-  if len(train_dataset_group) == 0:
+  if len(train_dataset) == 0:
     print("No data found. Please verify arguments / 画像がありません。引数指定を確認してください")
     return
 
-  if cache_latents:
-    assert train_dataset_group.is_latent_cacheable(), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
-
   # モデルに xformers とか memory efficient attention を組み込む
   train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
 
@@ -213,7 +186,7 @@ def train(args):
     vae.requires_grad_(False)
     vae.eval()
     with torch.no_grad():
-      train_dataset_group.cache_latents(vae)
+      train_dataset.cache_latents(vae)
     vae.to("cpu")
     if torch.cuda.is_available():
       torch.cuda.empty_cache()
@@ -225,14 +198,35 @@ def train(args):
 
   # 学習に必要なクラスを準備する
   print("prepare optimizer, data loader etc.")
+
+  # 8-bit Adamを使う
+  if args.use_8bit_adam:
+    try:
+      import bitsandbytes as bnb
+    except ImportError:
+      raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
+    print("use 8-bit Adam optimizer")
+    optimizer_class = bnb.optim.AdamW8bit
+  elif args.use_lion_optimizer:
+    try:
+      import lion_pytorch
+    except ImportError:
+      raise ImportError("No lion_pytorch / lion_pytorch がインストールされていないようです")
+    print("use Lion optimizer")
+    optimizer_class = lion_pytorch.Lion
+  else:
+    optimizer_class = torch.optim.AdamW
+
   trainable_params = text_encoder.get_input_embeddings().parameters()
-  _, _, optimizer = train_util.get_optimizer(args, trainable_params)
+
+  # betaやweight decayはdiffusers DreamBoothもDreamBooth SDもデフォルト値のようなのでオプションはとりあえず省略
+  optimizer = optimizer_class(trainable_params, lr=args.learning_rate)
 
   # dataloaderを準備する
   # 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, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
 
   # 学習ステップ数を計算する
   if args.max_train_epochs is not None:
@@ -240,9 +234,8 @@ 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,
-                                              num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
-                                              num_cycles=args.lr_scheduler_num_cycles, power=args.lr_scheduler_power)
+  lr_scheduler = diffusers.optimization.get_scheduler(
+      args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps, num_training_steps=args.max_train_steps * args.gradient_accumulation_steps)
 
   # acceleratorがなんかよろしくやってくれるらしい
   text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
@@ -290,8 +283,8 @@ def train(args):
   # 学習する
   total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
   print("running training / 学習開始")
-  print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
-  print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
+  print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset.num_train_images}")
+  print(f"  num reg images / 正則化画像の数: {train_dataset.num_reg_images}")
   print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
   print(f"  num epochs / epoch数: {num_train_epochs}")
   print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
@@ -310,11 +303,12 @@ def train(args):
 
   for epoch in range(num_train_epochs):
     print(f"epoch {epoch+1}/{num_train_epochs}")
-    train_dataset_group.set_current_epoch(epoch + 1)
+    train_dataset.set_current_epoch(epoch + 1)
 
     text_encoder.train()
 
     loss_total = 0
+    bef_epo_embs = unwrap_model(text_encoder).get_input_embeddings().weight[token_ids].data.detach().clone()
     for step, batch in enumerate(train_dataloader):
       with accelerator.accumulate(text_encoder):
         with torch.no_grad():
@@ -363,9 +357,9 @@ def train(args):
         loss = loss.mean()                # 平均なのでbatch_sizeで割る必要なし
 
         accelerator.backward(loss)
-        if accelerator.sync_gradients and args.max_grad_norm != 0.0:
+        if accelerator.sync_gradients:
           params_to_clip = text_encoder.get_input_embeddings().parameters()
-          accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+          accelerator.clip_grad_norm_(params_to_clip, 1.0)  # args.max_grad_norm)
 
         optimizer.step()
         lr_scheduler.step()
@@ -380,14 +374,9 @@ def train(args):
         progress_bar.update(1)
         global_step += 1
 
-        train_util.sample_images(accelerator, args, None, global_step, accelerator.device,
-                                 vae, tokenizer, text_encoder, unet, prompt_replacement)
-
       current_loss = loss.detach().item()
       if args.logging_dir is not None:
-        logs = {"loss": current_loss, "lr": float(lr_scheduler.get_last_lr()[0])}
-        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 = {"loss": current_loss, "lr": lr_scheduler.get_last_lr()[0]}
         accelerator.log(logs, step=global_step)
 
       loss_total += current_loss
@@ -405,6 +394,8 @@ def train(args):
     accelerator.wait_for_everyone()
 
     updated_embs = unwrap_model(text_encoder).get_input_embeddings().weight[token_ids].data.detach().clone()
+    # d = updated_embs - bef_epo_embs
+    # print(bef_epo_embs.size(), updated_embs.size(), d.mean(), d.min())
 
     if args.save_every_n_epochs is not None:
       model_name = train_util.DEFAULT_EPOCH_NAME if args.output_name is None else args.output_name
@@ -426,9 +417,6 @@ def train(args):
       if saving and args.save_state:
         train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)
 
-    train_util.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device,
-                             vae, tokenizer, text_encoder, unet, prompt_replacement)
-
     # end of epoch
 
   is_main_process = accelerator.is_main_process
@@ -503,8 +491,6 @@ if __name__ == '__main__':
   train_util.add_sd_models_arguments(parser)
   train_util.add_dataset_arguments(parser, True, True, False)
   train_util.add_training_arguments(parser, True)
-  train_util.add_optimizer_arguments(parser)
-  config_util.add_config_arguments(parser)
 
   parser.add_argument("--save_model_as", type=str, default="pt", choices=[None, "ckpt", "pt", "safetensors"],
                       help="format to save the model (default is .pt) / モデル保存時の形式（デフォルトはpt）")