Add multires_training arg for networks

- multires_training allows creation of buckets down to the set minimum bucket resolution defined.

Author: 67372a

finetune/prepare_buckets_latents.py

@@ -97,7 +97,7 @@ def main(args):
     ), f"illegal resolution (not 'width,height') / 画像サイズに誤りがあります。'幅,高さ'で指定してください: {args.max_resolution}"
 
     bucket_manager = train_util.BucketManager(
-        args.bucket_no_upscale, max_reso, args.min_bucket_reso, args.max_bucket_reso, args.bucket_reso_steps
+        args.bucket_no_upscale, max_reso, args.min_bucket_reso, args.max_bucket_reso, args.bucket_reso_steps, args.multires_training
     )
     if not args.bucket_no_upscale:
         bucket_manager.make_buckets()
@@ -242,6 +242,11 @@ def setup_parser() -> argparse.ArgumentParser:
         action="store_true",
         help="make bucket for each image without upscaling / 画像を拡大せずbucketを作成します",
     )
+    parser.add_argument(
+        "--multires_training",
+        action="store_true",
+        help="make buckets for all resolutions down to minimum resolution for multires training",
+    )
     parser.add_argument(
         "--mixed_precision",
         type=str,

library/config_util.py

@@ -112,6 +112,7 @@ class BaseDatasetParams:
     validation_seed: Optional[int] = None
     validation_split: float = 0.0
     resize_interpolation: Optional[str] = None
+    multires_training: bool = False
 
 @dataclass
 class DreamBoothDatasetParams(BaseDatasetParams):
@@ -251,6 +252,7 @@ def __validate_and_convert_scalar_or_twodim(klass, value: Union[float, Sequence]
         "resolution": functools.partial(__validate_and_convert_scalar_or_twodim.__func__, int),
         "network_multiplier": float,
         "resize_interpolation": str,
+        "multires_training": bool,
     }
 
     # options handled by argparse but not handled by user config
@@ -560,6 +562,7 @@ def print_info(_datasets, dataset_type: str):
                   max_bucket_reso: {dataset.max_bucket_reso}
                   bucket_reso_steps: {dataset.bucket_reso_steps}
                   bucket_no_upscale: {dataset.bucket_no_upscale}
+                  multires_training: {dataset.multires_training}
                 \n"""), "  ")
             else:
                 info += "\n"

library/model_util.py

@@ -1333,7 +1333,7 @@ def use_reflection_padding(vae):
 # endregion
 
 
-def make_bucket_resolutions(max_reso, min_size=256, max_size=1024, divisible=64):
+def make_bucket_resolutions(max_reso, min_size=256, max_size=1024, divisible=64, multires_training=False):
     max_width, max_height = max_reso
     max_area = max_width * max_height
 
@@ -1344,18 +1344,19 @@ def make_bucket_resolutions(max_reso, min_size=256, max_size=1024, divisible=64)
 
     width = min_size
     while width <= max_size:
-        height = min(max_size, int((max_area // width) // divisible) * divisible)
-        if height >= min_size:
-            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))
+        max_h = min(max_size, int((max_area // width) // divisible) * divisible)
+        
+        if not multires_training:
+            height = max_h
+            if height >= min_size:
+                resos.add((width, height))
+                resos.add((height, width))
+        else:
+            height = min_size
+            while height <= max_h:
+                resos.add((width, height))
+                resos.add((height, width))
+                height += divisible
 
         width += divisible
 

library/train_util.py

@@ -226,7 +226,8 @@ def __init__(
 
 
 class BucketManager:
-    def __init__(self, no_upscale, max_reso, min_size, max_size, reso_steps) -> None:
+    def __init__(self, no_upscale, max_reso, min_size, max_size, reso_steps, multires_training=False) -> None:
+        self.multires_training = multires_training
         if max_size is not None:
             if max_reso is not None:
                 assert max_size >= max_reso[0], "the max_size should be larger than the width of max_reso"
@@ -274,7 +275,7 @@ def sort(self):
         self.reso_to_id = sorted_reso_to_id
 
     def make_buckets(self):
-        resos = model_util.make_bucket_resolutions(self.max_reso, self.min_size, self.max_size, self.reso_steps)
+        resos = model_util.make_bucket_resolutions(self.max_reso, self.min_size, self.max_size, self.reso_steps, self.multires_training)
         self.set_predefined_resos(resos)
 
     def set_predefined_resos(self, resos):
@@ -304,9 +305,19 @@ def select_bucket(self, 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_errors = np.abs(self.predefined_aspect_ratios - aspect_ratio)
+                if getattr(self, "multires_training", False):
+                    min_ar_error = ar_errors.min()
+                    # filter out the closest aspect ratios
+                    closest_indices = np.where(ar_errors <= min_ar_error + 1e-4)[0]
+                    # from these, find the one with the closest area
+                    target_area = image_width * image_height
+                    areas = np.array([self.predefined_resos[i][0] * self.predefined_resos[i][1] for i in closest_indices])
+                    best_index = closest_indices[np.abs(areas - target_area).argmin()]
+                    reso = self.predefined_resos[best_index]
+                else:
+                    predefined_bucket_id = ar_errors.argmin()  # 当該解像度以外でaspect ratio errorが最も少ないもの
+                    reso = self.predefined_resos[predefined_bucket_id]
 
             ar_reso = reso[0] / reso[1]
             if aspect_ratio > ar_reso:  # 横が長い→縦を合わせる
@@ -1118,6 +1129,7 @@ def make_buckets(self):
                     self.min_bucket_reso,
                     self.max_bucket_reso,
                     self.bucket_reso_steps,
+                    getattr(self, "multires_training", False),
                 )
                 if not self.bucket_no_upscale:
                     self.bucket_manager.make_buckets()
@@ -2023,6 +2035,7 @@ def __init__(
         max_bucket_reso: int,
         bucket_reso_steps: int,
         bucket_no_upscale: bool,
+        multires_training: bool,
         prior_loss_weight: float,
         debug_dataset: bool,
         validation_split: float,
@@ -2050,11 +2063,13 @@ def __init__(
             self.max_bucket_reso = max_bucket_reso
             self.bucket_reso_steps = bucket_reso_steps
             self.bucket_no_upscale = bucket_no_upscale
+            self.multires_training = multires_training
         else:
             self.min_bucket_reso = None
             self.max_bucket_reso = None
             self.bucket_reso_steps = None  # この情報は使われない
             self.bucket_no_upscale = False
+            self.multires_training = False
 
         def read_caption(img_path, caption_extension, enable_wildcard):
             # captionの候補ファイル名を作る
@@ -2333,6 +2348,7 @@ def __init__(
         max_bucket_reso: int,
         bucket_reso_steps: int,
         bucket_no_upscale: bool,
+        multires_training: bool,
         debug_dataset: bool,
         validation_seed: int,
         validation_split: float,
@@ -2353,11 +2369,13 @@ def __init__(
             self.max_bucket_reso = max_bucket_reso
             self.bucket_reso_steps = bucket_reso_steps
             self.bucket_no_upscale = bucket_no_upscale
+            self.multires_training = multires_training
         else:
             self.min_bucket_reso = None
             self.max_bucket_reso = None
             self.bucket_reso_steps = None  # この情報は使われない
             self.bucket_no_upscale = False
+            self.multires_training = False
 
         self.num_train_images = 0
         self.num_reg_images = 0
@@ -2522,6 +2540,7 @@ def __init__(
         max_bucket_reso: int,
         bucket_reso_steps: int,
         bucket_no_upscale: bool,
+        multires_training: bool,
         debug_dataset: bool,
         validation_split: float,
         validation_seed: Optional[int],
@@ -2577,6 +2596,7 @@ def __init__(
             max_bucket_reso,
             bucket_reso_steps,
             bucket_no_upscale,
+            multires_training,
             1.0,
             debug_dataset,
             validation_split,
@@ -4821,6 +4841,11 @@ def add_dataset_arguments(
         action="store_true",
         help="make bucket for each image without upscaling / 画像を拡大せずbucketを作成します",
     )
+    parser.add_argument(
+        "--multires_training",
+        action="store_true",
+        help="make buckets for all resolutions down to minimum resolution for multires training",
+    )
     parser.add_argument(
         "--resize_interpolation",
         type=str,

train_network.py

@@ -1623,6 +1623,7 @@ def load_model_hook(models, input_dir):
                     "ss_shuffle_caption": bool(args.shuffle_caption),
                     "ss_enable_bucket": bool(dataset.enable_bucket),
                     "ss_bucket_no_upscale": bool(dataset.bucket_no_upscale),
+                    "ss_multires_training": bool(getattr(dataset, "multires_training", False)),
                     "ss_min_bucket_reso": dataset.min_bucket_reso,
                     "ss_max_bucket_reso": dataset.max_bucket_reso,
                     "ss_keep_tokens": args.keep_tokens,