Update app/configs.py

app/configs.py

@@ -1,134 +1,61 @@
-from pathlib import Path
-from typing import Dict, Any, List, Optional
-
 import mlflow
 import mlflow.pyfunc
-
-
-# =========================
-# MLflow Config
-# =========================
+from typing import Optional, Dict, Any, List
 
 MLFLOW_URI = "https://omarelrayes-mlflow-server.hf.space"
 
 mlflow.set_tracking_uri(MLFLOW_URI)
 
+CLASSIFICATION_MODEL_URI = "runs:/efcf3f12eacc409daffe6a50888fa759/model"
+SEGMENTATION_MODEL_URI = "runs:/d5de5322822b45349b7bd311e747c794/model"
 
-# =========================
-# Run IDs from MLflow
-# =========================
-
-CLASSIFICATION_MODEL_URI = (
-    "runs:/efcf3f12eacc409daffe6a50888fa759/model"
-)
-
-SEGMENTATION_MODEL_URI = (
-    "runs:/d5de5322822b45349b7bd311e747c794/model"
-)
-
+_classification_model = None
+_segmentation_model = None
 
-# =========================
-# Lazy Models (Cache)
-# =========================
 
-_classification_model: Optional[
-    mlflow.pyfunc.PyFuncModel
-] = None
+def load_model(uri: str):
+    print(f"Downloading model: {uri}")
 
-_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)
+    local_path = mlflow.artifacts.download_artifacts(uri)
 
     print(f"Model downloaded to: {local_path}")
 
-    model = mlflow.pyfunc.load_model(local_path)
-
-    print("Model loaded successfully")
+    return mlflow.pyfunc.load_model(local_path)
 
-    return model
-
-
-# =========================
-# Classification Model
-# =========================
 
 def get_classification_model():
-
     global _classification_model
 
     if _classification_model is None:
-
         print("Loading Classification Model...")
-
-        _classification_model = _load_mlflow_model(
-            CLASSIFICATION_MODEL_URI
-        )
+        _classification_model = load_model(CLASSIFICATION_MODEL_URI)
 
     return _classification_model
 
 
-# =========================
-# Segmentation Model
-# =========================
-
 def get_segmentation_model():
-
     global _segmentation_model
 
     if _segmentation_model is None:
-
         print("Loading Segmentation Model...")
-
-        _segmentation_model = _load_mlflow_model(
-            SEGMENTATION_MODEL_URI
-        )
+        _segmentation_model = load_model(SEGMENTATION_MODEL_URI)
 
     return _segmentation_model
 
 
-# =========================
-# Metadata
-# =========================
-
-model_classes: Dict[int, str] = {
+model_classes = {
     0: "benign",
-    1: "malicious"
+    1: "malignant"
 }
 
-
 request_history: List[Dict[str, Any]] = []
 
 
-# =========================
-# Storage
-# =========================
+from pathlib import Path
 
 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
-    )
+for d in [STORAGE_DIR, IMAGES_DIR, SEGMENTS_DIR]:
+    d.mkdir(parents=True, exist_ok=True)