Update app.py

app.py

@@ -587,15 +587,55 @@ class BaseModel(nn.Module):
         attended_features = self.attention(features)
         logits = self.classifier(attended_features)
         return logits, attended_features
-        
 
-model_yolo = YOLO('yolov8l.pt')
-num_classes = len(dog_breeds)
-model = BaseModel(num_classes=num_classes)
-model_path = '124_best_model_dog.pth'
-checkpoint = torch.load(model_path)
-model.load_state_dict(checkpoint['base_model'], strict=False)
-model.eval()
+
+class ModelManager:
+    """
+    模型管理器：負責AI模型的初始化和管理
+    使用單例模式確保只有一個實例在管理所有模型
+    """
+    _instance = None
+    _initialized = False
+    _yolo_model = None
+    _breed_model = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        # 避免重複初始化
+        if not ModelManager._initialized:
+            ModelManager._initialized = True
+    
+    @property
+    def yolo_model(self):
+        """
+        延遲初始化YOLO模型
+        只有在第一次使用時才會創建實例
+        """
+        if self._yolo_model is None:
+            self._yolo_model = YOLO('yolov8l.pt')
+        return self._yolo_model
+
+    @property
+    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.eval()
+        return self._breed_model
+
+model_manager = ModelManager()
             
 
 # Image preprocessing function
@@ -627,7 +667,7 @@ async def predict_single_dog(image):
     
     with torch.no_grad():
         # Get model outputs (只使用logits，不需要features)
-        logits = model(image_tensor)[0]  # 如果model仍返回tuple，取第一個元素
+        logits = model_manager.breed_model(image_tensor)[0]  # 如果model仍返回tuple，取第一個元素
         probs = F.softmax(logits, dim=1)
         
         # Classifier prediction
@@ -649,7 +689,9 @@ async def predict_single_dog(image):
 @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]
+    results = model_manager.yolo_model(image, conf=conf_threshold, 
+                                     iou=iou_threshold)[0]
+    
     dogs = []
     boxes = []
     for box in results.boxes: