Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from ultralytics import YOLO
-from huggingface_hub import hf_hub_download, hf_hub_url
+from huggingface_hub import hf_hub_download
 from PIL import Image
 import torch
 import torch.serialization
@@ -9,14 +9,12 @@ import hashlib
 from typing import Optional
 
 # ===== MODEL COMPONENTS =====
-# Import all model components that may appear in custom YOLOv8 models
 from torch.nn import Sequential, Conv2d, BatchNorm2d, SiLU, ReLU, LeakyReLU, MaxPool2d, Upsample, ModuleList
 from ultralytics.nn.modules.conv import Conv
 from ultralytics.nn.modules.block import Bottleneck, C2f, SPPF
 from ultralytics.nn.tasks import DetectionModel
 
 # ===== SAFE GLOBALS =====
-# Add all components to safe globals for secure loading
 torch.serialization.add_safe_globals([
     DetectionModel,
     Sequential,
@@ -37,91 +35,82 @@ torch.serialization.add_safe_globals([
 # ===== MODEL CONFIG =====
 MODEL_REPO = "Safi029/ABD-model"
 MODEL_FILE = "ABD.pt"
-EXPECTED_SHA256 = "a1b2c3d4..."  # Replace with actual hash of your model file
+EXPECTED_SHA256 = "c3335b0cc6c504c4ac74b62bf2bc9aa06ecf402fa71184ec88f40a1f37979859"
 
 # ===== HELPER FUNCTIONS =====
-def verify_model(file_path: str, expected_hash: str) -> bool:
+def verify_model(file_path: str) -> bool:
     """Verify model integrity using SHA256 hash"""
-    if not os.path.exists(file_path):
-        return False
-        
     sha256 = hashlib.sha256()
     with open(file_path, "rb") as f:
         while chunk := f.read(8192):
             sha256.update(chunk)
-    return sha256.hexdigest() == expected_hash.lower()
+    actual_hash = sha256.hexdigest()
+    print(f"🔍 Model SHA256: {actual_hash}")
+    return actual_hash == EXPECTED_SHA256.lower()
 
 def download_model() -> str:
-    """Download model with verification"""
+    """Download and verify model"""
     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, EXPECTED_SHA256):
+    if not os.path.exists(model_path) or not verify_model(model_path):
         print("⬇️ Downloading model...")
         hf_hub_download(
             repo_id=MODEL_REPO,
             filename=MODEL_FILE,
             local_dir="models",
-            local_dir_use_symlinks=False,
-            resume_download=True,
             force_download=True
         )
         
-        if not verify_model(model_path, EXPECTED_SHA256):
+        if not verify_model(model_path):
             raise ValueError("❌ Downloaded model failed verification!")
     
     return model_path
 
 def load_model(model_path: str) -> YOLO:
-    """Safely load YOLO model with error handling"""
-    try:
-        print("🔧 Loading model...")
-        model = YOLO(model_path)
-        # Verify model loaded correctly by running a dummy inference
-        dummy_input = torch.randn(1, 3, 640, 640)
-        model(dummy_input)  # Test forward pass
-        print("✅ Model loaded and verified!")
-        return model
-    except Exception as e:
-        raise RuntimeError(f"Model loading failed: {str(e)}")
+    """Safely load YOLO model"""
+    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
+
+# ===== GRADIO INTERFACE =====
+def create_interface(model):
+    def detect_structure(image: Image.Image) -> Image.Image:
+        """Run detection on input image"""
+        try:
+            results = model(image)
+            return Image.fromarray(results[0].plot())
+        except Exception as e:
+            print(f"❌ Inference error: {e}")
+            error_img = Image.new("RGB", (300, 100), color="red")
+            return error_img
+
+    return gr.Interface(
+        fn=detect_structure,
+        inputs=gr.Image(type="pil", label="Input Image"),
+        outputs=gr.Image(type="pil", label="Detection Results"),
+        title="YOLOv8 Molecular Structure Detector",
+        description="🔬 Detect atoms and bonds in molecular structures",
+        allow_flagging="never"
+    )
 
 # ===== MAIN APPLICATION =====
 def main():
     try:
-        # 1. Download and verify model
-        model_path = download_model()
+        print(f"PyTorch: {torch.__version__}")
+        print(f"CUDA: {torch.cuda.is_available()}")
         
-        # 2. Load model
+        # Download and load model
+        model_path = download_model()
         model = load_model(model_path)
         
-        # 3. Gradio interface
-        def detect_structure(image: Image.Image) -> Image.Image:
-            """Run detection on input image"""
-            try:
-                results = model(image)
-                return Image.fromarray(results[0].plot())
-            except Exception as e:
-                print(f"❌ Inference error: {e}")
-                error_img = Image.new("RGB", (300, 100), color="red")
-                # Add error text
-                error_img.text((10, 45), f"Error: {str(e)[:50]}...")
-                return error_img
-
-        # 4. Create interface
-        demo = gr.Interface(
-            fn=detect_structure,
-            inputs=gr.Image(type="pil", label="Input Image"),
-            outputs=gr.Image(type="pil", label="Detection Results"),
-            title="YOLOv8 Molecular Structure Detector",
-            description=(
-                "🔬 Detect atoms and bonds in molecular structures. "
-                "Upload an image of a chemical structure to see detection results."
-            ),
-            allow_flagging="never"
-        )
-
-        # 5. Launch
+        # Create and launch interface
+        demo = create_interface(model)
         print("🚀 Starting Gradio interface...")
         demo.launch(server_name="0.0.0.0", share=False)
 
@@ -130,12 +119,6 @@ def main():
         raise
 
 if __name__ == "__main__":
-    # Check PyTorch and CUDA availability
-    print(f"PyTorch version: {torch.__version__}")
-    print(f"CUDA available: {torch.cuda.is_available()}")
-    
     main()
 
 
-
-