Upload app.py with huggingface_hub

app.py

@@ -203,21 +203,37 @@ def load_predictor() -> MoEPredictor:
         )
 
     # Ensure expert checkpoints are resolvable in the Space (paths inside ckpt are absolute)
-    def ensure_expert(name: str) -> Path:
+    def ensure_expert(name: str, comm: str) -> Path:
         """Return a local path to the expert checkpoint, downloading if needed."""
+        fname = Path(name).name
+        comm_tag = comm.replace("/", "_")
         local_candidates = [
-            REPO_ROOT / "experts" / name,
-            REPO_ROOT / name,
+            REPO_ROOT / "experts" / fname,
+            REPO_ROOT / fname,
+            REPO_ROOT / "experts" / f"{comm_tag}_expert.pth",
+            REPO_ROOT / f"{comm_tag}_expert.pth",
         ]
         for cand in local_candidates:
             if cand.exists():
                 return cand
-        # Download from model repo
-        downloaded = hf_hub.hf_hub_download(
-            repo_id="wi-lab/lwm-spectro",
-            filename=f"experts/{name}",
-        )
-        return Path(downloaded)
+        # Download from model repo with multiple filename guesses
+        download_candidates = [
+            f"experts/{fname}",
+            f"experts/{comm_tag}_expert.pth",
+            fname,
+        ]
+        last_err = None
+        for rel in download_candidates:
+            try:
+                downloaded = hf_hub.hf_hub_download(
+                    repo_id="wi-lab/lwm-spectro",
+                    filename=rel,
+                )
+                return Path(downloaded)
+            except Exception as exc:  # pragma: no cover - network/permissions issues
+                last_err = exc
+                continue
+        raise RuntimeError(f"Could not resolve expert checkpoint for {comm} ({fname}): {last_err}")
 
     # Rewrite expert paths into a temp checkpoint so MoEPredictor loads cleanly
     import torch  # local import to keep top import list compact
@@ -231,7 +247,8 @@ def load_predictor() -> MoEPredictor:
             if not ckpt_field:
                 continue
             fname = Path(ckpt_field).name
-            local_path = ensure_expert(fname)
+            comm = expert.get("comm", "unknown")
+            local_path = ensure_expert(fname, comm)
             if str(local_path) != ckpt_field:
                 expert["checkpoint"] = str(local_path)
                 patched = True