Fix

1 files changed, 7 insertions, 14 deletions

diff --git a/train_ti.py b/train_ti.py
index 7f5fb49..45e730a 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -484,19 +484,13 @@ def parse_args():
         help="The weight of prior preservation loss."
     )
     parser.add_argument(
-        "--lora_r",
-        type=int,
-        default=8,
-        help="Lora rank, only used if use_lora is True"
-    )
-    parser.add_argument(
-        "--lora_alpha",
-        type=int,
-        default=32,
-        help="Lora alpha, only used if use_lora is True"
+        "--emb_alpha",
+        type=float,
+        default=1.0,
+        help="Embedding alpha"
     )
     parser.add_argument(
-        "--lora_dropout",
+        "--emb_dropout",
         type=float,
         default=0,
         help="Embedding dropout probability.",
@@ -669,9 +663,8 @@ def main():
 
     tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler, embeddings = get_models(
         args.pretrained_model_name_or_path,
-        args.lora_r,
-        args.lora_alpha,
-        args.lora_dropout
+        args.emb_alpha,
+        args.emb_dropout
     )
 
     tokenizer.set_use_vector_shuffle(args.vector_shuffle)