Update app.py

app.py

@@ -41,6 +41,31 @@ from ultralytics import YOLO
 import asyncio
 import traceback
 
+def setup_environment():
+    """配置適合 ZeroGPU 環境的設置"""
+    # 啟用 CUDA 錯誤的同步報告，幫助診斷問題
+    os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
+    
+    # 檢查 CUDA 是否可用
+    if torch.cuda.is_available():
+        # 顯示 GPU 信息
+        device_name = torch.cuda.get_device_name(0)
+        print(f"使用 GPU: {device_name}")
+        
+        # 針對 A100 的優化設置
+        if "A100" in device_name:
+            print("檢測到 A100 GPU，應用特殊優化...")
+            torch.backends.cudnn.benchmark = True
+            
+        # 清理 GPU 內存
+        torch.cuda.empty_cache()
+        return True
+    else:
+        print("CUDA 不可用，使用 CPU 模式")
+        return False
+
+cuda_available = setup_environment()
+
 history_manager = UserHistoryManager()
 
 class ModelManager:
@@ -70,28 +95,42 @@ class ModelManager:
             self._device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         return self._device
 
-    @property
-    def yolo_model(self):
-        if self._yolo_model is None:
-            self._yolo_model = YOLO('yolov8x.pt')
-        return self._yolo_model
-
     # @property
     # def yolo_model(self):
     #     if self._yolo_model is None:
     #         self._yolo_model = YOLO('yolov8x.pt')
-            
-    #         if hasattr(self._yolo_model, 'model') and hasattr(self._yolo_model.model, 'fuse'):
-    #             print("avoid CUDA error")
-                
-    #             def no_fuse(*args, **kwargs):
-    #                 print("skip the fuse step")
-    #                 return self._yolo_model.model
-                
-    #             self._yolo_model.model.fuse = no_fuse
-                
     #     return self._yolo_model
 
+    @property
+    def yolo_model(self):
+        if self._yolo_model is None:
+            try:
+                print("正在加載 YOLO 模型...")
+                # 不指定設備，讓 YOLO 自動選擇
+                self._yolo_model = YOLO('yolov8x.pt')
+                
+                # 禁用模型融合來避免 CUDA 錯誤
+                if hasattr(self._yolo_model, 'model') and hasattr(self._yolo_model.model, 'fuse'):
+                    print("在 ZeroGPU 環境下禁用模型融合以避免 CUDA 錯誤")
+                    
+                    # 備份原始融合方法
+                    original_fuse = self._yolo_model.model.fuse
+                    
+                    # 創建一個空的融合方法
+                    def no_fuse(*args, **kwargs):
+                        print("已跳過融合操作")
+                        return self._yolo_model.model
+                    
+                    # 替換融合方法
+                    self._yolo_model.model.fuse = no_fuse
+            except Exception as e:
+                print(f"加載 YOLO 模型時出錯: {str(e)}")
+                print("嘗試降級到較小的模型和 CPU 模式...")
+                # 降級到較小的模型並明確使用 CPU
+                self._yolo_model = YOLO('yolov8n.pt', device='cpu')
+        
+        return self._yolo_model
+
     @property
     def breed_model(self):
         if self._breed_model is None: