Empty GPU cache at the end of run

baseline.py

@@ -4,6 +4,7 @@ import torch
 import models
 from models import pipelines
 from shared import model_dict, DEFAULT_OVERALL_NEGATIVE_PROMPT
+import gc
 
 vae, tokenizer, text_encoder, unet, scheduler, dtype = model_dict.vae, model_dict.tokenizer, model_dict.text_encoder, model_dict.unet, model_dict.scheduler, model_dict.dtype
 
@@ -36,5 +37,8 @@ def run(prompt, scheduler_key='dpm_scheduler', bg_seed=1, num_inference_steps=20
         model_dict, latents, input_embeddings, num_inference_steps,  
         guidance_scale=guidance_scale, scheduler_key=scheduler_key
     )
+    
+    gc.collect()
+    torch.cuda.empty_cache()
 
     return images[0]

generation.py

@@ -7,6 +7,7 @@ import utils
 from models import pipelines, sam
 from utils import parse, latents
 from shared import model_dict, sam_model_dict, DEFAULT_SO_NEGATIVE_PROMPT, DEFAULT_OVERALL_NEGATIVE_PROMPT
+import gc
 
 verbose = False
 
@@ -194,5 +195,8 @@ def run(
 
         # display(Image.fromarray(images[0]), "img", run_ind)
         
+    gc.collect()
+    torch.cuda.empty_cache()
+        
     return images[0], so_img_list