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

Allow models to provide an additional loss

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-28 09:49:36 -06:00
parent 1e12b6b73f
commit 4e50535478
1 changed files with 6 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -944,6 +944,12 @@ class SDTrainer(BaseSDTrainProcess):
loss = loss + additional_loss loss = loss + additional_loss
if hasattr(self.sd, "get_additional_loss"):
additional_model_loss = self.sd.get_additional_loss(pred, target)
if additional_model_loss is not None:
loss = loss + additional_model_loss
self.additional_logs["additional_model_loss"] = additional_model_loss.item()
if self.train_config.max_loss_debug and self.train_config.max_loss is not None: if self.train_config.max_loss_debug and self.train_config.max_loss is not None:
if loss.item() > self.train_config.max_loss: if loss.item() > self.train_config.max_loss:
print_acc(f"Loss {loss.item()} is greater than max loss {self.train_config.max_loss}. Clipping to max loss.") print_acc(f"Loss {loss.item()} is greater than max loss {self.train_config.max_loss}. Clipping to max loss.")