Fix L1 loss weighting

RossM · Jun 11, 2024 · 67eaa58 · 67eaa58
1 parent 4c291b6
commit 67eaa58
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dreambooth/train_dreambooth.py b/dreambooth/train_dreambooth.py
@@ -1833,7 +1833,7 @@ def loss_fn(model_pred, target):
                             latent_loss = F.mse_loss(model_pred.float(), target.float(), reduction="none").mean(dim=(1,2,3))
 
                             if args.l1_weight != 0:
-                                l1_weight = args.l1_weight * (1 - (1 - alpha_prod) ** args.l1_gamma)
+                                l1_weight = args.l1_weight * (1 - (1 - alpha_prod) ** args.l1_gamma) * snr ** 0.5
                                 latent_loss = latent_loss + l1_weight * F.l1_loss(model_pred.float(), target.float(), reduction="none").mean(dim=(1,2,3))
 
                             loss = latent_loss