From 6970adaff742ac89adb3d85c803689210dc030e2 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Sat, 7 Jan 2023 15:05:39 +0100
Subject: Made aspect ratio bucketing configurable

---
 train_ti.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

(limited to 'train_ti.py')

diff --git a/train_ti.py b/train_ti.py
index 89c6672..38c9755 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -145,6 +145,17 @@ def parse_args():
             " resolution"
         ),
     )
+    parser.add_argument(
+        "--num_aspect_ratio_buckets",
+        type=int,
+        default=4,
+        help="Number of buckets in either direction (adds 64 pixels per step).",
+    )
+    parser.add_argument(
+        "--progressive_aspect_ratio_buckets",
+        action="store_true",
+        help="Include images in smaller buckets as well.",
+    )
     parser.add_argument(
         "--tag_dropout",
         type=float,
@@ -710,6 +721,8 @@ def main():
         class_subdir=args.class_image_dir,
         num_class_images=args.num_class_images,
         size=args.resolution,
+        num_aspect_ratio_buckets=args.num_aspect_ratio_buckets,
+        progressive_aspect_ratio_buckets=args.progressive_aspect_ratio_buckets,
         repeats=args.repeats,
         dropout=args.tag_dropout,
         template_key=args.train_data_template,
-- 
cgit v1.2.3-54-g00ecf