Fix

1 files changed, 10 insertions, 1 deletions

diff --git a/train_lora.py b/train_lora.py
index d5dde02..5c78664 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -507,6 +507,12 @@ def parse_args():
         help="The weight of prior preservation loss."
     )
     parser.add_argument(
+        "--emb_alpha",
+        type=float,
+        default=1.0,
+        help="Embedding alpha"
+    )
+    parser.add_argument(
         "--emb_dropout",
         type=float,
         default=0,
@@ -660,7 +666,10 @@ def main():
     save_args(output_dir, args)
 
     tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler, embeddings = get_models(
-        args.pretrained_model_name_or_path, args.emb_dropout)
+        args.pretrained_model_name_or_path,
+        args.emb_alpha,
+        args.emb_dropout
+    )
 
     unet_config = LoraConfig(
         r=args.lora_r,