reduce memory usage to fit in 12gb colab

ratwithacompiler · Oct 7, 2022 · c9f9c45 · c9f9c45
1 parent 71ec208
commit c9f9c45
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/convert_original_stable_diffusion_to_diffusers.py b/convert_original_stable_diffusion_to_diffusers.py
@@ -675,6 +675,7 @@ def convert_ldm_clip_checkpoint(checkpoint):
 
     unet = UNet2DConditionModel(**unet_config)
     unet.load_state_dict(converted_unet_checkpoint)
+    del converted_unet_checkpoint
 
     print("unet loaded dict")
 
@@ -686,6 +687,7 @@ def convert_ldm_clip_checkpoint(checkpoint):
 
     vae = AutoencoderKL(**vae_config)
     vae.load_state_dict(converted_vae_checkpoint)
+    del converted_vae_checkpoint
 
     print("vae loaded")
 
@@ -694,6 +696,7 @@ def convert_ldm_clip_checkpoint(checkpoint):
     if text_model_type == "FrozenCLIPEmbedder":
         print("converting clip")
         text_model = convert_ldm_clip_checkpoint(checkpoint)
+        del checkpoint
         print("converting clip done")
         tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
         print("safety checker")