Update app.py

app.py

@@ -1,17 +1,92 @@
-from huggingface_hub import snapshot_download
-import os, glob, torch
+import os
+import importlib.util
+import numpy as np
+import torch
+import gradio as gr
+from huggingface_hub import hf_hub_download
 
-local_dir = snapshot_download(repo_id="c1tr0n75/VoxelPathFinder")
+REPO_ID = "c1tr0n75/VoxelPathFinder"
 
-# preferred path in subfolder
-ckpt_path = os.path.join(local_dir, "training_outputs", "final_model.pth")
-if not os.path.exists(ckpt_path):
-  # fallback: find any .pth
-  pths = [p for p in glob.glob(os.path.join(local_dir, "**", "*.pth"), recursive=True)]
-  if not pths:
-    raise FileNotFoundError("No .pth found in repo snapshot.")
-  ckpt_path = pths[0]
+# 1) Make sure torch is imported, then define device BEFORE using it anywhere
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-ckpt = torch.load(ckpt_path, map_location=device)
+# 2) Download model code and weights from your model repo
+PY_PATH = hf_hub_download(repo_id=REPO_ID, filename="pathfinding_nn.py")
+CKPT_PATH = hf_hub_download(repo_id=REPO_ID, filename="training_outputs/final_model.pth")
+
+# 3) Dynamically import your model definitions
+spec = importlib.util.spec_from_file_location("pathfinding_nn", PY_PATH)
+mod = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(mod)
+PathfindingNetwork = mod.PathfindingNetwork
+create_voxel_input = mod.create_voxel_input
+
+# 4) Build and load model
+MODEL = PathfindingNetwork().to(DEVICE).eval()
+ckpt = torch.load(CKPT_PATH, map_location=DEVICE)
 state = ckpt.get("model_state_dict", ckpt)
-model.load_state_dict(state)
+MODEL.load_state_dict(state)
+
+ACTION_NAMES = ['FORWARD','BACK','LEFT','RIGHT','UP','DOWN']
+
+def decode(actions):
+    return [ACTION_NAMES[a] for a in actions if 0 <= a < 6]
+
+def infer_random(obstacle_prob=0.2, seed=None):
+    if seed is not None:
+        np.random.seed(int(seed))
+    voxel_dim = MODEL.voxel_dim
+    D, H, W = voxel_dim
+    obstacles = (np.random.rand(D, H, W) < float(obstacle_prob)).astype(np.float32)
+    free = np.argwhere(obstacles == 0)
+    if len(free) < 2:
+        return {"error": "Not enough free cells; lower obstacle_prob."}
+    s_idx, g_idx = np.random.choice(len(free), size=2, replace=False)
+    start = tuple(free[s_idx]); goal = tuple(free[g_idx])
+
+    voxel_np = create_voxel_input(obstacles, start, goal, voxel_dim=voxel_dim)
+    voxel = torch.from_numpy(voxel_np).float().unsqueeze(0).to(DEVICE)
+    pos = torch.tensor([[start, goal]], dtype=torch.long, device=DEVICE)
+
+    with torch.no_grad():
+        actions = MODEL(voxel, pos)[0].tolist()
+
+    return {
+        "start": start,
+        "goal": goal,
+        "num_actions": len([a for a in actions if 0 <= a < 6]),
+        "actions_ids": actions,
+        "actions_decoded": decode(actions)[:50],
+    }
+
+#def infer_npz(npz_file):
+#    if npz_file is None:
+#        return {"error": "Upload a .npz with 'voxel_data' and 'positions'."}
+#    data = np.load(npz_file.name)
+#    voxel = torch.from_numpy(data['voxel_data']).float().unsqueeze(0).to(DEVICE)
+#    pos = torch.from_numpy(data['positions']).long().unsqueeze(0).to(DEVICE)
+#    with torch.no_grad():
+#        actions = MODEL(voxel, pos)[0].tolist()
+#    return {
+#        "num_actions": len([a for a in actions if 0 <= a < 6]),
+#        "actions_ids": actions,
+#        "actions_decoded": decode(actions)[:50],
+#    }
+
+with gr.Blocks(title="Voxel Path Finder") as demo:
+    gr.Markdown("## 3D Voxel Path Finder — Inference")
+    with gr.Tab("Random environment"):
+        obstacle = gr.Slider(0.0, 0.9, value=0.2, step=0.05, label="Obstacle probability")
+        seed = gr.Number(value=None, label="Seed (optional)")
+        btn = gr.Button("Run inference")
+        out = gr.JSON(label="Result")
+        btn.click(infer_random, inputs=[obstacle, seed], outputs=out)
+
+    with gr.Tab("Upload .npz sample"):
+        file = gr.File(file_types=[".npz"], label="Upload sample (voxel_data, positions)")
+        btn2 = gr.Button("Run inference")
+        out2 = gr.JSON(label="Result")
+        btn2.click(infer_npz, inputs=file, outputs=out2)
+
+if __name__ == "__main__":
+    demo.launch()