revert commit

app.py

@@ -229,7 +229,31 @@ def gpu_memory_manager():
         torch.cuda.empty_cache()
         gc.collect()
 
-@spaces.GPU()
+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():
@@ -264,7 +288,8 @@ def generate_text(image, prompt, max_tokens):
         traceback.print_exc()
         return f"Error: {str(e)}", 0, 0
 
-@spaces.GPU()
+@spaces.GPU(duration=60)
+@cuda_error_handler
 def calculate_similarity(image1, image2):
     try:
         with gpu_memory_manager():