Update app/configs.py

app/configs.py

@@ -1,11 +1,10 @@
 # configs.py
 import mlflow
-import mlflow.artifacts
 import tensorflow as tf
 from pathlib import Path
 from typing import Dict, Any, List
 import os
-import glob
+from huggingface_hub import hf_hub_download
 
 # ======================
 # 🔥 Set PYTHONPATH for sitecustomize.py
@@ -13,7 +12,7 @@ import glob
 os.environ["PYTHONPATH"] = "/app:" + os.environ.get("PYTHONPATH", "")
 
 # ======================
-# 🔥 Import HF Plugin
+# 🔥 Import HF Plugin (for MLflow compatibility)
 # ======================
 try:
     from hf_artifact_repo import HuggingFaceArtifactRepository
@@ -44,16 +43,26 @@ SEGMENTATION_RUN_ID = os.getenv(
     "c321c4e944064127a1198e781c83cad8"
 )
 
+# ======================
+# HuggingFace Hub Configuration
+# ======================
+HF_REPO = "omarelrayes/mlflow-artifacts"
+
+# 🔥 المسارات الحقيقية في HF Hub
+CLASSIFICATION_MODEL_PATH = "1/models/m-7152eeee57b24d1eb3b0312c9de6354c/artifacts/data/model.keras"
+SEGMENTATION_MODEL_PATH = "2/models/m-7c6aea2dc3f64395a63a8b5ea90d6994/artifacts/data/model.keras"
+
 # ======================
 # MLflow Configuration
 # ======================
 mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
 
 print("="*60)
-print("🔧 MLflow Configuration:")
-print(f"   - Tracking URI: {MLFLOW_TRACKING_URI}")
+print("🔧 Configuration:")
+print(f"   - MLflow Tracking URI: {MLFLOW_TRACKING_URI}")
 print(f"   - Classification Run ID: {CLASSIFICATION_RUN_ID}")
 print(f"   - Segmentation Run ID: {SEGMENTATION_RUN_ID}")
+print(f"   - HF Repository: {HF_REPO}")
 print("="*60)
 
 # ======================
@@ -62,29 +71,6 @@ print("="*60)
 _classification_model = None
 _segmentation_model = None
 
-# ======================
-# Helper Function: Find model.keras in directory
-# ======================
-def find_model_file(directory):
-    """Search for model.keras in directory and subdirectories"""
-    print(f"🔍 Searching for model.keras in {directory}...")
-    
-    # Search recursively
-    matches = glob.glob(os.path.join(directory, "**", "model.keras"), recursive=True)
-    
-    if matches:
-        model_file = matches[0]
-        print(f"✅ Found model.keras at: {model_file}")
-        return model_file
-    
-    # List all files for debugging
-    print(f"📂 Files in {directory}:")
-    for root, dirs, files in os.walk(directory):
-        for file in files:
-            print(f"   - {os.path.join(root, file)}")
-    
-    raise FileNotFoundError(f"model.keras not found in {directory}")
-
 # ======================
 # Classification Model Loader
 # ======================
@@ -94,28 +80,40 @@ def get_classification_model():
     if _classification_model is None:
         print("="*60)
         print("🔄 Loading Classification Model...")
-        print(f"🔗 Run ID: {CLASSIFICATION_RUN_ID}")
+        print(f"📦 From HuggingFace Hub: {HF_REPO}")
+        print(f"📁 Path: {CLASSIFICATION_MODEL_PATH}")
+        print(f"🔗 MLflow Run ID: {CLASSIFICATION_RUN_ID}")
         print("="*60)
 
         try:
-            # 🔥 Download ALL artifacts from MLflow
-            local_path = mlflow.artifacts.download_artifacts(
-                run_id=CLASSIFICATION_RUN_ID,
-                dst_path="/tmp/mlflow_models/classification"
+            # 🔥 Step 1: Download from HuggingFace Hub
+            model_path = hf_hub_download(
+                repo_id=HF_REPO,
+                filename=CLASSIFICATION_MODEL_PATH,
+                repo_type="model",
+                token=HF_TOKEN,
+                cache_dir="/tmp/hf_cache"
             )
             
-            print(f"📥 Downloaded to: {local_path}")
-            
-            # 🔥 Find model.keras in the downloaded files
-            model_file = find_model_file(local_path)
+            print(f"📥 Model downloaded to: {model_path}")
             
-            # 🔥 Load model from local path
+            # 🔥 Step 2: Load model
             _classification_model = tf.keras.models.load_model(
-                model_file,
+                model_path,
                 compile=False
             )
             
             print("✅ Classification Model Loaded Successfully!")
+            
+            # 🔥 Step 3: Log to MLflow (for tracking)
+            with mlflow.start_run(run_id=CLASSIFICATION_RUN_ID, nested=True):
+                mlflow.log_param("model_source", "huggingface_hub")
+                mlflow.log_param("hf_repo", HF_REPO)
+                mlflow.log_param("model_path", CLASSIFICATION_MODEL_PATH)
+                mlflow.log_param("status", "loaded")
+            
+            print("✅ Logged to MLflow for tracking!")
+            
         except Exception as e:
             print(f"❌ Error loading classification model: {e}")
             import traceback
@@ -133,28 +131,40 @@ def get_segmentation_model():
     if _segmentation_model is None:
         print("="*60)
         print("🔄 Loading Segmentation Model...")
-        print(f"🔗 Run ID: {SEGMENTATION_RUN_ID}")
+        print(f"📦 From HuggingFace Hub: {HF_REPO}")
+        print(f"📁 Path: {SEGMENTATION_MODEL_PATH}")
+        print(f"🔗 MLflow Run ID: {SEGMENTATION_RUN_ID}")
         print("="*60)
 
         try:
-            # 🔥 Download ALL artifacts from MLflow
-            local_path = mlflow.artifacts.download_artifacts(
-                run_id=SEGMENTATION_RUN_ID,
-                dst_path="/tmp/mlflow_models/segmentation"
+            # 🔥 Step 1: Download from HuggingFace Hub
+            model_path = hf_hub_download(
+                repo_id=HF_REPO,
+                filename=SEGMENTATION_MODEL_PATH,
+                repo_type="model",
+                token=HF_TOKEN,
+                cache_dir="/tmp/hf_cache"
             )
             
-            print(f"📥 Downloaded to: {local_path}")
-            
-            # 🔥 Find model.keras in the downloaded files
-            model_file = find_model_file(local_path)
+            print(f"📥 Model downloaded to: {model_path}")
             
-            # 🔥 Load model from local path
+            # 🔥 Step 2: Load model
             _segmentation_model = tf.keras.models.load_model(
-                model_file,
+                model_path,
                 compile=False
             )
             
             print("✅ Segmentation Model Loaded Successfully!")
+            
+            # 🔥 Step 3: Log to MLflow (for tracking)
+            with mlflow.start_run(run_id=SEGMENTATION_RUN_ID, nested=True):
+                mlflow.log_param("model_source", "huggingface_hub")
+                mlflow.log_param("hf_repo", HF_REPO)
+                mlflow.log_param("model_path", SEGMENTATION_MODEL_PATH)
+                mlflow.log_param("status", "loaded")
+            
+            print("✅ Logged to MLflow for tracking!")
+            
         except Exception as e:
             print(f"❌ Error loading segmentation model: {e}")
             import traceback