1 files changed, 13 insertions, 6 deletions

diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index 00f3529..597abd0 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -15,10 +15,10 @@ from slugify import slugify
 
 from models.clip.tokenizer import MultiCLIPTokenizer
 from training.util import EMAModel
-from training.functional import TrainingCallbacks, save_samples
+from training.functional import TrainingStrategy, TrainingCallbacks, save_samples
 
 
-def textual_inversion_strategy(
+def textual_inversion_strategy_callbacks(
     accelerator: Accelerator,
     unet: UNet2DConditionModel,
     text_encoder: CLIPTextModel,
@@ -119,17 +119,18 @@ def textual_inversion_strategy(
         with ema_context():
             yield
 
-    @torch.no_grad()
     def on_after_optimize(lr: float):
+        if use_ema:
+            ema_embeddings.step(text_encoder.text_model.embeddings.temp_token_embedding.parameters())
+
+    @torch.no_grad()
+    def on_after_epoch(lr: float):
         if use_emb_decay:
             text_encoder.text_model.embeddings.normalize(
                 emb_decay_target,
                 min(1.0, max(0.0, emb_decay_factor * ((lr - emb_decay_start) / (learning_rate - emb_decay_start))))
             )
 
-        if use_ema:
-            ema_embeddings.step(text_encoder.text_model.embeddings.temp_token_embedding.parameters())
-
     def on_log():
         if use_ema:
             return {"ema_decay": ema_embeddings.decay}
@@ -157,7 +158,13 @@ def textual_inversion_strategy(
         on_train=on_train,
         on_eval=on_eval,
         on_after_optimize=on_after_optimize,
+        on_after_epoch=on_after_epoch,
         on_log=on_log,
         on_checkpoint=on_checkpoint,
         on_sample=on_sample,
     )
+
+
+textual_inversion_strategy = TrainingStrategy(
+    callbacks=textual_inversion_strategy_callbacks,
+)