1 files changed, 29 insertions, 9 deletions

diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index 14bdafd..d306f18 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -59,14 +59,11 @@ def textual_inversion_strategy_callbacks(
     save_samples_ = partial(
         save_samples,
         accelerator=accelerator,
-        unet=unet,
-        text_encoder=text_encoder,
         tokenizer=tokenizer,
         vae=vae,
         sample_scheduler=sample_scheduler,
         train_dataloader=train_dataloader,
         val_dataloader=val_dataloader,
-        dtype=weight_dtype,
         output_dir=sample_output_dir,
         seed=seed,
         batch_size=sample_batch_size,
@@ -94,7 +91,7 @@ def textual_inversion_strategy_callbacks(
         else:
             return nullcontext()
 
-    def on_model():
+    def on_accum_model():
         return text_encoder.text_model.embeddings.temp_token_embedding
 
     def on_prepare():
@@ -149,11 +146,29 @@ def textual_inversion_strategy_callbacks(
     @torch.no_grad()
     def on_sample(step):
         with ema_context():
-            save_samples_(step=step)
+            unet_ = accelerator.unwrap_model(unet)
+            text_encoder_ = accelerator.unwrap_model(text_encoder)
+
+            orig_unet_dtype = unet_.dtype
+            orig_text_encoder_dtype = text_encoder_.dtype
+
+            unet_.to(dtype=weight_dtype)
+            text_encoder_.to(dtype=weight_dtype)
+
+            save_samples_(step=step, unet=unet_, text_encoder=text_encoder_)
+
+            unet_.to(dtype=orig_unet_dtype)
+            text_encoder_.to(dtype=orig_text_encoder_dtype)
+
+            del unet_
+            del text_encoder_
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
     return TrainingCallbacks(
         on_prepare=on_prepare,
-        on_model=on_model,
+        on_accum_model=on_accum_model,
         on_train=on_train,
         on_eval=on_eval,
         on_before_optimize=on_before_optimize,
@@ -168,7 +183,11 @@ def textual_inversion_prepare(
     accelerator: Accelerator,
     text_encoder: CLIPTextModel,
     unet: UNet2DConditionModel,
-    *args
+    optimizer: torch.optim.Optimizer,
+    train_dataloader: DataLoader,
+    val_dataloader: Optional[DataLoader],
+    lr_scheduler: torch.optim.lr_scheduler._LRScheduler,
+    **kwargs
 ):
     weight_dtype = torch.float32
     if accelerator.state.mixed_precision == "fp16":
@@ -176,9 +195,10 @@ def textual_inversion_prepare(
     elif accelerator.state.mixed_precision == "bf16":
         weight_dtype = torch.bfloat16
 
-    prepped = accelerator.prepare(text_encoder, *args)
+    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)
-    return (prepped[0], unet) + prepped[1:]
+    return text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler, {}
 
 
 textual_inversion_strategy = TrainingStrategy(