Upload 35 files

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)
-'''

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)
 
@@ -1853,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?
@@ -1871,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
@@ -1964,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 = {}
@@ -2075,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):
@@ -2096,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()
@@ -2119,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)
@@ -2144,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:
@@ -2285,12 +2186,9 @@ def main(args):
 
   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
@@ -2321,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 = []
@@ -2393,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)
@@ -2410,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)
@@ -2439,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
@@ -2531,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]
@@ -2595,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)
@@ -2641,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.")
@@ -2654,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()
 
@@ -2700,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',
@@ -2713,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,
@@ -2729,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層目の出力を使う')
@@ -2751,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/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,93 +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
@@ -310,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
@@ -342,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():
@@ -425,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):
     '''
@@ -526,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]:
@@ -536,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)
@@ -566,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)
@@ -582,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])
@@ -597,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
@@ -605,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:
@@ -624,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)
 
@@ -653,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
@@ -665,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]
@@ -699,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になる
@@ -709,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))
@@ -740,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（解像度）指定は必須です"
 
@@ -764,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
@@ -787,171 +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, [], []
+
+      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, [], []
 
-      img_paths = glob_images(subset.image_dir, "*")
-      print(f"found directory {subset.image_dir} contains {len(img_paths)} image files")
+      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)
-          
-      self.set_tag_frequency(os.path.basename(subset.image_dir), captions)         # タグ頻度を記録
+        cap_for_img = read_caption(img_path)
+        captions.append(caption_by_folder if cap_for_img is None else cap_for_img)
 
-      return img_paths, captions
-
-    print("prepare images.")
-    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
+      self.set_tag_frequency(os.path.basename(dir), captions)         # タグ頻度を記録
 
-      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
+      return n_repeats, img_paths, captions
 
-      if subset.is_reg:
-        num_reg_images += subset.num_repeats * len(img_paths)
-      else:
-        num_train_images += subset.num_repeats * len(img_paths)
+    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, 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:
-          # わりといい加減だがいい方法が思いつかん
-          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)
-        assert caption is not None and len(caption) > 0, f"caption or tag is required / キャプションまたはタグは必須です:{abs_path}"
+    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:
@@ -963,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ファイルは使用されません")
@@ -1009,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'
 
@@ -1021,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
@@ -1033,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'])):
@@ -1491,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):
@@ -1543,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",
@@ -1550,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を上書きします）")
@@ -1570,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,
@@ -1608,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,
@@ -1634,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:
@@ -1663,249 +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}")
-
-    min_lr = lr
-    if type(trainable_params) == list and type(trainable_params[0]) == dict:
-      for group in trainable_params:
-        min_lr = min(min_lr, group.get("lr", lr))
-
-    if min_lr <= 0.1:
-      print(
-          f'learning rate is too low. If using dadaptation, set learning rate around 1.0 / 学習率が低すぎるようです。1.0前後の値を指定してください: {min_lr}')
-      print('recommend option: lr=1.0 / 推奨は1.0です')
-
-    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(',')])
@@ -1928,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
 
 
@@ -2000,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
@@ -2181,185 +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:
-      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
-
-  # ここでCUDAのキャッシュクリアとかしたほうがいいのか……
-
-  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):
-        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])
-            
-        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))
-
-  torch.set_rng_state(rng_state)
-  torch.cuda.set_rng_state(cuda_rng_state)
-  vae.to(org_vae_device)
-
 # endregion
 
 # region 前処理用

networks/lora.py

@@ -126,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

tools/convert_diffusers20_original_sd.py

