From ecb12378da48fc3a17539d5cc33edc561cf8a426 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Sun, 8 Jan 2023 20:33:04 +0100
Subject: Improved aspect ratio bucketing

---
 train_dreambooth.py | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

(limited to 'train_dreambooth.py')

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 79eede6..d396249 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -103,6 +103,29 @@ def parse_args():
         default=999999,
         help="Number of epochs the text encoder will be trained."
     )
+    parser.add_argument(
+        "--num_buckets",
+        type=int,
+        default=4,
+        help="Number of aspect ratio buckets in either direction.",
+    )
+    parser.add_argument(
+        "--progressive_buckets",
+        action="store_true",
+        help="Include images in smaller buckets as well.",
+    )
+    parser.add_argument(
+        "--bucket_step_size",
+        type=int,
+        default=64,
+        help="Step size between buckets.",
+    )
+    parser.add_argument(
+        "--bucket_max_pixels",
+        type=int,
+        default=None,
+        help="Maximum pixels per bucket.",
+    )
     parser.add_argument(
         "--tag_dropout",
         type=float,
@@ -734,6 +757,10 @@ def main():
         class_subdir=args.class_image_dir,
         num_class_images=args.num_class_images,
         size=args.resolution,
+        num_buckets=args.num_buckets,
+        progressive_buckets=args.progressive_buckets,
+        bucket_step_size=args.bucket_step_size,
+        bucket_max_pixels=args.bucket_max_pixels,
         dropout=args.tag_dropout,
         template_key=args.train_data_template,
         valid_set_size=args.valid_set_size,
-- 
cgit v1.2.3-54-g00ecf