Update app.py

app.py

@@ -6,14 +6,21 @@ import torch
 import torch.serialization
 import os
 import hashlib
+import warnings
 from typing import Optional
 
-# ===== IMPORT ALL NECESSARY MODULES =====
+# ===== IMPORT ALL ULTRALYTICS MODULES =====
 from torch.nn import Sequential, Conv2d, BatchNorm2d, SiLU, ReLU, LeakyReLU, MaxPool2d, Upsample, ModuleList
-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
+from ultralytics.nn.modules import (
+    Conv, Concat, 
+    Bottleneck, C2f, SPPF, 
+    Detect, DFL,  # Added DFL
+    C2fAttn, ImagePoolingAttn,  # Common attention modules
+    HGStem, HGBlock,  # Additional blocks
+    AIFI,  # Additional modules
+    Segment, Pose, Classify, RTDETRDecoder  # Task-specific heads
+)
+from ultralytics.nn.tasks import DetectionModel, SegmentationModel, PoseModel, ClassificationModel
 
 # ===== SAFE GLOBALS CONFIGURATION =====
 # Add all components to safe globals
@@ -23,10 +30,14 @@ torch.serialization.add_safe_globals([
     MaxPool2d, Upsample, ModuleList,
     
     # Ultralytics modules
-    DetectionModel,
+    DetectionModel, SegmentationModel, PoseModel, ClassificationModel,
     Conv, Concat,
     Bottleneck, C2f, SPPF,
-    Detect,  # Added Detect head
+    Detect, DFL,  # Added DFL
+    C2fAttn, ImagePoolingAttn,
+    HGStem, HGBlock,
+    AIFI,
+    Segment, Pose, Classify, RTDETRDecoder
 ])
 
 # ===== MODEL CONFIG =====
@@ -68,10 +79,16 @@ def load_model(model_path: str) -> YOLO:
     """Safely load YOLO model with error handling"""
     print("🔧 Loading model...")
     try:
-        # Temporary workaround for PyTorch 2.6+ weights_only=True issues
-        # Only use if you trust the model source!
+        # Temporary monkey patch for PyTorch 2.6+ weights_only restriction
+        # ONLY USE IF YOU TRUST THE MODEL SOURCE!
+        original_load = torch.load
+        torch.load = lambda *args, **kwargs: original_load(*args, **kwargs, weights_only=False)
+        
         model = YOLO(model_path, task='detect')
         
+        # Restore original torch.load
+        torch.load = original_load
+        
         # Test with small dummy input
         with torch.no_grad():
             dummy = torch.zeros(1, 3, 640, 640)
@@ -79,6 +96,9 @@ def load_model(model_path: str) -> YOLO:
         print("✅ Model loaded and verified!")
         return model
     except Exception as e:
+        # Ensure original torch.load is restored even if error occurs
+        if 'original_load' in locals():
+            torch.load = original_load
         raise RuntimeError(f"Model loading failed: {str(e)}")
 
 # ===== GRADIO INTERFACE =====
@@ -126,6 +146,8 @@ def main():
         raise
 
 if __name__ == "__main__":
+    # Suppress torch.load warnings
+    warnings.filterwarnings("ignore", category=UserWarning, message="torch.load")
     main()
-
+