@@ -0,0 +1,89 @@
+# convert Diffusers v1.x/v2.0 model to original Stable Diffusion
+
+import argparse
+import os
+import torch
+from diffusers import StableDiffusionPipeline
+
+import library.model_util as model_util
+
+
+def convert(args):
+  # 引数を確認する
+  load_dtype = torch.float16 if args.fp16 else None
+
+  save_dtype = None
+  if args.fp16:
+    save_dtype = torch.float16
+  elif args.bf16:
+    save_dtype = torch.bfloat16
+  elif args.float:
+    save_dtype = torch.float
+
+  is_load_ckpt = os.path.isfile(args.model_to_load)
+  is_save_ckpt = len(os.path.splitext(args.model_to_save)[1]) > 0
+
+  assert not is_load_ckpt or args.v1 != args.v2, f"v1 or v2 is required to load checkpoint / checkpointの読み込みにはv1/v2指定が必要です"
+  assert is_save_ckpt or args.reference_model is not None, f"reference model is required to save as Diffusers / Diffusers形式での保存には参照モデルが必要です"
+
+  # モデルを読み込む
+  msg = "checkpoint" if is_load_ckpt else ("Diffusers" + (" as fp16" if args.fp16 else ""))
+  print(f"loading {msg}: {args.model_to_load}")
+
+  if is_load_ckpt:
+    v2_model = args.v2
+    text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(v2_model, args.model_to_load)
+  else:
+    pipe = StableDiffusionPipeline.from_pretrained(args.model_to_load, torch_dtype=load_dtype, tokenizer=None, safety_checker=None)
+    text_encoder = pipe.text_encoder
+    vae = pipe.vae
+    unet = pipe.unet
+
+    if args.v1 == args.v2:
+      # 自動判定する
+      v2_model = unet.config.cross_attention_dim == 1024
+      print("checking model version: model is " + ('v2' if v2_model else 'v1'))
+    else:
+      v2_model = not args.v1
+
+  # 変換して保存する
+  msg = ("checkpoint" + ("" if save_dtype is None else f" in {save_dtype}")) if is_save_ckpt else "Diffusers"
+  print(f"converting and saving as {msg}: {args.model_to_save}")
+
+  if is_save_ckpt:
+    original_model = args.model_to_load if is_load_ckpt else None
+    key_count = model_util.save_stable_diffusion_checkpoint(v2_model, args.model_to_save, text_encoder, unet,
+                                                            original_model, args.epoch, args.global_step, save_dtype, vae)
+    print(f"model saved. total converted state_dict keys: {key_count}")
+  else:
+    print(f"copy scheduler/tokenizer config from: {args.reference_model}")
+    model_util.save_diffusers_checkpoint(v2_model, args.model_to_save, text_encoder, unet, args.reference_model, vae, args.use_safetensors)
+    print(f"model saved.")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--v1", action='store_true',
+                      help='load v1.x model (v1 or v2 is required to load checkpoint) / 1.xのモデルを読み込む')
+  parser.add_argument("--v2", action='store_true',
+                      help='load v2.0 model (v1 or v2 is required to load checkpoint) / 2.0のモデルを読み込む')
+  parser.add_argument("--fp16", action='store_true',
+                      help='load as fp16 (Diffusers only) and save as fp16 (checkpoint only) / fp16形式で読み込み（Diffusers形式のみ対応）、保存する（checkpointのみ対応）')
+  parser.add_argument("--bf16", action='store_true', help='save as bf16 (checkpoint only) / bf16形式で保存する（checkpointのみ対応）')
+  parser.add_argument("--float", action='store_true',
+                      help='save as float (checkpoint only) / float(float32)形式で保存する（checkpointのみ対応）')
+  parser.add_argument("--epoch", type=int, default=0, help='epoch to write to checkpoint / checkpointに記録するepoch数の値')
+  parser.add_argument("--global_step", type=int, default=0,
+                      help='global_step to write to checkpoint / checkpointに記録するglobal_stepの値')
+  parser.add_argument("--reference_model", type=str, default=None,
+                      help="reference model for schduler/tokenizer, required in saving Diffusers, copy schduler/tokenizer from this / scheduler/tokenizerのコピー元のDiffusersモデル、Diffusers形式で保存するときに必要")
+  parser.add_argument("--use_safetensors", action='store_true',
+                      help="use safetensors format to save Diffusers model (checkpoint depends on the file extension) / Duffusersモデルをsafetensors形式で保存する（checkpointは拡張子で自動判定）")
+
+  parser.add_argument("model_to_load", type=str, default=None,
+                      help="model to load: checkpoint file or Diffusers model's directory / 読み込むモデル、checkpointかDiffusers形式モデルのディレクトリ")
+  parser.add_argument("model_to_save", type=str, default=None,
+                      help="model to save: checkpoint (with extension) or Diffusers model's directory (without extension) / 変換後のモデル、拡張子がある場合はcheckpoint、ない場合はDiffusesモデルとして保存")
+
+  args = parser.parse_args()
+  convert(args)

