(ACTUAL) Min-SNR Weighting Strategy: Fixed SNR calculation to authors implementation

2026-04-08 22:35:09 +00:00 · 2023-03-23 12:34:49 +00:00
parent a3c7d711e4
commit 518a18aeff
5 changed files with 17 additions and 14 deletions
--- a/train_db.py
+++ b/train_db.py
@@ -293,7 +293,8 @@ def train(args):
                loss = loss * loss_weights

                if args.min_snr_gamma:
-                  loss = apply_snr_weight(loss, latents, noisy_latents, args.min_snr_gamma)
+                  loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma)
+

                loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし