Attach the ema to the base model so it can be used on specific models

2026-08-01 10:05:42 -06:00
parent c68e58083f
commit fc78b07332
1 changed files with 3 additions and 0 deletions
--- a/jobs/process/BaseSDTrainProcess.py
+++ b/jobs/process/BaseSDTrainProcess.py
@@ -794,6 +794,9 @@ class BaseSDTrainProcess(BaseTrainProcess):
                use_feedback=self.train_config.ema_config.use_feedback,
                param_multiplier=self.train_config.ema_config.param_multiplier,
            )
+            # expose to the model: models that run an EMA-teacher forward during training
+            # (e.g. wan21_pixel self_flow) read it from here
+            self.sd.ema = self.ema

    def before_dataset_load(self):
        pass