1 files changed, 9 insertions, 1 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 1a79b2b..c7899a0 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -105,6 +105,12 @@ def parse_args():
         help="How many class images to generate."
     )
     parser.add_argument(
+        "--class_image_dir",
+        type=str,
+        default="cls",
+        help="The directory where class images will be saved.",
+    )
+    parser.add_argument(
         "--repeats",
         type=int,
         default=1,
@@ -653,7 +659,7 @@ def main():
         data_file=args.train_data_file,
         batch_size=args.train_batch_size,
         prompt_processor=prompt_processor,
-        class_subdir="cls",
+        class_subdir=args.class_image_dir,
         num_class_images=args.num_class_images,
         size=args.resolution,
         repeats=args.repeats,
@@ -696,6 +702,8 @@ def main():
                     images = pipeline(
                         prompt=prompt,
                         negative_prompt=nprompt,
+                        height=args.sample_image_size,
+                        width=args.sample_image_size,
                         num_inference_steps=args.sample_steps
                     ).images