space gpu cuda fix

app.py

@@ -17,7 +17,7 @@ import gc
 from contextlib import contextmanager
 import os
 from loadimg import load_img
-# Add this near the top of the file
+import traceback
 
 title = "# **WIP / DEMO** 🙋🏻‍♂️Welcome to Tonic's Pixtral Model Demo"
 description = """
@@ -206,8 +206,31 @@ def gpu_memory_manager():
         torch.cuda.empty_cache()
         gc.collect()
 
+def cuda_error_handler(func):
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except RuntimeError as e:
+            if "CUDA" in str(e):
+                print(f"CUDA error occurred: {str(e)}")
+                print("Attempting to recover...")
+                torch.cuda.empty_cache()
+                gc.collect()
+                try:
+                    return func(*args, **kwargs)
+                except Exception as e2:
+                    print(f"Recovery failed. Error: {str(e2)}")
+                    return f"An error occurred: {str(e2)}", 0, 0
+            else:
+                raise
+        except Exception as e:
+            print(f"An unexpected error occurred: {str(e)}")
+            traceback.print_exc()
+            return f"An unexpected error occurred: {str(e)}", 0, 0
+    return wrapper
 
 @spaces.GPU(duration=120)
+@cuda_error_handler
 def generate_text(image, prompt, max_tokens):
     try:
         with gpu_memory_manager():
@@ -243,9 +266,12 @@ def generate_text(image, prompt, max_tokens):
         
         return generated_text, len(generated_ids[0]), 1
     except Exception as e:
+        print(f"Error in generate_text: {str(e)}")
+        traceback.print_exc()
         return f"Error: {str(e)}", 0, 0
 
 @spaces.GPU(duration=60)
+@cuda_error_handler
 def calculate_similarity(image1, image2):
     try:
         with gpu_memory_manager():
@@ -269,6 +295,8 @@ def calculate_similarity(image1, image2):
         
         return similarity
     except Exception as e:
+        print(f"Error in calculate_similarity: {str(e)}")
+        traceback.print_exc()
         return f"Error: {str(e)}"
 
 with gr.Blocks() as demo:
@@ -324,4 +352,8 @@ with gr.Blocks() as demo:
             )
 
 if __name__ == "__main__":
-    demo.launch()
+    try:
+        demo.launch()
+    except Exception as e:
+        print(f"An error occurred while launching the demo: {str(e)}")
+        traceback.print_exc()