Upload app.py with huggingface_hub

app.py

@@ -5,33 +5,45 @@ import tensorflow as tf
 from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
 import gradio as gr
 
-# Try to load a local SavedModel first, otherwise try to download from Hugging Face Hub
-MODEL_DIR = "saved_model"
+# Try to load model from various sources
 model = None
 
-if os.path.isdir(MODEL_DIR):
-    try:
-        model = tf.keras.models.load_model(MODEL_DIR)
-        print(f"Loaded model from local path: {MODEL_DIR}")
-    except Exception as e:
-        print(f"Failed to load local model: {e}")
+# Try local files first (for development)
+local_model_paths = ["saved_model", "best_model.h5", "final_model.h5"]
+for path in local_model_paths:
+    if os.path.exists(path):
+        try:
+            model = tf.keras.models.load_model(path, compile=False)
+            print(f"Loaded model from local path: {path}")
+            break
+        except Exception as e:
+            print(f"Failed to load local model from {path}: {e}")
 
+# If no local model, try to download from Hugging Face Hub
 if model is None:
-    # If HF_MODEL_ID is set, attempt to download model files from the Hub
-    HF_MODEL_ID = os.environ.get("HF_MODEL_ID")
-    if HF_MODEL_ID:
+    HF_MODEL_ID = os.environ.get("HF_MODEL_ID", "Sharris/age_detection_regression")
+    try:
+        from huggingface_hub import hf_hub_download
+        # Try to download the .h5 model file
+        model_path = hf_hub_download(repo_id=HF_MODEL_ID, filename="best_model.h5")
+        model = tf.keras.models.load_model(model_path, compile=False)
+        print(f"Loaded model from HF Hub: {HF_MODEL_ID}/best_model.h5")
+    except Exception as e:
+        print(f"Failed to load model from HF Hub ({HF_MODEL_ID}): {e}")
+        # Fallback: try to download entire repo and load from there
         try:
             from huggingface_hub import snapshot_download
             repo_dir = snapshot_download(repo_id=HF_MODEL_ID)
-            # Expecting a SavedModel dir inside the repo; try to load from repo root
-            model = tf.keras.models.load_model(repo_dir)
-            print(f"Loaded model from HF Hub repo: {HF_MODEL_ID}")
-        except Exception as e:
-            print(f"Failed to load model from HF Hub ({HF_MODEL_ID}): {e}")
+            model_file = os.path.join(repo_dir, "best_model.h5")
+            if os.path.exists(model_file):
+                model = tf.keras.models.load_model(model_file, compile=False)
+                print(f"Loaded model from downloaded repo: {model_file}")
+        except Exception as e2:
+            print(f"Fallback download also failed: {e2}")
 
 if model is None:
     raise RuntimeError(
-        "No model found. Place a SavedModel in './saved_model' or set HF_MODEL_ID env var to a Hugging Face model repo containing a SavedModel."
+        "No model found. Ensure 'best_model.h5' exists locally or set HF_MODEL_ID env var to a Hugging Face model repo containing the model."
     )
 
 INPUT_SIZE = (224, 224)