tools/detect_face_rotate.py

@@ -0,0 +1,239 @@
+# このスクリプトのライセンスは、train_dreambooth.pyと同じくApache License 2.0とします
+# (c) 2022 Kohya S. @kohya_ss
+
+# 横長の画像から顔検出して正立するように回転し、そこを中心に正方形に切り出す
+
+# v2: extract max face if multiple faces are found
+# v3: add crop_ratio option
+# v4: add multiple faces extraction and min/max size
+
+import argparse
+import math
+import cv2
+import glob
+import os
+from anime_face_detector import create_detector
+from tqdm import tqdm
+import numpy as np
+
+KP_REYE = 11
+KP_LEYE = 19
+
+SCORE_THRES = 0.90
+
+
+def detect_faces(detector, image, min_size):
+  preds = detector(image)                     # bgr
+  # print(len(preds))
+
+  faces = []
+  for pred in preds:
+    bb = pred['bbox']
+    score = bb[-1]
+    if score < SCORE_THRES:
+      continue
+
+    left, top, right, bottom = bb[:4]
+    cx = int((left + right) / 2)
+    cy = int((top + bottom) / 2)
+    fw = int(right - left)
+    fh = int(bottom - top)
+
+    lex, ley = pred['keypoints'][KP_LEYE, 0:2]
+    rex, rey = pred['keypoints'][KP_REYE, 0:2]
+    angle = math.atan2(ley - rey, lex - rex)
+    angle = angle / math.pi * 180
+
+    faces.append((cx, cy, fw, fh, angle))
+
+  faces.sort(key=lambda x: max(x[2], x[3]), reverse=True)         # 大きい順
+  return faces
+
+
+def rotate_image(image, angle, cx, cy):
+  h, w = image.shape[0:2]
+  rot_mat = cv2.getRotationMatrix2D((cx, cy), angle, 1.0)
+
+  # # 回転する分、すこし画像サイズを大きくする→とりあえず無効化
+  # nh = max(h, int(w * math.sin(angle)))
+  # nw = max(w, int(h * math.sin(angle)))
+  # if nh > h or nw > w:
+  #   pad_y = nh - h
+  #   pad_t = pad_y // 2
+  #   pad_x = nw - w
+  #   pad_l = pad_x // 2
+  #   m = np.array([[0, 0, pad_l],
+  #                 [0, 0, pad_t]])
+  #   rot_mat = rot_mat + m
+  #   h, w = nh, nw
+  #   cx += pad_l
+  #   cy += pad_t
+
+  result = cv2.warpAffine(image, rot_mat, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)
+  return result, cx, cy
+
+
+def process(args):
+  assert (not args.resize_fit) or args.resize_face_size is None, f"resize_fit and resize_face_size can't be specified both / resize_fitとresize_face_sizeはどちらか片方しか指定できません"
+  assert args.crop_ratio is None or args.resize_face_size is None, f"crop_ratio指定時はresize_face_sizeは指定できません"
+
+  # アニメ顔検出モデルを読み込む
+  print("loading face detector.")
+  detector = create_detector('yolov3')
+
+  # cropの引数を解析する
+  if args.crop_size is None:
+    crop_width = crop_height = None
+  else:
+    tokens = args.crop_size.split(',')
+    assert len(tokens) == 2, f"crop_size must be 'width,height' / crop_sizeは'幅,高さ'で指定してください"
+    crop_width, crop_height = [int(t) for t in tokens]
+
+  if args.crop_ratio is None:
+    crop_h_ratio = crop_v_ratio = None
+  else:
+    tokens = args.crop_ratio.split(',')
+    assert len(tokens) == 2, f"crop_ratio must be 'horizontal,vertical' / crop_ratioは'幅,高さ'の倍率で指定してください"
+    crop_h_ratio, crop_v_ratio = [float(t) for t in tokens]
+
+  # 画像を処理する
+  print("processing.")
+  output_extension = ".png"
+
+  os.makedirs(args.dst_dir, exist_ok=True)
+  paths = glob.glob(os.path.join(args.src_dir, "*.png")) + glob.glob(os.path.join(args.src_dir, "*.jpg")) + \
+      glob.glob(os.path.join(args.src_dir, "*.webp"))
+  for path in tqdm(paths):
+    basename = os.path.splitext(os.path.basename(path))[0]
+
+    # image = cv2.imread(path)        # 日本語ファイル名でエラーになる
+    image = cv2.imdecode(np.fromfile(path, np.uint8), cv2.IMREAD_UNCHANGED)
+    if len(image.shape) == 2:
+      image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+    if image.shape[2] == 4:
+      print(f"image has alpha. ignore / 画像の透明度が設定されているため無視します: {path}")
+      image = image[:, :, :3].copy()                    # copyをしないと内部的に透明度情報が付いたままになるらしい
+
+    h, w = image.shape[:2]
+
+    faces = detect_faces(detector, image, args.multiple_faces)
+    for i, face in enumerate(faces):
+      cx, cy, fw, fh, angle = face
+      face_size = max(fw, fh)
+      if args.min_size is not None and face_size < args.min_size:
+        continue
+      if args.max_size is not None and face_size >= args.max_size:
+        continue
+      face_suffix = f"_{i+1:02d}" if args.multiple_faces else ""
+
+      # オプション指定があれば回転する
+      face_img = image
+      if args.rotate:
+        face_img, cx, cy = rotate_image(face_img, angle, cx, cy)
+
+      # オプション指定があれば顔を中心に切り出す
+      if crop_width is not None or crop_h_ratio is not None:
+        cur_crop_width, cur_crop_height = crop_width, crop_height
+        if crop_h_ratio is not None:
+          cur_crop_width = int(face_size * crop_h_ratio + .5)
+          cur_crop_height = int(face_size * crop_v_ratio + .5)
+
+        # リサイズを必要なら行う
+        scale = 1.0
+        if args.resize_face_size is not None:
+          # 顔サイズを基準にリサイズする
+          scale = args.resize_face_size / face_size
+          if scale < cur_crop_width / w:
+            print(
+                f"image width too small in face size based resizing / 顔を基準にリサイズすると画像の幅がcrop sizeより小さい（顔が相対的に大きすぎる）ので顔サイズが変わります: {path}")
+            scale = cur_crop_width / w
+          if scale < cur_crop_height / h:
+            print(
+                f"image height too small in face size based resizing / 顔を基準にリサイズすると画像の高さがcrop sizeより小さい（顔が相対的に大きすぎる）ので顔サイズが変わります: {path}")
+            scale = cur_crop_height / h
+        elif crop_h_ratio is not None:
+          # 倍率指定の時にはリサイズしない
+          pass
+        else:
+          # 切り出しサイズ指定あり
+          if w < cur_crop_width:
+            print(f"image width too small/ 画像の幅がcrop sizeより小さいので画質が劣化します: {path}")
+            scale = cur_crop_width / w
+          if h < cur_crop_height:
+            print(f"image height too small/ 画像の高さがcrop sizeより小さいので画質が劣化します: {path}")
+            scale = cur_crop_height / h
+          if args.resize_fit:
+            scale = max(cur_crop_width / w, cur_crop_height / h)
+
+        if scale != 1.0:
+          w = int(w * scale + .5)
+          h = int(h * scale + .5)
+          face_img = cv2.resize(face_img, (w, h), interpolation=cv2.INTER_AREA if scale < 1.0 else cv2.INTER_LANCZOS4)
+          cx = int(cx * scale + .5)
+          cy = int(cy * scale + .5)
+          fw = int(fw * scale + .5)
+          fh = int(fh * scale + .5)
+
+        cur_crop_width = min(cur_crop_width, face_img.shape[1])
+        cur_crop_height = min(cur_crop_height, face_img.shape[0])
+
+        x = cx - cur_crop_width // 2
+        cx = cur_crop_width // 2
+        if x < 0:
+          cx = cx + x
+          x = 0
+        elif x + cur_crop_width > w:
+          cx = cx + (x + cur_crop_width - w)
+          x = w - cur_crop_width
+        face_img = face_img[:, x:x+cur_crop_width]
+
+        y = cy - cur_crop_height // 2
+        cy = cur_crop_height // 2
+        if y < 0:
+          cy = cy + y
+          y = 0
+        elif y + cur_crop_height > h:
+          cy = cy + (y + cur_crop_height - h)
+          y = h - cur_crop_height
+        face_img = face_img[y:y + cur_crop_height]
+
+      # # debug
+      # print(path, cx, cy, angle)
+      # crp = cv2.resize(image, (image.shape[1]//8, image.shape[0]//8))
+      # cv2.imshow("image", crp)
+      # if cv2.waitKey() == 27:
+      #   break
+      # cv2.destroyAllWindows()
+
+      # debug
+      if args.debug:
+        cv2.rectangle(face_img, (cx-fw//2, cy-fh//2), (cx+fw//2, cy+fh//2), (255, 0, 255), fw//20)
+
+      _, buf = cv2.imencode(output_extension, face_img)
+      with open(os.path.join(args.dst_dir, f"{basename}{face_suffix}_{cx:04d}_{cy:04d}_{fw:04d}_{fh:04d}{output_extension}"), "wb") as f:
+        buf.tofile(f)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--src_dir", type=str, help="directory to load images / 画像を読み込むディレクトリ")
+  parser.add_argument("--dst_dir", type=str, help="directory to save images / 画像を保存するディレクトリ")
+  parser.add_argument("--rotate", action="store_true", help="rotate images to align faces / 顔が正立するように画像を回転する")
+  parser.add_argument("--resize_fit", action="store_true",
+                      help="resize to fit smaller side after cropping / 切り出し後の画像の短辺がcrop_sizeにあうようにリサイズする")
+  parser.add_argument("--resize_face_size", type=int, default=None,
+                      help="resize image before cropping by face size / 切り出し前に顔がこのサイズになるようにリサイズする")
+  parser.add_argument("--crop_size", type=str, default=None,
+                      help="crop images with 'width,height' pixels, face centered / 顔を中心として'幅,高さ'のサイズで切り出す")
+  parser.add_argument("--crop_ratio", type=str, default=None,
+                      help="crop images with 'horizontal,vertical' ratio to face, face centered / 顔を中心として顔サイズの'幅倍率,高さ倍率'のサイズで切り出す")
+  parser.add_argument("--min_size", type=int, default=None,
+                      help="minimum face size to output (included) / 処理対象とする顔の最小サイズ（この値以上）")
+  parser.add_argument("--max_size", type=int, default=None,
+                      help="maximum face size to output (excluded) / 処理対象とする顔の最大サイズ（この値未満）")
+  parser.add_argument("--multiple_faces", action="store_true",
+                      help="output each faces / 複数の顔が見つかった場合、それぞれを切り出す")
+  parser.add_argument("--debug", action="store_true", help="render rect for face / 処理後画像の顔位置に矩形を描画します")
+  args = parser.parse_args()
+
+  process(args)

tools/resize_images_to_resolution.py

@@ -0,0 +1,122 @@
+import glob
+import os
+import cv2
+import argparse
+import shutil
+import math
+from PIL import Image
+import numpy as np
+
+
+def resize_images(src_img_folder, dst_img_folder, max_resolution="512x512", divisible_by=2, interpolation=None, save_as_png=False, copy_associated_files=False):
+  # Split the max_resolution string by "," and strip any whitespaces
+  max_resolutions = [res.strip() for res in max_resolution.split(',')]
+
+  # # Calculate max_pixels from max_resolution string
+  # max_pixels = int(max_resolution.split("x")[0]) * int(max_resolution.split("x")[1])
+
+  # Create destination folder if it does not exist
+  if not os.path.exists(dst_img_folder):
+    os.makedirs(dst_img_folder)
+
+  # Select interpolation method
+  if interpolation == 'lanczos4':
+    cv2_interpolation = cv2.INTER_LANCZOS4
+  elif interpolation == 'cubic':
+    cv2_interpolation = cv2.INTER_CUBIC
+  else:
+    cv2_interpolation = cv2.INTER_AREA
+
+  # Iterate through all files in src_img_folder
+  img_exts = (".png", ".jpg", ".jpeg", ".webp", ".bmp")                   # copy from train_util.py
+  for filename in os.listdir(src_img_folder):
+    # Check if the image is png, jpg or webp etc...
+    if not filename.endswith(img_exts):
+      # Copy the file to the destination folder if not png, jpg or webp etc (.txt or .caption or etc.)
+      shutil.copy(os.path.join(src_img_folder, filename), os.path.join(dst_img_folder, filename))
+      continue
+
+    # Load image
+    # img = cv2.imread(os.path.join(src_img_folder, filename))
+    image = Image.open(os.path.join(src_img_folder, filename))
+    if not image.mode == "RGB":
+      image = image.convert("RGB")
+    img = np.array(image, np.uint8)
+
+    base, _ = os.path.splitext(filename)
+    for max_resolution in max_resolutions:
+      # Calculate max_pixels from max_resolution string
+      max_pixels = int(max_resolution.split("x")[0]) * int(max_resolution.split("x")[1])
+
+      # Calculate current number of pixels
+      current_pixels = img.shape[0] * img.shape[1]
+
+      # Check if the image needs resizing
+      if current_pixels > max_pixels:
+        # Calculate scaling factor
+        scale_factor = max_pixels / current_pixels
+
+        # Calculate new dimensions
+        new_height = int(img.shape[0] * math.sqrt(scale_factor))
+        new_width = int(img.shape[1] * math.sqrt(scale_factor))
+
+        # Resize image
+        img = cv2.resize(img, (new_width, new_height), interpolation=cv2_interpolation)
+      else:
+        new_height, new_width = img.shape[0:2]
+
+      # Calculate the new height and width that are divisible by divisible_by (with/without resizing)
+      new_height = new_height if new_height % divisible_by == 0 else new_height - new_height % divisible_by
+      new_width = new_width if new_width % divisible_by == 0 else new_width - new_width % divisible_by
+
+      # Center crop the image to the calculated dimensions
+      y = int((img.shape[0] - new_height) / 2)
+      x = int((img.shape[1] - new_width) / 2)
+      img = img[y:y + new_height, x:x + new_width]
+
+      # Split filename into base and extension
+      new_filename = base + '+' + max_resolution + ('.png' if save_as_png else '.jpg')
+
+      # Save resized image in dst_img_folder
+      # cv2.imwrite(os.path.join(dst_img_folder, new_filename), img, [cv2.IMWRITE_JPEG_QUALITY, 100])
+      image = Image.fromarray(img)
+      image.save(os.path.join(dst_img_folder, new_filename), quality=100)
+
+      proc = "Resized" if current_pixels > max_pixels else "Saved"
+      print(f"{proc} image: {filename} with size {img.shape[0]}x{img.shape[1]} as {new_filename}")
+
+    # If other files with same basename, copy them with resolution suffix
+    if copy_associated_files:
+      asoc_files = glob.glob(os.path.join(src_img_folder, base + ".*"))
+      for asoc_file in asoc_files:
+        ext = os.path.splitext(asoc_file)[1]
+        if ext in img_exts:
+          continue
+        for max_resolution in max_resolutions:
+          new_asoc_file = base + '+' + max_resolution + ext
+          print(f"Copy {asoc_file} as {new_asoc_file}")
+          shutil.copy(os.path.join(src_img_folder, asoc_file), os.path.join(dst_img_folder, new_asoc_file))
+
+
+def main():
+  parser = argparse.ArgumentParser(
+      description='Resize images in a folder to a specified max resolution(s) / 指定されたフォルダ内の画像を指定した最大画像サイズ（面積）以下にアスペクト比を維持したままリサイズします')
+  parser.add_argument('src_img_folder', type=str, help='Source folder containing the images / 元画像のフォルダ')
+  parser.add_argument('dst_img_folder', type=str, help='Destination folder to save the resized images / リサイズ後の画像を保存するフォルダ')
+  parser.add_argument('--max_resolution', type=str,
+                      help='Maximum resolution(s) in the format "512x512,384x384, etc, etc" / 最大画像サイズをカンマ区切りで指定 ("512x512,384x384, etc, etc" など)', default="512x512,384x384,256x256,128x128")
+  parser.add_argument('--divisible_by', type=int,
+                      help='Ensure new dimensions are divisible by this value / リサイズ後の画像のサイズをこの値で割り切れるようにします', default=1)
+  parser.add_argument('--interpolation', type=str, choices=['area', 'cubic', 'lanczos4'],
+                      default='area', help='Interpolation method for resizing / リサイズ時の補完方法')
+  parser.add_argument('--save_as_png', action='store_true', help='Save as png format / png形式で保存')
+  parser.add_argument('--copy_associated_files', action='store_true',
+                      help='Copy files with same base name to images (captions etc) / 画像と同じファイル名（拡張子を除く）のファイルもコピーする')
+
+  args = parser.parse_args()
+  resize_images(args.src_img_folder, args.dst_img_folder, args.max_resolution,
+                args.divisible_by, args.interpolation, args.save_as_png, args.copy_associated_files)
+
+
+if __name__ == '__main__':
+  main()

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,47 +118,35 @@ 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)
@@ -117,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)
 
