Update

1 files changed, 17 insertions, 0 deletions

diff --git a/train_lora.py b/train_lora.py
index cc7c1ec..70fbae4 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -27,6 +27,7 @@ from data.csv import VlpnDataModule, keyword_filter
 from training.functional import train, add_placeholder_tokens, get_models
 from training.strategy.lora import lora_strategy
 from training.optimization import get_scheduler
+from training.sampler import create_named_schedule_sampler
 from training.util import AverageMeter, save_args
 
 # https://github.com/huggingface/peft/blob/main/examples/lora_dreambooth/train_dreambooth.py
@@ -410,6 +411,19 @@ def parse_args():
         help="Minimum learning rate in the lr scheduler."
     )
     parser.add_argument(
+        "--min_snr_gamma",
+        type=int,
+        default=5,
+        help="MinSNR gamma."
+    )
+    parser.add_argument(
+        "--schedule_sampler",
+        type=str,
+        default="uniform",
+        choices=["uniform", "loss-second-moment"],
+        help="Noise schedule sampler."
+    )
+    parser.add_argument(
         "--optimizer",
         type=str,
         default="adan",
@@ -708,6 +722,7 @@ def main():
         args.emb_alpha,
         args.emb_dropout
     )
+    schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, noise_scheduler.config.num_train_timesteps)
 
     unet_config = LoraConfig(
         r=args.lora_r,
@@ -923,6 +938,8 @@ def main():
         tokenizer=tokenizer,
         vae=vae,
         noise_scheduler=noise_scheduler,
+        schedule_sampler=schedule_sampler,
+        min_snr_gamma=args.min_snr_gamma,
         dtype=weight_dtype,
         seed=args.seed,
         compile_unet=args.compile_unet,