Upload streamlit_app.py

src/streamlit_app.py

@@ -968,17 +968,34 @@ def submit_uploaded_correction_pairs(nuc_files, myo_files, artifacts: dict, sele
 
 @st.cache_resource
 def load_model(device: str):
-    # MYOSEG_MODEL_PATH lets you run offline / test a local checkpoint.
-    local = os.environ.get("MYOSEG_MODEL_PATH") or hf_hub_download(
-        repo_id=MODEL_REPO_ID,
-        filename=MODEL_FILENAME,
-        force_download=False,
-    )
+    # Where the weights come from, in order:
+    #   1. MYOSEG_MODEL_PATH (env var / Space variable)
+    #   2. model_final.pt committed in the Space repo itself (next to src/ or inside it)
+    #   3. the Hugging Face model repo MODEL_REPO_ID (needs HF_TOKEN if that repo is private)
+    app_dir = Path(__file__).resolve().parent
+    candidates = [os.environ.get("MYOSEG_MODEL_PATH"),
+                  str(app_dir.parent / MODEL_FILENAME), str(app_dir / MODEL_FILENAME)]
+    local = next((c for c in candidates if c and os.path.exists(c)), None)
+    source = local
+    if local is None:
+        try:
+            local = hf_hub_download(repo_id=MODEL_REPO_ID, filename=MODEL_FILENAME,
+                                    force_download=False, token=os.environ.get("HF_TOKEN"))
+            source = f"hf://{MODEL_REPO_ID}/{MODEL_FILENAME}"
+        except Exception as exc:
+            st.error(
+                f"Could not load the model. No `{MODEL_FILENAME}` was found in the Space, and "
+                f"downloading it from `{MODEL_REPO_ID}` failed ({type(exc).__name__}). "
+                "Fix: add model_final.pt to the Space repo, or make the model repo public, "
+                "or add a valid HF_TOKEN secret."
+            )
+            st.stop()
     file_sha = sha256_file(local)
     info = {
         "sha256": file_sha,
         "modified": time.ctime(os.path.getmtime(local)),
         "size_mb": round(os.path.getsize(local) / 1e6, 2),
+        "source": source,
     }
     ckpt  = torch.load(local, map_location=device, weights_only=False)
     state = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
@@ -992,8 +1009,7 @@ def show_model_details(info: dict):
     """Drawn on every run (outside the cache) so it is always visible."""
     with st.sidebar.expander("Model details", expanded=False):
         st.caption(f"App version: `{APP_VERSION}`")
-        st.caption(f"Repo: `{MODEL_REPO_ID}`")
-        st.caption(f"File: `{MODEL_FILENAME}`")
+        st.caption(f"Loaded from: `{info.get('source', MODEL_REPO_ID)}`")
         st.caption(f"Size: {info['size_mb']:.2f} MB")
         st.caption(f"Modified: {info['modified']}")
         st.code(info["sha256"], language=None)