generate a folder on the fly for prior preservation

train_dreambooth_lora_sdxl_advanced.py

@@ -22,6 +22,7 @@ import math
 import os
 import random
 import re
+import uuid
 import shutil
 import warnings
 from contextlib import nullcontext
@@ -827,17 +828,17 @@ def parse_args(input_args=None):
     if env_local_rank != -1 and env_local_rank != args.local_rank:
         args.local_rank = env_local_rank
 
-    # if args.with_prior_preservation:
-    #     if args.class_data_dir is None:
-    #         raise ValueError("You must specify a data directory for class images.")
-    #     if args.class_prompt is None:
-    #         raise ValueError("You must specify prompt for class images.")
-    # else:
-    #     # logger is not available yet
-    #     if args.class_data_dir is not None:
-    #         warnings.warn("You need not use --class_data_dir without --with_prior_preservation.")
-    #     if args.class_prompt is not None:
-    #         warnings.warn("You need not use --class_prompt without --with_prior_preservation.")
+    if args.with_prior_preservation:
+        if args.class_data_dir is None:
+            raise ValueError("You must specify a data directory for class images.")
+        if args.class_prompt is None:
+            raise ValueError("You must specify prompt for class images.")
+    else:
+        # logger is not available yet
+        if args.class_data_dir is not None:
+            warnings.warn("You need not use --class_data_dir without --with_prior_preservation.")
+        if args.class_prompt is not None:
+            warnings.warn("You need not use --class_prompt without --with_prior_preservation.")
 
     return args
 
@@ -1321,6 +1322,9 @@ def main(args):
 
     # Generate class images if prior preservation is enabled.
     if args.with_prior_preservation:
+        if arg.class_data_dir == None:
+            class_folder = str(uuid.uuid4())
+            args.class_data_dir = os.path.join("ariadne", class_folder)
         class_images_dir = Path(args.class_data_dir)
         if not class_images_dir.exists():
             class_images_dir.mkdir(parents=True)