Update app.py

app.py

@@ -586,31 +586,15 @@ class BaseModel(nn.Module):
         attended_features = self.attention(features)
         logits = self.classifier(attended_features)
         return logits, attended_features
+        
 
-
-class ModelManager:
-    _instance = None
-    _initialized = False
-    
-    def __init__(self):
-        self.model_yolo = None
-        self.model = None
-    
-    @classmethod
-    async def get_instance(cls):
-        if not cls._instance:
-            cls._instance = cls()
-        return cls._instance
-    
-    async def initialize(self):
-        if not self._initialized:
-            self.model_yolo = YOLO('yolov8l.pt')
-            num_classes = len(dog_breeds)
-            self.model = BaseModel(num_classes=num_classes)
-            checkpoint = torch.load('124_best_model_dog.pth')
-            self.model.load_state_dict(checkpoint['base_model'], strict=False)
-            self.model.eval()
-            self._initialized = True
+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()
             
 
 # Image preprocessing function
@@ -637,8 +621,6 @@ async def predict_single_dog(image):
     Returns:
         tuple: (top1_prob, topk_breeds, relative_probs)
     """
-    manager = await ModelManager.get_instance()
-    await manager.initialize()
     
     image_tensor = preprocess_image(image).to(device)
     
@@ -665,10 +647,8 @@ async def predict_single_dog(image):
         
 @spaces.GPU
 async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
-    manager = await ModelManager.get_instance()
-    await manager.initialize()
     
-    results = manager.model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
+    results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     boxes = []
     for box in results.boxes:
@@ -759,8 +739,6 @@ async def predict(image):
     Returns:
         tuple: (html_output, annotated_image, initial_state)
     """
-    manager = await ModelManager.get_instance()
-    await manager.initialize()
     
     if image is None:
         return format_warning_html("Please upload an image to start."), None, None
@@ -903,7 +881,6 @@ def show_details_html(choice, previous_output, initial_state):
 
 def main():
     with gr.Blocks(css=get_css_styles()) as iface:
-        spaces.GPU.setup()
 
         gr.HTML("""
         <header style='text-align: center; padding: 20px; margin-bottom: 20px;'>