Update main.py

main.py

@@ -1,5 +1,5 @@
 # main.py
-
+from huggingface_hub import hf_hub_download
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import JSONResponse, HTMLResponse
 import tensorflow as tf
@@ -207,14 +207,21 @@ def predict(payload: PredictRequest):
         # ---------------------------------------------------------------------
         # Run model (raw confidence)
         # ---------------------------------------------------------------------
-        MODEL_PATH = "/app/models/Forest_Segmentation_Best.keras"
-
+        MODEL_REPO = "prshntdxt/Forest_Segmentation_Best"
+        MODEL_FILE = "Forest_Segmentation_Best.keras"
+        
+        MODEL_PATH = hf_hub_download(
+            repo_id=MODEL_REPO,
+            filename=MODEL_FILE,
+        )
+        
         model = tf.keras.models.load_model(
             MODEL_PATH,
             compile=False
         )
 
 
+
         confidence_map = model.predict(
             input_tensor, verbose=0
         )[0, :, :, 0]