Update app/configs.py

app/configs.py

@@ -11,9 +11,7 @@ import mlflow.pyfunc
 
 MLFLOW_URI = "https://omarelrayes-mlflow-server.hf.space"
 
-mlflow.set_tracking_uri(
-    MLFLOW_URI
-)
+mlflow.set_tracking_uri(MLFLOW_URI)
 
 
 # =========================
@@ -24,26 +22,50 @@ CLASSIFICATION_MODEL_URI = (
     "runs:/9893b5af3e414a75863d18340391078c/model"
 )
 
-
 SEGMENTATION_MODEL_URI = (
     "runs:/85a78da18bc94263848e61825b9e5104/model"
 )
 
 
 # =========================
-# Lazy Models
+# Lazy Models (Cache)
 # =========================
 
 _classification_model: Optional[
     mlflow.pyfunc.PyFuncModel
 ] = None
 
-
 _segmentation_model: Optional[
     mlflow.pyfunc.PyFuncModel
 ] = None
 
 
+# =========================
+# Helpers (IMPORTANT FIX)
+# =========================
+
+def _load_mlflow_model(model_uri: str):
+    """
+    Safe MLflow loader for Hugging Face / Docker environments:
+    runs:/ -> download_artifacts -> local path -> load_model
+    """
+
+    print(f"Downloading model: {model_uri}")
+
+    local_path = mlflow.artifacts.download_artifacts(model_uri)
+
+    print(f"Model downloaded to: {local_path}")
+
+    model = mlflow.pyfunc.load_model(local_path)
+
+    print("Model loaded successfully")
+
+    return model
+
+
+# =========================
+# Classification Model
+# =========================
 
 def get_classification_model():
 
@@ -51,17 +73,18 @@ def get_classification_model():
 
     if _classification_model is None:
 
-        print(
-            "Loading Classification Model..."
-        )
+        print("Loading Classification Model...")
 
-        _classification_model = mlflow.pyfunc.load_model(
+        _classification_model = _load_mlflow_model(
             CLASSIFICATION_MODEL_URI
         )
 
     return _classification_model
 
 
+# =========================
+# Segmentation Model
+# =========================
 
 def get_segmentation_model():
 
@@ -69,52 +92,42 @@ def get_segmentation_model():
 
     if _segmentation_model is None:
 
-        print(
-            "Loading Segmentation Model..."
-        )
+        print("Loading Segmentation Model...")
 
-        _segmentation_model = mlflow.pyfunc.load_model(
+        _segmentation_model = _load_mlflow_model(
             SEGMENTATION_MODEL_URI
         )
 
     return _segmentation_model
 
 
-
 # =========================
 # Metadata
 # =========================
 
 model_classes: Dict[int, str] = {
-
     0: "benign",
     1: "malicious"
-
 }
 
 
 request_history: List[Dict[str, Any]] = []
 
 
-
 # =========================
 # Storage
 # =========================
 
 STORAGE_DIR = Path("storage")
-
 IMAGES_DIR = STORAGE_DIR / "images"
-
 SEGMENTS_DIR = STORAGE_DIR / "segments"
 
 
-
 for dir_path in [
     STORAGE_DIR,
     IMAGES_DIR,
     SEGMENTS_DIR
 ]:
-
     dir_path.mkdir(
         parents=True,
         exist_ok=True