Update

1 files changed, 17 insertions, 0 deletions

diff --git a/train_ti.py b/train_ti.py
index ae73639..26f7941 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -23,6 +23,7 @@ from data.csv import VlpnDataModule, keyword_filter
 from training.functional import train, add_placeholder_tokens, get_models
 from training.strategy.ti import textual_inversion_strategy
 from training.optimization import get_scheduler
+from training.sampler import create_named_schedule_sampler
 from training.util import AverageMeter, save_args
 
 logger = get_logger(__name__)
@@ -359,6 +360,19 @@ def parse_args():
         default=0.9999
     )
     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",
@@ -682,6 +696,7 @@ def main():
         args.emb_alpha,
         args.emb_dropout
     )
+    schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, noise_scheduler.config.num_train_timesteps)
 
     tokenizer.set_use_vector_shuffle(args.vector_shuffle)
     tokenizer.set_dropout(args.vector_dropout)
@@ -837,6 +852,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,