Upload app.py

app.py

@@ -462,17 +462,39 @@ with gr.Blocks(css=css) as demo:
 
 
 if __name__ == "__main__":
+    # Configure GPU/CPU handling
+    import torch
+    
+    # Function to initialize CUDA safely and verify it's working
+    def is_cuda_working():
+        try:
+            if torch.cuda.is_available():
+                # Test CUDA initialization with a small operation
+                test_tensor = torch.tensor([1.0], device="cuda")
+                _ = test_tensor * 2
+                return True
+            return False
+        except Exception as e:
+            print(f"CUDA initialization test failed: {e}")
+            return False
+    
     # Check if running on Hugging Face Spaces
     if "SPACE_ID" in os.environ:
-        # Running on HF Spaces - use minimal configuration
-        demo.launch(
-            server_port=7860,     # Match the port HF Spaces expects
-            share=False,          # Explicitly disable sharing
-            debug=False,          # Disable debug mode in production
-            show_error=True       # Show detailed errors if they occur
-        )
+        cuda_working = is_cuda_working()
+        if cuda_working:
+            print("GPU is available and working. Using CUDA.")
+            # You might want to set some environment variables or configurations here
+            os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
+        else:
+            print("CUDA not working properly. Forcing CPU mode.")
+            os.environ["CUDA_VISIBLE_DEVICES"] = ""
+            torch.backends.cudnn.enabled = False
+        
+        # Launch with minimal parameters on Spaces
+        demo.launch()
     else:
         # Running locally - use custom server settings and share
+        print(f"Running locally with device: {'cuda' if torch.cuda.is_available() else 'cpu'}")
         demo.launch(
             server_name="0.0.0.0",
             server_port=7860,