Update app.py

app.py

@@ -8,30 +8,25 @@ import os
 import hashlib
 from typing import Optional
 
-# ===== MODEL COMPONENTS =====
+# ===== IMPORT ALL NECESSARY MODULES =====
 from torch.nn import Sequential, Conv2d, BatchNorm2d, SiLU, ReLU, LeakyReLU, MaxPool2d, Upsample, ModuleList
-from ultralytics.nn.modules.conv import Conv, Concat  # Added Concat
+from ultralytics.nn.modules.conv import Conv, Concat
 from ultralytics.nn.modules.block import Bottleneck, C2f, SPPF
+from ultralytics.nn.modules.head import Detect  # Added Detect
 from ultralytics.nn.tasks import DetectionModel
 
-# ===== SAFE GLOBALS =====
-# Add all components including Concat to safe globals
+# ===== SAFE GLOBALS CONFIGURATION =====
+# Add all components to safe globals
 torch.serialization.add_safe_globals([
+    # Torch modules
+    Sequential, Conv2d, BatchNorm2d, SiLU, ReLU, LeakyReLU, 
+    MaxPool2d, Upsample, ModuleList,
+    
+    # Ultralytics modules
     DetectionModel,
-    Sequential,
-    Conv,
-    Concat,  # Explicitly added
-    Conv2d,
-    BatchNorm2d,
-    SiLU,
-    ReLU,
-    LeakyReLU,
-    Bottleneck,
-    C2f,
-    SPPF,
-    ModuleList,
-    MaxPool2d,
-    Upsample
+    Conv, Concat,
+    Bottleneck, C2f, SPPF,
+    Detect,  # Added Detect head
 ])
 
 # ===== MODEL CONFIG =====
@@ -73,7 +68,10 @@ def load_model(model_path: str) -> YOLO:
     """Safely load YOLO model with error handling"""
     print("🔧 Loading model...")
     try:
-        model = YOLO(model_path)
+        # Temporary workaround for PyTorch 2.6+ weights_only=True issues
+        # Only use if you trust the model source!
+        model = YOLO(model_path, task='detect')
+        
         # Test with small dummy input
         with torch.no_grad():
             dummy = torch.zeros(1, 3, 640, 640)