Update app.py

app.py

@@ -591,13 +591,14 @@ class BaseModel(nn.Module):
 
 class ModelManager:
     """
-    模型管理器：負責AI模型的初始化和管理
-    使用單例模式確保只有一個實例在管理所有模型
+    模型管理器：負責AI模型的初始化、設備管理和資源控制
+    使用單例模式確保整個應用程序中只有一個實例
     """
     _instance = None
     _initialized = False
     _yolo_model = None
     _breed_model = None
+    _device = None
 
     def __new__(cls):
         if cls._instance is None:
@@ -607,8 +608,20 @@ class ModelManager:
     def __init__(self):
         # 避免重複初始化
         if not ModelManager._initialized:
+            # 初始化設備，這會在第一次創建實例時執行
+            self._device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
             ModelManager._initialized = True
     
+    @property
+    def device(self):
+        """
+        提供對設備的訪問
+        確保在需要時設備已經被初始化
+        """
+        if self._device is None:
+            self._device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        return self._device
+
     @property
     def yolo_model(self):
         """
@@ -623,18 +636,23 @@ class ModelManager:
     def breed_model(self):
         """
         延遲初始化品種分類模型
-        只有在第一次使用時才會創建實例
+        只有在第一次使用時才會創建實例並移動到正確的設備上
         """
         if self._breed_model is None:
-            self._breed_model = BaseModel(num_classes=len(dog_breeds), 
-                                        device=device).to(device)
-            checkpoint = torch.load('124_best_model_dog.pth', 
-                                  map_location=device)
-            self._breed_model.load_state_dict(checkpoint['base_model'], 
-                                            strict=False)
+            self._breed_model = BaseModel(
+                num_classes=len(dog_breeds), 
+                device=self.device  # 使用我們的device屬性
+            ).to(self.device)
+            
+            checkpoint = torch.load(
+                '124_best_model_dog.pth', 
+                map_location=self.device  # 確保checkpoint加載到正確的設備
+            )
+            self._breed_model.load_state_dict(checkpoint['base_model'], strict=False)
             self._breed_model.eval()
         return self._breed_model
 
+
 model_manager = ModelManager()
             
 
@@ -663,7 +681,7 @@ def predict_single_dog(image):
         tuple: (top1_prob, topk_breeds, relative_probs)
     """
     
-    image_tensor = preprocess_image(image).to(device)
+    image_tensor = preprocess_image(image).to(model_manager.device)
     
     with torch.no_grad():
         # Get model outputs (只使用logits，不需要features)