Use L1 pixel loss

RossM · May 20, 2024 · abbc723 · abbc723
1 parent 72213a8
commit abbc723
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dreambooth/train_dreambooth.py b/dreambooth/train_dreambooth.py
@@ -1829,10 +1829,10 @@ def loss_fn(model_pred, target):
                                 pixel_target = vae.decode(latents[pixel_mask].to(weight_dtype)).sample
                                 pixel_pred = vae.decode(pred_latents[pixel_mask].to(weight_dtype)).sample
 
-                                pixel_loss[pixel_mask] = F.mse_loss(pixel_pred.float(), pixel_target.float(), reduction="none").mean(dim=(1,2,3))
+                                pixel_loss[pixel_mask] = F.l1_loss(pixel_pred.float(), pixel_target.float(), reduction="none").mean(dim=(1,2,3))
 
                                 # Ad hoc weight factor to make losses have approximately the same scale
-                                pixel_loss = 700 * pixel_loss
+                                pixel_loss = 25 * pixel_loss
 
                             else:
                                 pixel_loss = latent_loss.new_zeros([])