Add a dino version of DFE

2026-03-04 08:20:37 -07:00
parent 9dee42fc09
commit b04c64e0f8
2 changed files with 167 additions and 2 deletions
--- a/extensions_built_in/sd_trainer/SDTrainer.py
+++ b/extensions_built_in/sd_trainer/SDTrainer.py
@@ -654,7 +654,7 @@ class SDTrainer(BaseSDTrainProcess):
                    dfe_loss += torch.nn.functional.mse_loss(pred_feature_list[i], target_feature_list[i], reduction="mean")
                
                additional_loss += dfe_loss * self.train_config.diffusion_feature_extractor_weight * 100.0
-            elif self.dfe.version in [3, 4, 5]:
+            elif self.dfe.version in [3, 4, 5, 6]:
                dfe_loss = self.dfe(
                    noise=noise,
                    noise_pred=noise_pred,