Update app.py

app.py

@@ -33,9 +33,10 @@ from html_templates import (
 )
 from urllib.parse import quote
 from ultralytics import YOLO
-from device_manager import DeviceManager
+from device_manager import DeviceManager, device_handler
 import asyncio
 import traceback
+import spaces
 
 
 # model_yolo = YOLO('yolov8l.pt')
@@ -622,6 +623,7 @@ def preprocess_image(image):
 
     return transform(image).unsqueeze(0)
     
+@device_handler    
 async def predict_single_dog(image):
     """
     Predicts the dog breed using only the classifier.
@@ -655,7 +657,8 @@ async def predict_single_dog(image):
             return await predict_single_dog(image)  # 遞迴調用，使用 CPU
         raise e
         
-
+        
+@device_handler
 async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []