1 files changed, 12 insertions, 10 deletions

diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index 2038e34..10bc6d7 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -78,6 +78,7 @@ def textual_inversion_strategy_callbacks(
             power=ema_power,
             max_value=ema_max_decay,
         )
+        ema_embeddings.to(accelerator.device)
     else:
         ema_embeddings = None
 
@@ -92,15 +93,6 @@ def textual_inversion_strategy_callbacks(
     def on_accum_model():
         return text_encoder.text_model.embeddings.temp_token_embedding
 
-    def on_prepare():
-        text_encoder.text_model.embeddings.temp_token_embedding.requires_grad_(True)
-
-        if ema_embeddings is not None:
-            ema_embeddings.to(accelerator.device)
-
-        if gradient_checkpointing:
-            unet.train()
-
     @contextmanager
     def on_train(epoch: int):
         tokenizer.train()
@@ -177,7 +169,6 @@ def textual_inversion_strategy_callbacks(
             torch.cuda.empty_cache()
 
     return TrainingCallbacks(
-        on_prepare=on_prepare,
         on_accum_model=on_accum_model,
         on_train=on_train,
         on_eval=on_eval,
@@ -197,6 +188,7 @@ def textual_inversion_prepare(
     train_dataloader: DataLoader,
     val_dataloader: Optional[DataLoader],
     lr_scheduler: torch.optim.lr_scheduler._LRScheduler,
+    gradient_checkpointing: bool = False,
     **kwargs
 ):
     weight_dtype = torch.float32
@@ -207,7 +199,17 @@ def textual_inversion_prepare(
 
     text_encoder, optimizer, train_dataloader, val_dataloader, lr_scheduler = accelerator.prepare(
         text_encoder, optimizer, train_dataloader, val_dataloader, lr_scheduler)
+
     unet.to(accelerator.device, dtype=weight_dtype)
+    unet.requires_grad_(False)
+    unet.eval()
+    if gradient_checkpointing:
+        unet.train()
+
+    text_encoder.text_model.encoder.requires_grad_(False)
+    text_encoder.text_model.final_layer_norm.requires_grad_(False)
+    text_encoder.eval()
+
     return text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler, {}