Update app.py

app.py

@@ -93,6 +93,7 @@ def generate_image(setup_args, num_iterations):
     clean_dir(save_dir)
     
     try:
+        torch.cuda.empty_cache()  # Free up cached memory
         steps_completed = []
         result_container = {"best_image": None, "total_init_rewards": None, "total_best_rewards": None}
 
@@ -130,11 +131,14 @@ def generate_image(setup_args, num_iterations):
         final_image_path = os.path.join(save_dir, "best_image.png")
         if os.path.exists(final_image_path):
             iter_images = list_iter_images(save_dir)
+            torch.cuda.empty_cache()  # Free up cached memory
             yield (final_image_path, f"Final image saved at {final_image_path}", iter_images)
         else:
+            torch.cuda.empty_cache()  # Free up cached memory
             yield (None, "Image generation completed, but no final image was found.", None)
 
     except Exception as e:
+        torch.cuda.empty_cache()  # Free up cached memory
         yield (None, f"An error occurred: {str(e)}", None)
 
 def show_gallery_output(gallery_state):