Update app/configs.py

app/configs.py

@@ -5,6 +5,7 @@ import tensorflow as tf
 from pathlib import Path
 from typing import Dict, Any, List
 import os
+import glob
 
 # ======================
 # 🔥 Set PYTHONPATH for sitecustomize.py
@@ -61,6 +62,29 @@ 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
 # ======================
@@ -74,21 +98,18 @@ def get_classification_model():
         print("="*60)
 
         try:
-            # 🔥 Download artifacts from MLflow (will use HF Hub via plugin)
+            # 🔥 Download ALL artifacts from MLflow
             local_path = mlflow.artifacts.download_artifacts(
                 run_id=CLASSIFICATION_RUN_ID,
-                artifact_path="model",
                 dst_path="/tmp/mlflow_models/classification"
             )
             
             print(f"📥 Downloaded to: {local_path}")
             
-            # 🔥 Load model from local path
-            model_file = os.path.join(local_path, "data", "model.keras")
-            
-            if not os.path.exists(model_file):
-                raise FileNotFoundError(f"Model file not found: {model_file}")
+            # 🔥 Find model.keras in the downloaded files
+            model_file = find_model_file(local_path)
             
+            # 🔥 Load model from local path
             _classification_model = tf.keras.models.load_model(
                 model_file,
                 compile=False
@@ -116,21 +137,18 @@ def get_segmentation_model():
         print("="*60)
 
         try:
-            # 🔥 Download artifacts from MLflow (will use HF Hub via plugin)
+            # 🔥 Download ALL artifacts from MLflow
             local_path = mlflow.artifacts.download_artifacts(
                 run_id=SEGMENTATION_RUN_ID,
-                artifact_path="model",
                 dst_path="/tmp/mlflow_models/segmentation"
             )
             
             print(f"📥 Downloaded to: {local_path}")
             
-            # 🔥 Load model from local path
-            model_file = os.path.join(local_path, "data", "model.keras")
-            
-            if not os.path.exists(model_file):
-                raise FileNotFoundError(f"Model file not found: {model_file}")
+            # 🔥 Find model.keras in the downloaded files
+            model_file = find_model_file(local_path)
             
+            # 🔥 Load model from local path
             _segmentation_model = tf.keras.models.load_model(
                 model_file,
                 compile=False