Update app.py

app.py

@@ -1,4 +1,5 @@
 import torch
+import gc
 import gradio as gr
 from main import setup, execute_task
 from arguments import parse_args
@@ -51,6 +52,7 @@ def clean_dir(save_dir):
 
 def start_over(gallery_state, loaded_model_setup):
     torch.cuda.empty_cache()  # Free up cached memory
+    gc.collect()
     if gallery_state is not None:
         gallery_state = None
     if loaded_model_setup is not None:
@@ -63,6 +65,7 @@ def setup_model(prompt, model, seed, num_iterations, enable_hps, hps_w, enable_i
         
     """Clear CUDA memory before starting the training."""
     torch.cuda.empty_cache()  # Free up cached memory
+    gc.collect()
     
     # Set up arguments
     args = parse_args()
@@ -108,6 +111,7 @@ def setup_model(prompt, model, seed, num_iterations, enable_hps, hps_w, enable_i
 
 def generate_image(setup_args, num_iterations):
     torch.cuda.empty_cache()  # Free up cached memory
+    gc.collect()
 
     args = setup_args[0]
     trainer = setup_args[1]
@@ -125,6 +129,7 @@ def generate_image(setup_args, num_iterations):
     
     try:
         torch.cuda.empty_cache()  # Free up cached memory
+        gc.collect()
         steps_completed = []
         result_container = {"best_image": None, "total_init_rewards": None, "total_best_rewards": None}
         error_status = {"error_occurred": False}  # Shared dictionary to track error status
@@ -175,6 +180,7 @@ def generate_image(setup_args, num_iterations):
 
             if error_status["error_occurred"]:
                 torch.cuda.empty_cache()  # Free up cached memory
+                gc.collect()
                 yield (None, "CUDA out of memory. Please reduce your batch size or image resolution.", None)
             else:
                 main_thread.join()  # Ensure thread completion
@@ -182,13 +188,16 @@ def generate_image(setup_args, num_iterations):
                 if os.path.exists(final_image_path):
                     iter_images = list_iter_images(save_dir)
                     torch.cuda.empty_cache()  # Free up cached memory
+                    gc.collect()
                     time.sleep(0.5)
                     yield (final_image_path, f"Final image saved at {final_image_path}", iter_images)
                 else:
                     torch.cuda.empty_cache()  # Free up cached memory
+                    gc.collect()
                     yield (None, "Image generation completed, but no final image was found.", None)
 
         torch.cuda.empty_cache()  # Free up cached memory
+        gc.collect()
 
     except torch.cuda.OutOfMemoryError as e:
         print(f"Global CUDA Out of Memory Error: {e}")