From f4445cd78c6d6f9d3a7d65267878beb16aeaa8d4 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 19 Apr 2026 09:51:46 -0600 Subject: [PATCH] Added psuedo_huber loss --- extensions_built_in/sd_trainer/SDTrainer.py | 5 ++++- toolkit/config_modules.py | 4 +++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index d1d6507..c886115 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -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": diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 57edae0..5ed490e 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -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: