Improve memory management and error handling in image processing

app.py

@@ -1,6 +1,7 @@
 import shutil
 import time
 from pathlib import Path
+from xml.parsers.expat import model
 
 import cv2
 import gradio as gr
@@ -246,8 +247,6 @@ with gr.Blocks() as demo:
     with gr.Row():
         run_button = gr.Button("Run")
 
-    model = Model()
-
     def preprocess_image(
         image: PIL.Image.Image,
         image_size: int = IMAGE_SIZE,
@@ -276,6 +275,9 @@ with gr.Blocks() as demo:
         num_blocks_on_gpu,
         seed,
     ):
+        if DEVICE == "cuda":
+            torch.cuda.empty_cache()
+
         pil_images = []
         for contents in images:
             for content in contents:
@@ -285,15 +287,21 @@ with gr.Blocks() as demo:
                     break
 
         try:
-            output_image = model.compute(
-                pil_images,
-                prompt,
-                negative_prompt,
-                true_cfg_scale,
-                num_inference_steps,
-                num_blocks_on_gpu,
-                seed,
-            )
+            model = Model()
+
+            try:
+                output_image = model.compute(
+                    pil_images,
+                    prompt,
+                    negative_prompt,
+                    true_cfg_scale,
+                    num_inference_steps,
+                    num_blocks_on_gpu,
+                    seed,
+                )
+            except Exception:
+                if DEVICE == "cuda":
+                    torch.cuda.empty_cache()
 
             # Save the output image for download
             timestamp = int(time.time())