Update app.py

app.py

@@ -34,6 +34,7 @@ from html_templates import (
 from urllib.parse import quote
 from ultralytics import YOLO
 import traceback
+import spaces
 
 model_yolo = YOLO('yolov8l.pt')
 
@@ -150,7 +151,8 @@ def preprocess_image(image):
     ])
 
     return transform(image).unsqueeze(0)
-
+    
+@spaces.GPU
 async def predict_single_dog(image):
     """
     Predicts the dog breed using only the classifier.
@@ -181,8 +183,9 @@ async def predict_single_dog(image):
             print(f"{breed}: {prob:.4f}")
             
         return probabilities[0], breeds[:3], relative_probs
+        
 
-
+@spaces.GPU
 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 = []
@@ -263,8 +266,9 @@ def create_breed_comparison(breed1: str, breed2: str) -> dict:
         }
 
     return comparison_data
+    
 
-
+@spaces.GPU
 async def predict(image):
     """
     Main prediction function that handles both single and multiple dog detection.