Sharris commited on
Commit
d549347
·
verified ·
1 Parent(s): de3c81a

Upload app.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. app.py +29 -17
app.py CHANGED
@@ -5,33 +5,45 @@ import tensorflow as tf
5
  from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
6
  import gradio as gr
7
 
8
- # Try to load a local SavedModel first, otherwise try to download from Hugging Face Hub
9
- MODEL_DIR = "saved_model"
10
  model = None
11
 
12
- if os.path.isdir(MODEL_DIR):
13
- try:
14
- model = tf.keras.models.load_model(MODEL_DIR)
15
- print(f"Loaded model from local path: {MODEL_DIR}")
16
- except Exception as e:
17
- print(f"Failed to load local model: {e}")
 
 
 
 
18
 
 
19
  if model is None:
20
- # If HF_MODEL_ID is set, attempt to download model files from the Hub
21
- HF_MODEL_ID = os.environ.get("HF_MODEL_ID")
22
- if HF_MODEL_ID:
 
 
 
 
 
 
 
23
  try:
24
  from huggingface_hub import snapshot_download
25
  repo_dir = snapshot_download(repo_id=HF_MODEL_ID)
26
- # Expecting a SavedModel dir inside the repo; try to load from repo root
27
- model = tf.keras.models.load_model(repo_dir)
28
- print(f"Loaded model from HF Hub repo: {HF_MODEL_ID}")
29
- except Exception as e:
30
- print(f"Failed to load model from HF Hub ({HF_MODEL_ID}): {e}")
 
31
 
32
  if model is None:
33
  raise RuntimeError(
34
- "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."
35
  )
36
 
37
  INPUT_SIZE = (224, 224)
 
5
  from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
6
  import gradio as gr
7
 
8
+ # Try to load model from various sources
 
9
  model = None
10
 
11
+ # Try local files first (for development)
12
+ local_model_paths = ["saved_model", "best_model.h5", "final_model.h5"]
13
+ for path in local_model_paths:
14
+ if os.path.exists(path):
15
+ try:
16
+ model = tf.keras.models.load_model(path, compile=False)
17
+ print(f"Loaded model from local path: {path}")
18
+ break
19
+ except Exception as e:
20
+ print(f"Failed to load local model from {path}: {e}")
21
 
22
+ # If no local model, try to download from Hugging Face Hub
23
  if model is None:
24
+ HF_MODEL_ID = os.environ.get("HF_MODEL_ID", "Sharris/age_detection_regression")
25
+ try:
26
+ from huggingface_hub import hf_hub_download
27
+ # Try to download the .h5 model file
28
+ model_path = hf_hub_download(repo_id=HF_MODEL_ID, filename="best_model.h5")
29
+ model = tf.keras.models.load_model(model_path, compile=False)
30
+ print(f"Loaded model from HF Hub: {HF_MODEL_ID}/best_model.h5")
31
+ except Exception as e:
32
+ print(f"Failed to load model from HF Hub ({HF_MODEL_ID}): {e}")
33
+ # Fallback: try to download entire repo and load from there
34
  try:
35
  from huggingface_hub import snapshot_download
36
  repo_dir = snapshot_download(repo_id=HF_MODEL_ID)
37
+ model_file = os.path.join(repo_dir, "best_model.h5")
38
+ if os.path.exists(model_file):
39
+ model = tf.keras.models.load_model(model_file, compile=False)
40
+ print(f"Loaded model from downloaded repo: {model_file}")
41
+ except Exception as e2:
42
+ print(f"Fallback download also failed: {e2}")
43
 
44
  if model is None:
45
  raise RuntimeError(
46
+ "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."
47
  )
48
 
49
  INPUT_SIZE = (224, 224)