Implement pseudo Huber loss for Flux and SD3

2026-04-09 06:45:09 +00:00 · 2024-11-27 18:11:51 +00:00
parent 2a61fc0784
commit 420a180d93
15 changed files with 76 additions and 61 deletions
--- a/flux_train.py
+++ b/flux_train.py
@@ -667,7 +667,7 @@ def train(args):

                # calculate loss
                loss = train_util.conditional_loss(
-                    model_pred.float(), target.float(), reduction="none", loss_type=args.loss_type, huber_c=None
+                    args, noise_pred.float(), target.float(), timesteps, "none", noise_scheduler
                )
                if weighting is not None:
                    loss = loss * weighting