Empty cache after each generation

generation.py

@@ -65,6 +65,9 @@ def generate_single_object_with_box_batch(prompts, bboxes, phrases, words, input
             return_box_vis=True, save_all_latents=True, batched_condition=True, scheduler_key=scheduler_key
         )
         
+        gc.collect()
+        torch.cuda.empty_cache()
+        
         single_object_images.append(single_object_images_batch)
         single_object_pil_images_box_ann.append(single_object_pil_images_box_ann_batch)
         latents_all.append(latents_all_batch)
@@ -81,6 +84,9 @@ def generate_single_object_with_box_batch(prompts, bboxes, phrases, words, input
     
     latents_all = latents_all.transpose(0,1)[:,:,None,...]
     
+    gc.collect()
+    torch.cuda.empty_cache()
+    
     return latents_all, mask_selected_tensor, single_object_pil_images_box_ann
 
 def get_masked_latents_all_list(so_prompt_phrase_word_box_list, input_latents_list, so_input_embeddings, verbose=False, **kwargs):