Update app/configs.py

app/configs.py

@@ -1,48 +1,74 @@
 import mlflow
 import mlflow.pyfunc
-from typing import Optional, Dict, Any, List
+from typing import Optional
+
 
 MLFLOW_URI = "https://omarelrayes-mlflow-server.hf.space"
 
-mlflow.set_tracking_uri(MLFLOW_URI)
+mlflow.set_tracking_uri(
+    MLFLOW_URI
+)
 
-CLASSIFICATION_MODEL_URI = "runs:/efcf3f12eacc409daffe6a50888fa759/model"
-SEGMENTATION_MODEL_URI = "runs:/d5de5322822b45349b7bd311e747c794/model"
 
-_classification_model = None
-_segmentation_model = None
+CLASSIFICATION_MODEL_URI = (
+    "runs:/efcf3f12eacc409daffe6a50888fa759/model"
+)
 
 
-def load_model(uri: str):
-    print(f"Downloading model: {uri}")
+SEGMENTATION_MODEL_URI = (
+    "runs:/d5de5322822b45349b7bd311e747c794/model"
+)
 
-    local_path = mlflow.artifacts.download_artifacts(uri)
 
-    print(f"Model downloaded to: {local_path}")
 
-    return mlflow.pyfunc.load_model(local_path)
+_classification_model = None
+_segmentation_model = None
+
 
 
 def get_classification_model():
+
     global _classification_model
 
+
     if _classification_model is None:
+
         print("Loading Classification Model...")
-        _classification_model = load_model(CLASSIFICATION_MODEL_URI)
+
+        _classification_model = mlflow.pyfunc.load_model(
+            CLASSIFICATION_MODEL_URI
+        )
+
+        print("Classification Model Loaded")
+
 
     return _classification_model
 
 
+
+
+
 def get_segmentation_model():
+
     global _segmentation_model
 
+
     if _segmentation_model is None:
+
         print("Loading Segmentation Model...")
-        _segmentation_model = load_model(SEGMENTATION_MODEL_URI)
+
+        _segmentation_model = mlflow.pyfunc.load_model(
+            SEGMENTATION_MODEL_URI
+        )
+
+        print("Segmentation Model Loaded")
+
 
     return _segmentation_model
 
 
+
+
 model_classes = {
     0: "benign",
     1: "malignant"