Update app/configs.py

app/configs.py

@@ -1,83 +1,89 @@
 import mlflow
-import mlflow.pyfunc
+import mlflow.tensorflow
 
 from pathlib import Path
 from typing import Dict, Any, List
 
+# ======================
+# MLflow Setup
+# ======================
 
 MLFLOW_URI = "https://omarelrayes-mlflow-server.hf.space"
+mlflow.set_tracking_uri(MLFLOW_URI)
 
-mlflow.set_tracking_uri(
-    MLFLOW_URI
-)
+# ======================
+# Model URIs
+# ======================
 
+CLASSIFICATION_MODEL_URI = "runs:/7766854bcd074c8e856aea7c47680bf4/model"
 
-CLASSIFICATION_MODEL_URI = (
-    "https://omarelrayes-mlflow-server.hf.space/"
-    "artifacts/1/efcf3f12eacc409daffe6a50888fa759/artifacts/model"
-)
-
-
-SEGMENTATION_MODEL_URI = (
-    "https://omarelrayes-mlflow-server.hf.space/"
-    "artifacts/1/85a78da18bc94263848e61825b9e5104/artifacts/model"
-)
+SEGMENTATION_MODEL_URI = "runs:/35f0bea9bf774ed59f2298a897c8cdf5/model"
 
+# ======================
+# Lazy Loaded Models
+# ======================
 
 _classification_model = None
 _segmentation_model = None
 
 
+# ======================
+# Classification Model
+# ======================
 
 def get_classification_model():
-
     global _classification_model
 
     if _classification_model is None:
-
         print("Loading Classification Model...")
 
-        _classification_model = mlflow.pyfunc.load_model(
+        _classification_model = mlflow.tensorflow.load_model(
             CLASSIFICATION_MODEL_URI
         )
 
-        print("Classification Loaded")
+        print("Classification Model Loaded Successfully")
 
     return _classification_model
 
 
-
+# ======================
+# Segmentation Model
+# ======================
 
 def get_segmentation_model():
-
     global _segmentation_model
 
-
     if _segmentation_model is None:
-
         print("Loading Segmentation Model...")
 
-        _segmentation_model = mlflow.pyfunc.load_model(
+        _segmentation_model = mlflow.tensorflow.load_model(
             SEGMENTATION_MODEL_URI
         )
 
-        print("Segmentation Model Loaded")
-
+        print("Segmentation Model Loaded Successfully")
 
     return _segmentation_model
 
 
-
+# ======================
+# Classes
+# ======================
 
 model_classes = {
     0: "benign",
     1: "malignant"
 }
 
+# ======================
+# Request History
+# ======================
+
 request_history: List[Dict[str, Any]] = []
 
 
-from pathlib import Path
+# ======================
+# Storage Setup
+# ======================
 
 STORAGE_DIR = Path("storage")
 IMAGES_DIR = STORAGE_DIR / "images"