Added flat snr gamma vs min. Fixes timestep timing

toolkit/config_modules.py

@@ -168,6 +168,7 @@ class TrainConfig:
         self.train_unet = kwargs.get('train_unet', True)
         self.train_text_encoder = kwargs.get('train_text_encoder', True)
         self.min_snr_gamma = kwargs.get('min_snr_gamma', None)
+        self.snr_gamma = kwargs.get('snr_gamma', None)
         self.noise_offset = kwargs.get('noise_offset', 0.0)
         self.skip_first_sample = kwargs.get('skip_first_sample', False)
         self.gradient_checkpointing = kwargs.get('gradient_checkpointing', True)

toolkit/train_tools.py

@@ -688,13 +688,17 @@ def apply_snr_weight(
         loss,
         timesteps,
         noise_scheduler: Union['DDPMScheduler'],
-        gamma
+        gamma,
+        fixed=False,
 ):
-    # will get it form noise scheduler if exist or will calculate it if not
+    # will get it from noise scheduler if exist or will calculate it if not
     all_snr = get_all_snr(noise_scheduler, loss.device)
 
     snr = torch.stack([all_snr[t] for t in timesteps])
     gamma_over_snr = torch.div(torch.ones_like(snr) * gamma, snr)
-    snr_weight = torch.minimum(gamma_over_snr, torch.ones_like(gamma_over_snr)).float().to(loss.device)  # from paper
+    if fixed:
+        snr_weight = gamma_over_snr.float().to(loss.device)  # directly using gamma over snr
+    else:
+        snr_weight = torch.minimum(gamma_over_snr, torch.ones_like(gamma_over_snr)).float().to(loss.device)
     loss = loss * snr_weight
     return loss