2 files changed, 9 insertions, 9 deletions

diff --git a/pipelines/stable_diffusion/vlpn_stable_diffusion.py b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
index dbd262f..ea2a656 100644
--- a/pipelines/stable_diffusion/vlpn_stable_diffusion.py
+++ b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
@@ -375,8 +375,7 @@ class VlpnStableDiffusion(DiffusionPipeline):
 
     def decode_latents(self, latents):
         latents = 1 / self.vae.config.scaling_factor * latents
-        # image = self.vae.decode(latents.to(dtype=self.vae.dtype)).sample
-        image = self.vae.decode(latents).sample
+        image = self.vae.decode(latents.to(dtype=self.vae.dtype)).sample
         image = (image / 2 + 0.5).clamp(0, 1)
         # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
         image = image.cpu().permute(0, 2, 3, 1).float().numpy()
diff --git a/training/strategy/lora.py b/training/strategy/lora.py
index 1e32114..8905171 100644
--- a/training/strategy/lora.py
+++ b/training/strategy/lora.py
@@ -12,6 +12,7 @@ from accelerate import Accelerator
 from transformers import CLIPTextModel
 from diffusers import AutoencoderKL, UNet2DConditionModel, DPMSolverMultistepScheduler
 from peft import get_peft_model_state_dict
+from safetensors.torch import save_file
 
 from models.clip.tokenizer import MultiCLIPTokenizer
 from training.functional import TrainingStrategy, TrainingCallbacks, save_samples
@@ -78,6 +79,9 @@ def lora_strategy_callbacks(
 
     @torch.no_grad()
     def on_checkpoint(step, postfix):
+        if postfix != "end":
+            return
+
         print(f"Saving checkpoint for step {step}...")
 
         unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=False)
@@ -94,26 +98,23 @@ def lora_strategy_callbacks(
         state_dict.update(text_encoder_state_dict)
         lora_config["text_encoder_peft_config"] = text_encoder_.get_peft_config_as_dict(inference=True)
 
-        accelerator.print(state_dict)
-        accelerator.save(state_dict, checkpoint_output_dir / f"{step}_{postfix}.pt")
+        save_file(state_dict, checkpoint_output_dir / f"{step}_{postfix}.safetensors")
         with open(checkpoint_output_dir / "lora_config.json", "w") as f:
             json.dump(lora_config, f)
 
         del unet_
         del text_encoder_
 
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
     @torch.no_grad()
     def on_sample(step):
-        vae_dtype = vae.dtype
-        vae.to(dtype=text_encoder.dtype)
-
         unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=True)
         text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=True)
 
         save_samples_(step=step, unet=unet_, text_encoder=text_encoder_)
 
-        vae.to(dtype=vae_dtype)
-
         del unet_
         del text_encoder_