Update app.py

app.py

@@ -539,36 +539,44 @@ from ultralytics import YOLO
 import asyncio
 import traceback
 
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['HF_ZERO_GPU'] = '1'  # 明確告訴系統我們要使用 ZeroGPU
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
+
 def get_device():
     print("Initializing device configuration...")
     
-    # 首先嘗試使用 CUDA，但要更謹慎地處理初始化
-    if torch.cuda.is_available():
+    # 特別針對 ZeroGPU 的檢測邏輯
+    if 'HF_ZERO_GPU' in os.environ and torch.cuda.is_available():
         try:
-            # 設置環境變量，告訴 PyTorch 在沒有 GPU 時自動回退到 CPU
-            os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
+            # 強制進行 CUDA 初始化
+            torch.cuda.init()
+            # 等待一小段時間讓系統完成初始化
+            import time
+            time.sleep(2)
             
             device = torch.device('cuda')
-            # 使用 try-except 來處理 GPU 信息獲取
-            try:
-                print(f"Using GPU: {torch.cuda.get_device_name(0)}")
-                print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
-            except Exception as e:
-                print("GPU detected but couldn't get detailed information")
-                
-            # 進行一個小的測試計算來驗證 GPU 功能
+            
+            # 執行一個小的測試來確認 GPU 功能
             test_tensor = torch.rand(1).to(device)
             _ = test_tensor * test_tensor
-            print("GPU test calculation successful")
+            
+            print("ZeroGPU initialization successful")
+            print(f"Using device: {device}")
+            if torch.cuda.is_available():
+                print(f"GPU: {torch.cuda.get_device_name(0)}")
             
             return device
             
         except Exception as e:
-            print(f"GPU initialization failed: {str(e)}")
+            print(f"ZeroGPU initialization failed: {str(e)}")
             print("Falling back to CPU")
             return torch.device('cpu')
     else:
-        print("CUDA not available, using CPU")
+        if not torch.cuda.is_available():
+            print("CUDA not available, using CPU")
+        elif 'HF_ZERO_GPU' not in os.environ:
+            print("HF_ZERO_GPU not set, using CPU")
         return torch.device('cpu')
 
 device = get_device()