Update app.py

app.py

@@ -8,7 +8,8 @@ import numpy as np
 
 app = FastAPI()
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-model = YOLO('nailong_yolo11.onnx').to(device)
+# 移除 .to(device)，在预测时指定设备
+model = YOLO('nailong_yolo11.onnx')
 
 @spaces.GPU
 def predict(img):
@@ -25,8 +26,8 @@ def predict(img):
     # 转换为numpy数组并进行预测
     img_array = np.array(input_image)
     
-    # 进行预测
-    results = model.predict(img_array)
+    # 在predict时指定device
+    results = model.predict(img_array, device=device)
     result = results[0]
     
     # 获取预测结果
@@ -36,7 +37,7 @@ def predict(img):
     info = {
         "detected": len(result.boxes) > 0,
         "count": len(result.boxes),
-        "detections": []  # 存储每个检测目标的详细信息
+        "detections": []
     }
     
     if info["detected"]:
@@ -44,7 +45,7 @@ def predict(img):
         for box in result.boxes:
             conf = float(box.conf[0])
             cls = int(box.cls[0])
-            cls_name = result.names[cls]  # 获取类别名称
+            cls_name = result.names[cls]
             
             detection_info = {
                 "class": cls_name,