try handle OOM errors

app.py

@@ -108,14 +108,25 @@ def generate_image(setup_args, num_iterations):
 
         # Function to run main in a separate thread
         def run_main():
-            result_container["best_image"], result_container["total_init_rewards"], result_container["total_best_rewards"] = execute_task(args, trainer, device, dtype, shape, enable_grad, settings, progress_callback)
+            try:
+                # Call main and handle any potential OOM errors
+                result_container["best_image"], result_container["total_init_rewards"], result_container["total_best_rewards"] = execute_task(args, trainer, device, dtype, shape, enable_grad, settings, progress_callback)
+            except torch.cuda.OutOfMemoryError as e:
+                # Handle CUDA OOM error
+                print("CUDA Out of Memory Error: ", e)
+                status["error_occurred"] = True  # Update status on error
+            except RuntimeError as e:
+                if 'out of memory' in str(e):
+                    status["error_occurred"] = True  # Update status on error
+                else:
+                    raise  # Reraise if it's not a CUDA OOM error
         
         # Start main in a separate thread
         main_thread = threading.Thread(target=run_main)
         main_thread.start()
 
         last_step_yielded = 0
-        while main_thread.is_alive() or last_step_yielded < num_iterations:
+        while main_thread.is_alive() and not status["error_occurred"]:
             # Check if new steps have been completed
             if steps_completed and steps_completed[-1] > last_step_yielded:
                 last_step_yielded = steps_completed[-1]
@@ -130,21 +141,35 @@ def generate_image(setup_args, num_iterations):
                 # Small sleep to prevent busy waiting
                 time.sleep(0.1)
 
-        main_thread.join()
-
-        # After main is complete, yield the final image
-        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)
+        # If an error occurred, clean up resources and stop
+        if status["error_occurred"]:
             torch.cuda.empty_cache()  # Free up cached memory
-            yield (final_image_path, f"Final image saved at {final_image_path}", iter_images)
+            yield (None, "CUDA out of memory. Please reduce your batch size or image resolution.", None)
         else:
-            torch.cuda.empty_cache()  # Free up cached memory
-            yield (None, "Image generation completed, but no final image was found.", None)
+            main_thread.join()
+
+            # After main is complete, yield the final image
+            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
-        return (None, f"An error occurred: {str(e)}", None)
+
+    except torch.cuda.OutOfMemoryError as e:
+        # Handle CUDA OOM error globally
+        yield (None, "CUDA out of memory.", None)
+    except RuntimeError as e:
+        if 'out of memory' in str(e):
+            yield (None, "CUDA out of memory.", None)
+        else:
+            yield (None, f"An error occurred: {str(e)}", None)
+    except Exception as e:
+        yield (None, f"An unexpected error occurred: {str(e)}", None)
 
 def show_gallery_output(gallery_state):
     if gallery_state is not None: