Update app.py

app.py

@@ -5,34 +5,49 @@ from PIL import Image
 import torch
 import torch.serialization
 
-# ✅ Add required classes for safe unpickling
-from ultralytics.nn.tasks import DetectionModel
+# ✅ Import all model components that may appear in custom YOLOv8 models
+from torch.nn import Sequential, Conv2d, BatchNorm2d, SiLU, ReLU, LeakyReLU
 from ultralytics.nn.modules.conv import Conv
-from torch.nn import Sequential, Conv2d, BatchNorm2d
+from ultralytics.nn.modules.block import Bottleneck, C2f, SPPF
+from ultralytics.nn.tasks import DetectionModel
 
+# ✅ Add all components to safe globals
 torch.serialization.add_safe_globals([
-    DetectionModel, Sequential, Conv, Conv2d, BatchNorm2d
+    DetectionModel,
+    Sequential,
+    Conv,
+    Conv2d,
+    BatchNorm2d,
+    SiLU,
+    ReLU,
+    LeakyReLU,
+    Bottleneck,
+    C2f,
+    SPPF,
 ])
 
-# 📥 Load model from Hugging Face
+# 📥 Download model
 model_path = hf_hub_download(repo_id="Safi029/ABD-model", filename="ABD.pt")
+
+# ✅ Load model safely
 model = YOLO(model_path)
 
-# 🧠 Define detection function
+# 🧠 Inference function
 def detect_structure(image):
     results = model(image)
     return Image.fromarray(results[0].plot())
 
-# 🎛️ Build Gradio interface
+# 🎛️ Gradio UI
 demo = gr.Interface(
     fn=detect_structure,
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
     title="YOLO Molecular Detector",
-    description="Upload a molecule image. YOLOv8 detects atoms or bonds."
+    description="Upload a molecular structure image. YOLOv8 will detect atoms and bonds."
 )
 
 demo.launch()
 
 
 
+