Some work on sd3 training. Not working

extensions_built_in/sd_trainer/SDTrainer.py

@@ -29,6 +29,7 @@ import gc
 import torch
 from jobs.process import BaseSDTrainProcess
 from torchvision import transforms
+import math
 
 
 
@@ -366,7 +367,29 @@ class SDTrainer(BaseSDTrainProcess):
             loss_per_element = (weighing.float() * (denoised_latents.float() - target.float()) ** 2)
             loss = loss_per_element
         else:
-            if self.train_config.loss_type == "mae":
+            # handle flow matching ref https://github.com/huggingface/diffusers/blob/ec068f9b5bf7c65f93125ec889e0ff1792a00da1/examples/dreambooth/train_dreambooth_lora_sd3.py#L1485C17-L1495C100
+            if self.sd.is_v3:
+                target = noisy_latents.detach()
+                bsz = pred.shape[0]
+                # todo implement others
+                # weighing_scheme =
+                # 3 just do mode for now?
+                # if args.weighting_scheme == "sigma_sqrt":
+                sigmas = self.sd.noise_scheduler.get_sigmas(timesteps, pred.ndim, dtype, self.device_torch)
+                weighting = (sigmas ** -2.0).float()
+                # elif args.weighting_scheme == "logit_normal":
+                #     # See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$).
+                #     u = torch.normal(mean=args.logit_mean, std=args.logit_std, size=(bsz,), device=accelerator.device)
+                #     weighting = torch.nn.functional.sigmoid(u)
+                # elif args.weighting_scheme == "mode":
+                # mode_scale = 1.29
+                # See sec 3.1 in the SD3 paper (20).
+                # u = torch.rand(size=(bsz,), device=pred.device)
+                # weighting = 1 - u - mode_scale * (torch.cos(math.pi * u / 2) ** 2 - 1 + u)
+
+                loss = (weighting.float() * (pred.float() - target.float()) ** 2).reshape(target.shape[0], -1)
+
+            elif self.train_config.loss_type == "mae":
                 loss = torch.nn.functional.l1_loss(pred.float(), target.float(), reduction="none")
             else:
                 loss = torch.nn.functional.mse_loss(pred.float(), target.float(), reduction="none")