bugfix

main.py

@@ -13,20 +13,18 @@ from typing import Any, Dict
 import os
 import pkgutil
 
-
 from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_url
+import requests
+import tempfile
+import shutil
 from typing import Any, Dict
 
 import tensorflow as tf
-
-
-
+import traceback
+import logging
 from tensorflow import keras
 
-
-
-
-
 app = FastAPI(title="1.3 - AI Model Deployment")
 ''' browser: http://localhost:8000/docs'''
 
@@ -40,10 +38,6 @@ app.add_middleware(
 )
 
 
-
-
-
-
 ANIMALS = ['Cat', 'Dog', 'Panda'] # Animal names here, these represent the labels of the images that we trained our model on.
 
 # 1) download your SavedModel from the Hub into a writable directory (Spaces often
@@ -57,24 +51,64 @@ try:
 except PermissionError:
     raise RuntimeError(f"Cannot create model directory '{local_model_dir}'. Ensure the process has write access or set HF_MODEL_DIR to a writable path.")
 
-# download files into local_model_dir
-hf_hub_download(repo_id, filename="config.json", repo_type="model", local_dir=local_model_dir)
-hf_hub_download(repo_id, filename="metadata.json", repo_type="model", local_dir=local_model_dir)
-hf_hub_download(repo_id, filename="model.weights.h5", repo_type="model", local_dir=local_model_dir)
-
-# 2) load it
-# If the hub returns an unpacked model folder, point load_model at that directory.
+# download files into local_model_dir and load model with resilient error handling
+model = None
+model_load_error = None
 try:
-    model = tf.keras.models.load_model(local_model_dir)
-except (IOError, OSError, ValueError) as e:
-    raise RuntimeError(f"Failed to load model from '{local_model_dir}': {e}")
-
+    # First try using a cache dir so downloads happen in a shared cache and final move
+    # into local_model_dir is less likely to require risky renames inside the repo.
+    cache_dir = '/tmp/.cache/huggingface'
+    os.makedirs(cache_dir, exist_ok=True)
 
+    hf_hub_download(repo_id, filename="config.json", repo_type="model", local_dir=local_model_dir, cache_dir=cache_dir)
+    hf_hub_download(repo_id, filename="metadata.json", repo_type="model", local_dir=local_model_dir, cache_dir=cache_dir)
+    hf_hub_download(repo_id, filename="model.weights.h5", repo_type="model", local_dir=local_model_dir, cache_dir=cache_dir)
 
+    # 2) load it
+    model = tf.keras.models.load_model(local_model_dir)
+    logging.info(f"Model loaded successfully from {local_model_dir}")
+except Exception as e:
+    # Primary download attempt failed -> try streamed fallback which writes directly
+    # into local_model_dir (avoids internal temp/move when those operations are blocked).
+    primary_err = e
+    logging.error("Primary hf_hub_download failed, attempting streamed fallback: %s", e)
+    try:
+        for filename in ("config.json", "metadata.json", "model.weights.h5"):
+            url = hf_hub_url(repo_id=repo_id, filename=filename, repo_type='model')
+            logging.info(f"Streaming {filename} from {url} into {local_model_dir}")
+            resp = requests.get(url, stream=True, timeout=60)
+            resp.raise_for_status()
+            # write to a temp file inside the target dir then move atomically
+            with tempfile.NamedTemporaryFile(dir=local_model_dir, delete=False) as tmpf:
+                for chunk in resp.iter_content(chunk_size=8192):
+                    if chunk:
+                        tmpf.write(chunk)
+                tmp_path = tmpf.name
+            final_path = os.path.join(local_model_dir, filename)
+            try:
+                shutil.move(tmp_path, final_path)
+            except Exception:
+                # if atomic move fails, try copy+remove
+                shutil.copy(tmp_path, final_path)
+                os.remove(tmp_path)
+
+        # After streamed download, try loading
+        model = tf.keras.models.load_model(local_model_dir)
+        logging.info(f"Model loaded successfully from {local_model_dir} after streamed fallback")
+    except Exception as e2:
+        model_load_error = f"primary: {primary_err}; fallback: {e2}"
+        tb = traceback.format_exc()
+        logging.error("Streamed fallback failed: %s", e2)
+        logging.error(tb)
+        model = None
 
 
 @app.post('/upload/image')
 async def uploadImage(img: UploadFile = File(...)):
+    if model is None:
+        # Model isn't available — return a helpful error to the caller instead of crashing.
+        return fastapi.Response(status_code=503, content=f"Model not loaded: {model_load_error}")
+
     original_image = Image.open(img.file) # Read the bytes and process as an image
     if original_image.mode == 'RGBA':
         original_image = original_image.convert('RGB')
@@ -111,6 +145,8 @@ def versions() -> Dict[str, Any]:
 def predict_stub() -> Dict[str, Any]:
 
     # This is a stub, so we're not doing a real prediction
+    if model is None:
+        return {"prediction": "model not loaded", "error": model_load_error}
     return {"prediction": "stub, we're not doing a real prediction"}