Add DFE8 with partial step

2026-04-17 17:40:16 -06:00
parent 7c4f18ce51
commit beb40ae29b
2 changed files with 47 additions and 12 deletions
--- a/extensions_built_in/sd_trainer/SDTrainer.py
+++ b/extensions_built_in/sd_trainer/SDTrainer.py
@@ -664,7 +664,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, 6, 7]:
+            elif self.dfe.version in [3, 4, 5, 6, 7, 8]:
                dfe_loss = self.dfe(
                    noise=noise,
                    noise_pred=noise_pred,