Update app.py

app.py

@@ -35,6 +35,7 @@ from urllib.parse import quote
 from ultralytics import YOLO
 import traceback
 import spaces
+from functools import wraps
 
 # model_yolo = YOLO('yolov8l.pt')
 
@@ -533,6 +534,12 @@ dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staff
               "Wire-Haired_Fox_Terrier"]
 
 
+def gpu_wrapper(f):
+    @wraps(f)
+    async def wrapped(*args, **kwargs):
+        return await spaces.GPU(f)(*args, **kwargs)
+    return wrapped
+
 class MultiHeadAttention(nn.Module):
 
     def __init__(self, in_dim, num_heads=8):
@@ -612,7 +619,7 @@ def preprocess_image(image):
 
     return transform(image).unsqueeze(0)
     
-@spaces.GPU 
+@gpu_wrapper
 async def predict_single_dog(image):
     """
     Predicts the dog breed using only the classifier.
@@ -645,7 +652,7 @@ async def predict_single_dog(image):
             
         return probabilities[0], breeds[:3], relative_probs
         
-@spaces.GPU
+@gpu_wrapper
 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]
@@ -728,7 +735,7 @@ def create_breed_comparison(breed1: str, breed2: str) -> dict:
 
     return comparison_data
     
-@spaces.GPU
+@gpu_wrapper
 async def predict(image):
     """
     Main prediction function that handles both single and multiple dog detection.