1 files changed, 1 insertions, 45 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index a70c80e..5a4c47b 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -74,26 +74,6 @@ def parse_args():
         help="The name of the current project.",
     )
     parser.add_argument(
-        "--placeholder_tokens",
-        type=str,
-        nargs='*',
-        default=[],
-        help="A token to use as a placeholder for the concept.",
-    )
-    parser.add_argument(
-        "--initializer_tokens",
-        type=str,
-        nargs='*',
-        default=[],
-        help="A token to use as initializer word."
-    )
-    parser.add_argument(
-        "--num_vectors",
-        type=int,
-        nargs='*',
-        help="Number of vectors per embedding."
-    )
-    parser.add_argument(
         "--exclude_collections",
         type=str,
         nargs='*',
@@ -436,30 +416,6 @@ def parse_args():
     if args.project is None:
         raise ValueError("You must specify --project")
 
-    if isinstance(args.placeholder_tokens, str):
-        args.placeholder_tokens = [args.placeholder_tokens]
-
-    if isinstance(args.initializer_tokens, str):
-        args.initializer_tokens = [args.initializer_tokens] * len(args.placeholder_tokens)
-
-    if len(args.initializer_tokens) == 0:
-        raise ValueError("You must specify --initializer_tokens")
-
-    if len(args.placeholder_tokens) == 0:
-        args.placeholder_tokens = [f"<*{i}>" for i in range(len(args.initializer_tokens))]
-
-    if len(args.placeholder_tokens) != len(args.initializer_tokens):
-        raise ValueError("--placeholder_tokens and --initializer_tokens must have the same number of items")
-
-    if args.num_vectors is None:
-        args.num_vectors = 1
-
-    if isinstance(args.num_vectors, int):
-        args.num_vectors = [args.num_vectors] * len(args.initializer_tokens)
-
-    if len(args.placeholder_tokens) != len(args.num_vectors):
-        raise ValueError("--placeholder_tokens and --num_vectors must have the same number of items")
-
     if isinstance(args.collection, str):
         args.collection = [args.collection]
 
@@ -503,7 +459,7 @@ def main():
 
     vae.enable_slicing()
     vae.set_use_memory_efficient_attention_xformers(True)
-    unet.set_use_memory_efficient_attention_xformers(True)
+    unet.enable_xformers_memory_efficient_attention()
 
     if args.gradient_checkpointing:
         unet.enable_gradient_checkpointing()