Code simplifications, avoid autocast

1 files changed, 7 insertions, 1 deletions

diff --git a/train_ti.py b/train_ti.py
index 5f37d54..a228795 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -361,6 +361,7 @@ def parse_args():
 class Checkpointer(CheckpointerBase):
     def __init__(
         self,
+        weight_dtype,
         datamodule,
         accelerator,
         vae,
@@ -387,6 +388,7 @@ class Checkpointer(CheckpointerBase):
             sample_batch_size=sample_batch_size
         )
 
+        self.weight_dtype = weight_dtype
         self.accelerator = accelerator
         self.vae = vae
         self.unet = unet
@@ -417,8 +419,9 @@ class Checkpointer(CheckpointerBase):
     @torch.no_grad()
     def save_samples(self, step, num_inference_steps, guidance_scale=7.5, eta=0.0):
         text_encoder = self.accelerator.unwrap_model(self.text_encoder)
+        orig_dtype = text_encoder.dtype
+        text_encoder.to(dtype=self.weight_dtype)
 
-        # Save a sample image
         pipeline = VlpnStableDiffusion(
             text_encoder=text_encoder,
             vae=self.vae,
@@ -430,6 +433,8 @@ class Checkpointer(CheckpointerBase):
 
         super().save_samples(pipeline, step, num_inference_steps, guidance_scale, eta)
 
+        text_encoder.to(dtype=orig_dtype)
+
         del text_encoder
         del pipeline
 
@@ -739,6 +744,7 @@ def main():
     max_acc_val = 0.0
 
     checkpointer = Checkpointer(
+        weight_dtype=weight_dtype,
         datamodule=datamodule,
         accelerator=accelerator,
         vae=vae,