scale v-pred loss like noise pred

train_db.py

@@ -26,8 +26,10 @@ import library.custom_train_functions as custom_train_functions
 from library.custom_train_functions import (
     apply_snr_weight,
     get_weighted_text_embeddings,
+    prepare_scheduler_for_custom_training,
     pyramid_noise_like,
     apply_noise_offset,
+    scale_v_prediction_loss_like_noise_prediction,
 )
 
 # perlin_noise,
@@ -240,6 +242,7 @@ def train(args):
     noise_scheduler = DDPMScheduler(
         beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
     )
+    prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device)
 
     if accelerator.is_main_process:
         accelerator.init_trackers("dreambooth" if args.log_tracker_name is None else args.log_tracker_name)
@@ -327,6 +330,8 @@ def train(args):
 
                 if args.min_snr_gamma:
                     loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma)
+                if args.scale_v_pred_loss_like_noise_pred:
+                    loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
 
                 loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし