Delete append_module.py

append_module.py

@@ -1,504 +0,0 @@
-import argparse
-import json
-import shutil
-import time
-from typing import Dict, List, NamedTuple, Tuple
-from accelerate import Accelerator
-from torch.autograd.function import Function
-import glob
-import math
-import os
-import random
-import hashlib
-from io import BytesIO
-
-from tqdm import tqdm
-import torch
-from torchvision import transforms
-from transformers import CLIPTokenizer
-import diffusers
-from diffusers import DDPMScheduler, StableDiffusionPipeline
-import albumentations as albu
-import numpy as np
-from PIL import Image
-import cv2
-from einops import rearrange
-from torch import einsum
-import safetensors.torch
-
-import library.model_util as model_util
-import library.train_util as train_util
-
-#============================================================================================================
-#AdafactorScheduleに暫定的にinitial_lrを層別に適用できるようにしたもの
-#============================================================================================================
-from torch.optim.lr_scheduler import LambdaLR
-class AdafactorSchedule_append(LambdaLR):
-    """
-    Since [`~optimization.Adafactor`] performs its own scheduling, if the training loop relies on a scheduler (e.g.,
-    for logging), this class creates a proxy object that retrieves the current lr values from the optimizer.
-
-    It returns `initial_lr` during startup and the actual `lr` during stepping.
-    """
-
-    def __init__(self, optimizer, initial_lr=0.0):
-        def lr_lambda(_):
-            return initial_lr
-
-        for group in optimizer.param_groups:
-            if not type(initial_lr)==list:
-                group["initial_lr"] = initial_lr
-            else:
-                group["initial_lr"] = initial_lr.pop(0)
-        super().__init__(optimizer, lr_lambda)
-        for group in optimizer.param_groups:
-            del group["initial_lr"]
-
-    def get_lr(self):
-        opt = self.optimizer
-        lrs = [
-            opt._get_lr(group, opt.state[group["params"][0]])
-            for group in opt.param_groups
-            if group["params"][0].grad is not None
-        ]
-        if len(lrs) == 0:
-            lrs = self.base_lrs  # if called before stepping
-        return lrs
-
-#============================================================================================================
-#model_util 内より
-#============================================================================================================
-def make_bucket_resolutions_fix(max_reso, min_reso, min_size=256, max_size=1024, divisible=64, step=1):
-  max_width, max_height = max_reso
-  max_area = (max_width // divisible) * (max_height // divisible)
-
-  min_widht, min_height = min_reso
-  min_area = (min_widht // divisible) * (min_height // divisible)
-  
-  area_size_list = []
-  area_size_resos_list = []
-  _max_area = max_area
-
-  while True:
-    resos = set()
-    size = int(math.sqrt(_max_area)) * divisible
-    resos.add((size, size))
-
-    size = min_size
-    while size <= max_size:
-      width = size
-      height = min(max_size, (_max_area // (width // divisible)) * divisible)
-      resos.add((width, height))
-      resos.add((height, width))
-
-      # # make additional resos
-      # if width >= height and width - divisible >= min_size:
-      #   resos.add((width - divisible, height))
-      #   resos.add((height, width - divisible))
-      # if height >= width and height - divisible >= min_size:
-      #   resos.add((width, height - divisible))
-      #   resos.add((height - divisible, width))
-
-      size += divisible
-
-    resos = list(resos)
-    resos.sort()
-
-    #aspect_ratios = [w / h for w, h in resos]
-    area_size_list.append(_max_area)
-    area_size_resos_list.append(resos)
-    #area_size_ratio_list.append(aspect_ratios)
-
-    _max_area -= step
-    if _max_area < min_area:
-      break
-  return area_size_resos_list, area_size_list
-
-#============================================================================================================
-#train_util 内より
-#============================================================================================================
-class BucketManager_append(train_util.BucketManager):
-  def __init__(self, no_upscale, max_reso, min_size, max_size, reso_steps, min_reso=None, area_step=None) -> None:
-    super().__init__(no_upscale, max_reso, min_size, max_size, reso_steps)
-    print("BucketManager_appendを作成しました")
-    if min_reso is None:
-      self.min_reso = None
-      self.min_area = None
-    else:
-      self.min_reso = min_reso
-      self.min_area = min_reso[0] * min_reso[1]
-    self.area_step = area_step
-    self.area_sizes_flag = False
-
-  def make_buckets(self):
-    if self.min_reso:
-      print(f"make_resolution append")
-      resos, area_sizes = make_bucket_resolutions_fix(self.max_reso, self.min_reso, self.min_size, self.max_size, self.reso_steps, self.area_step)
-      self.set_predefined_resos(resos, area_sizes)
-    else:
-      resos = model_util.make_bucket_resolutions(self.max_reso, self.min_size, self.max_size, self.reso_steps)
-      self.set_predefined_resos(resos)
-
-  def set_predefined_resos(self, resos, area_sizes=None):
-    # 規定サイズから選ぶ場合の解像度、aspect ratioの情報を格納しておく
-    if area_sizes:
-      self.area_sizes_flag = True
-      self.predefined_area_sizes = np.array(area_sizes.copy())
-      self.predefined_resos_list = resos.copy()
-      self.predefined_resos_set_list = [set(reso) for reso in resos]
-      self.predefined_aspect_ratios_list = [np.array([w/h for w,h in reso]) for reso in resos]
-      self.predefined_resos = None
-      self.predefined_resos_set = None
-      self.predefined_aspect_ratios = None
-    else:
-      self.area_sizes_flag = False
-      self.predefined_area_sizes = None
-      self.predefined_resos = resos.copy()
-      self.predefined_resos_set = set(resos)
-      self.predefined_aspect_ratios = np.array([w / h for w, h in resos])
-
-  def select_bucket(self, image_width, image_height):
-    # 画像サイズを算出する
-    area_size = (image_width//64) * (image_height//64)
-    aspect_ratio = image_width / image_height
-    bucket_size_id = None
-    # 拡張したバケットサイズを使うために画像サイズのエリアを決定する
-    if self.area_sizes_flag:
-        size_errors = self.predefined_area_sizes - area_size
-        bucket_size_id = np.abs(size_errors).argmin()
-        #一定の範囲を探索して使用する画像サイズを確定する
-        serch_size_range = 1
-        bucket_size_id_list = [bucket_size_id]
-        for i in range(serch_size_range):
-          if bucket_size_id - i <0:
-            bucket_size_id_list.append(bucket_size_id + i + 1)
-          elif bucket_size_id + 1 + i >= len(self.predefined_resos_list):
-            bucket_size_id_list.append(bucket_size_id - i - 1)
-          else:
-            bucket_size_id_list.append(bucket_size_id - i - 1)
-            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:
-          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()
-          if _min_error > ar_error:
-            _min_error = ar_error
-            _min_id = now_size_id
-          if _min_error == 0.:
-            break
-        bucket_size_id = _min_id
-        del _min_error, _min_id, ar_error #余計なものは掃除
-        self.predefined_resos = self.predefined_resos_list[bucket_size_id]
-        self.predefined_resos_set = self.predefined_resos_set_list[bucket_size_id]
-        self.predefined_aspect_ratios = self.predefined_aspect_ratios_list[bucket_size_id]
-    # --ここから処理はそのまま
-    if not self.no_upscale:
-      # 同じaspect ratioがあるかもしれないので（fine tuningで、no_upscale=Trueで前処理した場合）、解像度が同じものを優先する
-      reso = (image_width, image_height)
-      if reso in self.predefined_resos_set:
-        pass
-      else:
-        ar_errors = self.predefined_aspect_ratios - aspect_ratio
-        predefined_bucket_id = np.abs(ar_errors).argmin()          # 当該解像度以外でaspect ratio errorが最も少ないもの
-        reso = self.predefined_resos[predefined_bucket_id]
-
-      ar_reso = reso[0] / reso[1]
-      if aspect_ratio > ar_reso:                   # 横が長い→縦を合わせる
-        scale = reso[1] / image_height
-      else:
-        scale = reso[0] / image_width
-
-      resized_size = (int(image_width * scale + .5), int(image_height * scale + .5))
-      # print("use predef", image_width, image_height, reso, resized_size)
-    else:
-      if image_width * image_height > self.max_area:
-        # 画像が大きすぎるのでアスペクト比を保ったまま縮小することを前提にbucketを決める
-        resized_width = math.sqrt(self.max_area * aspect_ratio)
-        resized_height = self.max_area / resized_width
-        assert abs(resized_width / resized_height - aspect_ratio) < 1e-2, "aspect is illegal"
-
-        # リサイズ後の短辺または長辺をreso_steps単位にする：aspect ratioの差が少ないほうを選ぶ
-        # 元のbucketingと同じロジック
-        b_width_rounded = self.round_to_steps(resized_width)
-        b_height_in_wr = self.round_to_steps(b_width_rounded / aspect_ratio)
-        ar_width_rounded = b_width_rounded / b_height_in_wr
-
-        b_height_rounded = self.round_to_steps(resized_height)
-        b_width_in_hr = self.round_to_steps(b_height_rounded * aspect_ratio)
-        ar_height_rounded = b_width_in_hr / b_height_rounded
-
-        # print(b_width_rounded, b_height_in_wr, ar_width_rounded)
-        # print(b_width_in_hr, b_height_rounded, ar_height_rounded)
-
-        if abs(ar_width_rounded - aspect_ratio) < abs(ar_height_rounded - aspect_ratio):
-          resized_size = (b_width_rounded, int(b_width_rounded / aspect_ratio + .5))
-        else:
-          resized_size = (int(b_height_rounded * aspect_ratio + .5), b_height_rounded)
-        # print(resized_size)
-      else:
-        resized_size = (image_width, image_height)              # リサイズは不要
-
-      # 画像のサイズ未満をbucketのサイズとする（paddingせずにcroppingする）
-      bucket_width = resized_size[0] - resized_size[0] % self.reso_steps
-      bucket_height = resized_size[1] - resized_size[1] % self.reso_steps
-      # print("use arbitrary", image_width, image_height, resized_size, bucket_width, bucket_height)
-
-      reso = (bucket_width, bucket_height)
-
-    self.add_if_new_reso(reso)
-
-    ar_error = (reso[0] / reso[1]) - aspect_ratio
-    return reso, resized_size, ar_error
-
-class DreamBoothDataset(train_util.DreamBoothDataset):
-  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__(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を作る）
-    min_size and max_size are ignored when enable_bucket is False
-    '''
-    print("loading image sizes.")
-    for info in tqdm(self.image_data.values()):
-      if info.image_size is None:
-        info.image_size = self.get_image_size(info.absolute_path)
-
-    if self.enable_bucket:
-      print("make buckets")
-    else:
-      print("prepare dataset")
-
-    # bucketを作成し、画像をbucketに振り分ける
-    if self.enable_bucket:
-      if self.bucket_manager is None:                         # fine tuningの場合でmetadataに定義がある場合は、すでに初期化済み
-        #======================================================================change
-        if self.min_resolution:
-          self.bucket_manager = BucketManager_append(self.bucket_no_upscale, (self.width, self.height),
-                                              self.min_bucket_reso, self.max_bucket_reso, self.bucket_reso_steps, self.min_resolution, self.area_step)
-        else:
-          self.bucket_manager = train_util.BucketManager(self.bucket_no_upscale, (self.width, self.height),
-                                              self.min_bucket_reso, self.max_bucket_reso, self.bucket_reso_steps)
-        #======================================================================change
-        if not self.bucket_no_upscale:
-          self.bucket_manager.make_buckets()
-        else:
-          print("min_bucket_reso and max_bucket_reso are ignored if bucket_no_upscale is set, because bucket reso is defined by image size automatically / bucket_no_upscaleが指定された場合は、bucketの解像度は画像サイズから自動計算されるため、min_bucket_resoとmax_bucket_resoは無視されます")
-
-      img_ar_errors = []
-      for image_info in self.image_data.values():
-        image_width, image_height = image_info.image_size
-        image_info.bucket_reso, image_info.resized_size, ar_error = self.bucket_manager.select_bucket(image_width, image_height)
-
-        # print(image_info.image_key, image_info.bucket_reso)
-        img_ar_errors.append(abs(ar_error))
-
-      self.bucket_manager.sort()
-    else:
-      self.bucket_manager = train_util.BucketManager(False, (self.width, self.height), None, None, None)
-      self.bucket_manager.set_predefined_resos([(self.width, self.height)])  # ひとつの固定サイズbucketのみ
-      for image_info in self.image_data.values():
-        image_width, image_height = image_info.image_size
-        image_info.bucket_reso, image_info.resized_size, _ = self.bucket_manager.select_bucket(image_width, image_height)
-
-    for image_info in self.image_data.values():
-      for _ in range(image_info.num_repeats):
-        self.bucket_manager.add_image(image_info.bucket_reso, image_info.image_key)
-
-    # bucket情報を表示、格納する
-    if self.enable_bucket:
-      self.bucket_info = {"buckets": {}}
-      print("number of images (including repeats) / 各bucketの画像枚数（繰り返し回数を含む）")
-      for i, (reso, bucket) in enumerate(zip(self.bucket_manager.resos, self.bucket_manager.buckets)):
-        count = len(bucket)
-        if count > 0:
-          self.bucket_info["buckets"][i] = {"resolution": reso, "count": len(bucket)}
-          print(f"bucket {i}: resolution {reso}, count: {len(bucket)}")
-
-      img_ar_errors = np.array(img_ar_errors)
-      mean_img_ar_error = np.mean(np.abs(img_ar_errors))
-      self.bucket_info["mean_img_ar_error"] = mean_img_ar_error
-      print(f"mean ar error (without repeats): {mean_img_ar_error}")
-
-    # データ参照用indexを作る。このindexはdatasetのshuffleに用いられる
-    self.buckets_indices: List(train_util.BucketBatchIndex) = []
-    for bucket_index, bucket in enumerate(self.bucket_manager.buckets):
-      batch_count = int(math.ceil(len(bucket) / self.batch_size))
-      for batch_index in range(batch_count):
-        self.buckets_indices.append(train_util.BucketBatchIndex(bucket_index, self.batch_size, batch_index))
-
-      # ↓以下はbucketごとのbatch件数があまりにも増えて混乱を招くので元に戻す
-      # 　学習時はステップ数がランダムなので、同一画像が同一batch内にあってもそれほど悪影響はないであろう、と考えられる
-      #
-      # # bucketが細分化されることにより、ひとつのbucketに一種類の画像のみというケースが増え、つまりそれは
-      # # ひとつのbatchが同じ画像で占められることになるので、さすがに良くないであろう
-      # # そのためバッチサイズを画像種類までに制限する
-      # # ただそれでも同一画像が同一バッチに含まれる可能性はあるので、繰り返し回数が少ないほうがshuffleの品質は良くなることは間違いない？
-      # # TO DO 正則化画像をepochまたがりで利用する仕組み
-      # num_of_image_types = len(set(bucket))
-      # bucket_batch_size = min(self.batch_size, num_of_image_types)
-      # batch_count = int(math.ceil(len(bucket) / bucket_batch_size))
-      # # print(bucket_index, num_of_image_types, bucket_batch_size, batch_count)
-      # for batch_index in range(batch_count):
-      #   self.buckets_indices.append(BucketBatchIndex(bucket_index, bucket_batch_size, batch_index))
-      # ↑ここまで
-
-    self.shuffle_buckets()
-    self._length = len(self.buckets_indices)
-
-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):
-  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:
-          if lora_name in lora.lora_name:
-            params.extend(lora.parameters())
-        else:
-          params.extend(lora.parameters())
-      return params
-
-    self.requires_grad_(True)
-    all_params = []
-    ret_scheduler_lr = []
-    
-    if loranames is not None:
-      textencoder_names = [None]
-      unet_names = [None]
-      if "text_encoder" in loranames:
-        textencoder_names = loranames["text_encoder"]
-      if "unet" in loranames:
-        unet_names = loranames["unet"]
-    
-    if self.text_encoder_loras:
-      for textencoder_name in textencoder_names:
-        param_data = {'params': enumerate_params(self.text_encoder_loras, lora_name=textencoder_name)}
-        if text_encoder_lr is not None:
-          param_data['lr'] = text_encoder_lr
-        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 unet_lr is not None:
-          param_data['lr'] = unet_lr
-        if scheduler_lr is not None:
-          ret_scheduler_lr.append(scheduler_lr[1])
-        all_params.append(param_data)
-
-    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("--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("--min_resolution", type=str, default=None)
-  parser.add_argument("--area_step", type=int, default=1)
-  parser.add_argument("--config", type=str, default=None)
-
-def create_split_names(split_flag, split_level):
-  split_names = None
-  if split_flag:
-    split_names = {}
-    text_encoder_names = [None]
-    unet_names = ["lora_unet_mid_block"]
-    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:
-      if split_level==2:
-        text_encoder_names = []
-        for i in range(12):
-          text_encoder_names.append(f"lora_te_text_model_encoder_layers_{i}_")
-      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
-
-def get_config(parser):
-  args = parser.parse_args()
-  if args.config is not None and (not args.config==""):
-    import yaml
-    import datetime
-    if os.path.splitext(args.config)[-1] == ".yaml":
-      args.config = os.path.splitext(args.config)[0]
-    config_path = f"./{args.config}.yaml"
-    if os.path.exists(config_path):
-      print(f"{config_path} から設定を読み込み中...")
-      margs, rest = parser.parse_known_args()
-      with open(config_path, mode="r") as f:
-          configs = yaml.unsafe_load(f)
-      #変数でのやり取りをするためargparserからDict型を取り出す
-      args_dic = vars(args)
-      #デフォから引数指定で変更があるものを確認
-      change_def_dic = {}
-      args_type_dic = {}
-      for key, v in args_dic.items():
-        if not parser.get_default(key) == v:
-          change_def_dic[key] = v
-      #デフォ指定されてるデータ型を取得する
-      for key, act in parser._option_string_actions.items():
-        if key=="-h": continue
-        key = key[2:]
-        args_type_dic[key] = act.type
-      #データタイプの確認とargsにkeyの内容を代入していく
-      for key, v in configs.items():
-        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
-      #最後にデフォから指定が変わってるものを変更する
-      for key, v in change_def_dic.items():
-        args_dic[key] = v
-    else:
-      print(f"{config_path} が見つかりませんでした")
-  return args