github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Added psuedo_huber loss

Browse Source
This commit is contained in:
Jaret Burkett
2026-04-19 09:51:46 -06:00
parent 488878f354
commit f4445cd78c
2 changed files with 7 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -781,7 +781,10 @@ class SDTrainer(BaseSDTrainProcess):
t0 = noisy_latents - tv * noise_pred
target = batch.latents.detach()
pred = t0
if self.train_config.loss_type == "pseudo_huber":
diff = pred - target
c=0.0
loss =(torch.sqrt(diff.pow(2) + c ** 2) - c)
if self.train_config.loss_type == "mae":
loss = torch.nn.functional.l1_loss(pred.float(), target.float(), reduction="none")
elif self.train_config.loss_type == "wavelet":

4
toolkit/config_modules.py
View File

@@ -499,7 +499,7 @@ class TrainConfig:
self.correct_pred_norm = kwargs.get('correct_pred_norm', False)
self.correct_pred_norm_multiplier = kwargs.get('correct_pred_norm_multiplier', 1.0)
self.loss_type = kwargs.get('loss_type', 'mse') # mse, mae, wavelet, pixelspace, mean_flow
self.loss_type = kwargs.get('loss_type', 'mse') # mse, mae, wavelet, pixelspace, mean_flow, pseudo_huber
# do the loss on a timestep to 0 prediction
self.t0_loss_target = kwargs.get('t0_loss_target', False)
@@ -707,6 +707,8 @@ class ModelConfig:
# model paths for models that support it
self.model_paths = kwargs.get("model_paths", {})
self.in_context = kwargs.get("in_context", False)
# allow frontend to pass arch with a color like arch:tag
# but remove the tag
if self.arch is not None: