Refactor GPU initialization and device management for Hugging Face Spaces compatibility

main.py

@@ -1,3 +1,4 @@
+import spaces
 from uuid import UUID
 from fastapi import FastAPI, HTTPException
 from transformers import pipeline
@@ -12,25 +13,35 @@ from VectorDB import *
 from pydantic import BaseModel
 import threading
 import uvicorn
-import spaces
-
-zero = torch.Tensor([0]).cuda()
-print(zero.device) # <-- 'cpu' 🤔
 
+# Use spaces.GPU decorator for GPU operations
 @spaces.GPU
-def greet(n):
-    print(zero.device) # <-- 'cuda:0' 🤗
-    return f"Hello {zero + n} Tensor"
+def initialize_gpu():
+    zero = torch.Tensor([0]).cuda()
+    print(f"GPU initialized: {zero.device}")
+    return zero.device
 
+# Call this function after app startup
+device_info = None
 
 # Pick the best available device - MPS (Mac), CUDA (NVIDIA), or CPU
-if torch.backends.mps.is_available():
-    device = torch.device("mps")
-elif torch.cuda.is_available():
-    device = torch.device("cuda")
-else:
+try:
+    # For Hugging Face Spaces, let spaces handle GPU initialization
+    if "spaces" in sys.modules:
+        print("Running in Hugging Face Spaces, using spaces.GPU for device management")
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        # Local development path
+        if torch.backends.mps.is_available():
+            device = torch.device("mps")
+        elif torch.cuda.is_available():
+            device = torch.device("cuda")
+        else:
+            device = torch.device("cpu")
+    print(f"Using device: {device}")
+except Exception as e:
+    print(f"Error detecting device, falling back to CPU: {str(e)}")
     device = torch.device("cpu")
-#print(device)
 
 initRAG(device)
 supabase: Client = initSupabase()
@@ -341,17 +352,6 @@ async def status():
 
 
 
-# ...existing code...
-
-# Remove this section
-# def run_fastapi():
-#     print("Starting FastAPI server on http://0.0.0.0:8000")
-#     uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info", reload=False)
-# 
-# # Start FastAPI in a separate thread
-# fastapi_thread = threading.Thread(target=run_fastapi, daemon=True)
-# fastapi_thread.start()
-
 # Gradio Interface
 def chatbot_interface(user_input, history):
     return f"You said: {user_input}"  # Replace with actual chatbot logic