Update app.py

app.py

@@ -534,12 +534,6 @@ 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):
@@ -619,7 +613,7 @@ def preprocess_image(image):
 
     return transform(image).unsqueeze(0)
     
-@gpu_wrapper
+@spaces.GPU
 async def predict_single_dog(image):
     """
     Predicts the dog breed using only the classifier.
@@ -652,7 +646,7 @@ async def predict_single_dog(image):
             
         return probabilities[0], breeds[:3], relative_probs
         
-@gpu_wrapper
+@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]
@@ -735,7 +729,7 @@ def create_breed_comparison(breed1: str, breed2: str) -> dict:
 
     return comparison_data
     
-@gpu_wrapper
+@spaces.GPU
 async def predict(image):
     """
     Main prediction function that handles both single and multiple dog detection.
@@ -888,6 +882,7 @@ def show_details_html(choice, previous_output, initial_state):
 
 def main():
     with gr.Blocks(css=get_css_styles()) as iface:
+        spaces.init()
 
         gr.HTML("""
         <header style='text-align: center; padding: 20px; margin-bottom: 20px;'>