@@ -127,7 +171,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,14 +208,36 @@ 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:
@@ -179,9 +245,11 @@ 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(
+  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:
@@ -249,19 +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
   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 / バッチサイズ: {', '.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"  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}")
 
-  # TODO refactor metadata creation and move to util
   metadata = {
       "ss_session_id": session_id,            # random integer indicating which group of epochs the model came from
       "ss_training_started_at": training_started_at,          # unix timestamp
@@ -269,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,
@@ -284,149 +352,29 @@ 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),
-    })
-
   # uncomment if another network is added
   # for key, value in net_kwargs.items():
   #   metadata["ss_arg_" + key] = value
@@ -462,7 +410,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)
 
     metadata["ss_epoch"] = str(epoch+1)
 
@@ -499,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:
@@ -517,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()
@@ -530,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)
@@ -562,7 +508,6 @@ 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, None if args.no_metadata else metadata)
 
@@ -577,12 +522,9 @@ def train(args):
       if saving and args.save_state:
         train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)
 
-    train_util.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
-
     # 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:
@@ -613,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"],
@@ -622,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,25 +163,20 @@ def train(args):
     captions = []
     for tmpl in templates:
       captions.append(tmpl.format(replace_to))
-    train_dataset_group.add_replacement("", captions)
-  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.add_replacement("", captions)
+  elif args.num_vectors_per_token > 1:
+    replace_to = " ".join(token_strings)
+    train_dataset.add_replacement(args.token_string, replace_to)
+
+  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)
 
@@ -208,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()
@@ -220,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:
@@ -235,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(
@@ -285,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}")
@@ -305,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():
@@ -358,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()
@@ -375,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
@@ -400,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
@@ -421,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
@@ -498,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）")