Added target_norm_std which is a game changer

2024-07-28 16:08:33 -06:00
parent 0bc4d555c7
commit e81e19fd0f
4 changed files with 11 additions and 2 deletions
--- a/extensions_built_in/sd_trainer/SDTrainer.py
+++ b/extensions_built_in/sd_trainer/SDTrainer.py
@@ -454,6 +454,13 @@ class SDTrainer(BaseSDTrainProcess):
            loss = loss + self.adapter.additional_loss.mean()
            self.adapter.additional_loss = None

+        if self.train_config.target_norm_std:
+            # seperate out the batch and channels
+            pred_std = noise_pred.std([2, 3], keepdim=True)
+            norm_std_loss = torch.abs(1.0 - pred_std).mean()
+            loss = loss + norm_std_loss
+
+
        return loss

    def preprocess_batch(self, batch: 'DataLoaderBatchDTO'):