Refactor model loading logic and improve error handling in app.py; add new .codex file

app.py

@@ -13,6 +13,7 @@ import logging
 import numpy as np
 import requests
 from io import BytesIO
+from threading import Lock
 
 
 # =============================================================================
@@ -28,19 +29,38 @@ logger = logging.getLogger(__name__)
 # =============================================================================
 
 MODEL_NAME = "datdevsteve/dinov2-nivra-finetuned"
+image_processor = None
+model = None
+id2label = {}
+model_load_lock = Lock()
 
-logger.info(f"[i] Loading model: {MODEL_NAME}")
 
-try:
-    image_processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
-    model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
-    model.eval()
-    logger.info("[i] Model loaded successfully")
-except Exception as e:
-    logger.error(f"[!] Error loading model: {e}")
-    raise
+def ensure_model_loaded() -> None:
+    global image_processor, model, id2label
 
-id2label = model.config.id2label if hasattr(model.config, "id2label") else {}
+    if image_processor is not None and model is not None:
+        return
+
+    with model_load_lock:
+        if image_processor is not None and model is not None:
+            return
+
+        logger.info(f"[i] Loading model: {MODEL_NAME}")
+
+        try:
+            image_processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
+            model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
+            model.eval()
+            id2label = (
+                model.config.id2label if hasattr(model.config, "id2label") else {}
+            )
+            logger.info("[i] Model loaded successfully")
+        except Exception as e:
+            logger.error(f"[!] Error loading model: {e}")
+            raise RuntimeError(
+                f"Unable to load model '{MODEL_NAME}'. Check Space network access "
+                f"and repository visibility. Original error: {e}"
+            ) from e
 
 
 # =============================================================================
@@ -62,6 +82,7 @@ def classify_confidence(score: float) -> str:
 
 def predict_medical_image(image: Image.Image, top_k: int = 5) -> Dict[str, Any]:
     try:
+        ensure_model_loaded()
         inputs = image_processor(images=image, return_tensors="pt")
 
         with torch.no_grad():
@@ -187,6 +208,10 @@ def create_demo():
     with gr.Blocks(title="Nivra Medical Image Classifier") as demo:
 
         gr.Markdown("# 🏥 Nivra Medical Image Classifier")
+        gr.Markdown(
+            "Upload an image or provide a public URL. "
+            "The model is loaded on the first prediction so the Space can start quickly."
+        )
 
         with gr.Row():
             with gr.Column(scale=2):

requirements.txt

@@ -9,9 +9,10 @@ transformers==4.41.2
 numpy<2
 
 Pillow>=10.0
+requests>=2.31.0
 
 # Validation / typing (3.11-compatible wheels exist)
 pydantic>=2.7,<3
 
 # Auth / crypto
-python-jose[cryptography]>=3.3.0
+python-jose[cryptography]>=3.3.0