Update app.py

app.py

@@ -10,15 +10,17 @@ from typing import Optional
 
 # ===== MODEL COMPONENTS =====
 from torch.nn import Sequential, Conv2d, BatchNorm2d, SiLU, ReLU, LeakyReLU, MaxPool2d, Upsample, ModuleList
-from ultralytics.nn.modules.conv import Conv
+from ultralytics.nn.modules.conv import Conv, Concat  # Added Concat
 from ultralytics.nn.modules.block import Bottleneck, C2f, SPPF
 from ultralytics.nn.tasks import DetectionModel
 
 # ===== SAFE GLOBALS =====
+# Add all components including Concat to safe globals
 torch.serialization.add_safe_globals([
     DetectionModel,
     Sequential,
     Conv,
+    Concat,  # Explicitly added
     Conv2d,
     BatchNorm2d,
     SiLU,
@@ -53,7 +55,6 @@ def download_model() -> str:
     os.makedirs("models", exist_ok=True)
     model_path = os.path.join("models", MODEL_FILE)
     
-    # Download if not exists or verification fails
     if not os.path.exists(model_path) or not verify_model(model_path):
         print("⬇️ Downloading model...")
         hf_hub_download(
@@ -69,14 +70,18 @@ def download_model() -> str:
     return model_path
 
 def load_model(model_path: str) -> YOLO:
-    """Safely load YOLO model"""
+    """Safely load YOLO model with error handling"""
     print("🔧 Loading model...")
-    model = YOLO(model_path)
-    # Verify model works with dummy input
-    dummy_input = torch.randn(1, 3, 640, 640)
-    model(dummy_input)  # Test forward pass
-    print("✅ Model loaded and verified!")
-    return model
+    try:
+        model = YOLO(model_path)
+        # Test with small dummy input
+        with torch.no_grad():
+            dummy = torch.zeros(1, 3, 640, 640)
+            model(dummy)
+        print("✅ Model loaded and verified!")
+        return model
+    except Exception as e:
+        raise RuntimeError(f"Model loading failed: {str(e)}")
 
 # ===== GRADIO INTERFACE =====
 def create_interface(model):
@@ -96,7 +101,7 @@ def create_interface(model):
         outputs=gr.Image(type="pil", label="Detection Results"),
         title="YOLOv8 Molecular Structure Detector",
         description="🔬 Detect atoms and bonds in molecular structures",
-        allow_flagging="never"
+        examples=[["example.jpg"]] if os.path.exists("example.jpg") else None
     )
 
 # ===== MAIN APPLICATION =====
@@ -112,7 +117,11 @@ def main():
         # Create and launch interface
         demo = create_interface(model)
         print("🚀 Starting Gradio interface...")
-        demo.launch(server_name="0.0.0.0", share=False)
+        demo.launch(
+            server_name="0.0.0.0",
+            share=False,
+            server_port=7860
+        )
 
     except Exception as e:
         print(f"❌ Fatal error: {str(e)}")