Upload app.py with huggingface_hub

app.py

@@ -202,6 +202,44 @@ def load_predictor() -> MoEPredictor:
             hf_hub.hf_hub_download(repo_id="wi-lab/lwm-spectro", filename="moe_checkpoint.pth")
         )
 
+    # Ensure expert checkpoints are resolvable in the Space (paths inside ckpt are absolute)
+    def ensure_expert(name: str) -> Path:
+        """Return a local path to the expert checkpoint, downloading if needed."""
+        local_candidates = [
+            REPO_ROOT / "experts" / name,
+            REPO_ROOT / name,
+        ]
+        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)
+
+    # Rewrite expert paths into a temp checkpoint so MoEPredictor loads cleanly
+    import torch  # local import to keep top import list compact
+
+    raw_ckpt = torch.load(ckpt_path, map_location="cpu")
+    experts = raw_ckpt.get("experts", [])
+    if experts:
+        patched = False
+        for expert in experts:
+            ckpt_field = expert.get("checkpoint")
+            if not ckpt_field:
+                continue
+            fname = Path(ckpt_field).name
+            local_path = ensure_expert(fname)
+            if str(local_path) != ckpt_field:
+                expert["checkpoint"] = str(local_path)
+                patched = True
+        if patched:
+            tmp_path = Path("/tmp/moe_checkpoint_patched.pth")
+            torch.save(raw_ckpt, tmp_path)
+            ckpt_path = tmp_path
+
     _predictor = MoEPredictor.from_checkpoint(ckpt_path)
     return _predictor