Update app.py

app.py

@@ -33,7 +33,7 @@ from html_templates import (
 )
 from urllib.parse import quote
 from ultralytics import YOLO
-from device_manager import DeviceManager, device_handler
+from device_manager import DeviceManager
 import asyncio
 import traceback
 
@@ -536,6 +536,8 @@ dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staff
               "Wire-Haired_Fox_Terrier"]
 
 
+device_mgr = DeviceManager()
+
 class MultiHeadAttention(nn.Module):
 
     def __init__(self, in_dim, num_heads=8):
@@ -568,7 +570,6 @@ class BaseModel(nn.Module):
     
     def __init__(self, num_classes, device='cuda' if torch.cuda.is_available() else 'cpu'):
         super().__init__()
-        self.device_mgr = DeviceManager()
         self.device = self.device_mgr.get_optimal_device()
         self.backbone = efficientnet_v2_m(weights=EfficientNet_V2_M_Weights.IMAGENET1K_V1)
         self.feature_dim = self.backbone.classifier[1].in_features
@@ -585,7 +586,6 @@ class BaseModel(nn.Module):
 
         self.to(device)
 
-    @device_handler
     def forward(self, x):
         x = x.to(self.device)
         features = self.backbone(x)
@@ -594,7 +594,6 @@ class BaseModel(nn.Module):
         return logits, attended_features
 
 # Initialize model
-device_mgr = DeviceManager()
 num_classes = len(dog_breeds)
 
 # Initialize base model
@@ -623,31 +622,40 @@ def preprocess_image(image):
 
     return transform(image).unsqueeze(0)
     
-@device_handler
 async def predict_single_dog(image):
-    image_tensor = preprocess_image(image)
-    
-    with torch.no_grad():
-        outputs = model(image_tensor)
-        logits = outputs[0] if isinstance(outputs, tuple) else outputs
-        probs = F.softmax(logits, dim=1)
-        
-        # 其餘代碼保持不變
-        top5_prob, top5_idx = torch.topk(probs, k=5)
-        breeds = [dog_breeds[idx.item()] for idx in top5_idx[0]]
-        probabilities = [prob.item() for prob in top5_prob[0]]
-        
-        sum_probs = sum(probabilities[:3])
-        relative_probs = [f"{(prob/sum_probs * 100):.2f}%" for prob in probabilities[:3]]
+    """
+    Predicts the dog breed using only the classifier.
+    """
+    try:
+        image_tensor = preprocess_image(image).to(device_mgr.get_optimal_device())
         
-        print("\nClassifier Predictions:")
-        for breed, prob in zip(breeds[:5], probabilities[:5]):
-            print(f"{breed}: {prob:.4f}")
+        with torch.no_grad():
+            outputs = model(image_tensor)  # 同步調用
+            logits = outputs[0] if isinstance(outputs, tuple) else outputs
+            probs = F.softmax(logits, dim=1)
+            
+            top5_prob, top5_idx = torch.topk(probs, k=5)
+            breeds = [dog_breeds[idx.item()] for idx in top5_idx[0]]
+            probabilities = [prob.item() for prob in top5_prob[0]]
+            
+            sum_probs = sum(probabilities[:3])
+            relative_probs = [f"{(prob/sum_probs * 100):.2f}%" for prob in probabilities[:3]]
+            
+            print("\nClassifier Predictions:")
+            for breed, prob in zip(breeds[:5], probabilities[:5]):
+                print(f"{breed}: {prob:.4f}")
+                
+            return probabilities[0], breeds[:3], relative_probs
             
-        return probabilities[0], breeds[:3], relative_probs
+    except RuntimeError as e:
+        if "out of memory" in str(e) or "CUDA" in str(e):
+            logger.warning("ZeroGPU unavailable, falling back to CPU")
+            device_mgr._current_device = torch.device('cpu')
+            model.to('cpu')
+            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 = []