local and hg setup

app.py

@@ -1,5 +1,10 @@
 from __future__ import annotations
 
+from env import IS_HF_SPACE
+
+if IS_HF_SPACE:
+    import spaces  # noqa: F401
+
 from app_ui import build_demo
 
 demo = build_demo()

app_ui.py

@@ -8,12 +8,12 @@ from config import (
     K_PER_ANGLE_DEFAULT,
     TOPK_DEFAULT,
 )
+from env import gpu_decorator
 from runtime import RUNTIME
 from search_service import run_search
-from spaces_compat import GPU
 
 
-@GPU(duration=120)
+@gpu_decorator(duration=120)
 def run_search_entry(
     image,
     top_k,
@@ -41,7 +41,6 @@ def run_search_entry(
         a270=a270,
     )
 
-
 def build_demo() -> gr.Blocks:
     with gr.Blocks(title="MouseBrain 2D→3D Demo") as demo:
         gr.Markdown(

config.py

@@ -3,55 +3,10 @@ from __future__ import annotations
 import os
 from pathlib import Path
 
-import torch
-
-
-def _env_flag(name: str, default: bool = False) -> bool:
-    value = os.getenv(name)
-    if value is None:
-        return default
-    return value.strip().lower() in {"1", "true", "yes", "on"}
-
-
-def _resolve_device() -> str:
-    requested = os.getenv("MOUSE_DEMO_DEVICE", "auto").strip().lower()
-
-    if requested == "auto":
-        return "cuda" if torch.cuda.is_available() else "cpu"
-
-    if requested == "cuda":
-        return "cuda" if torch.cuda.is_available() else "cpu"
-
-    if requested in {"cpu", "mps"}:
-        return requested
-
-    return "cpu"
-
-
 APP_DIR = Path(__file__).resolve().parent
 CACHE_DIR = Path(os.getenv("MOUSE_DEMO_CACHE", APP_DIR / ".cache"))
 CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
-# Environment
-IS_HF_SPACE = _env_flag("SPACE_ID", False) or _env_flag("HF_SPACE", False)
-LOCAL_DEV = _env_flag("MOUSE_DEMO_LOCAL_DEV", not IS_HF_SPACE)
-
-# Runtime
-DEVICE = _resolve_device()
-TOPK_DEFAULT = int(os.getenv("TOPK_DEFAULT", "12"))
-K_PER_ANGLE_DEFAULT = int(os.getenv("K_PER_ANGLE_DEFAULT", "100"))
-DEFAULT_RERANK_TOPK = int(os.getenv("DEFAULT_RERANK_TOPK", "100"))
-DEFAULT_RERANK_ALPHA = float(os.getenv("DEFAULT_RERANK_ALPHA", "1.0"))
-DEFAULT_ANGLES = (0.0, 90.0, 180.0, 270.0)
-
-# Unload behavior
-# In local dev, keep models loaded unless explicitly overridden.
-if "IDLE_UNLOAD_SECONDS" in os.environ:
-    IDLE_UNLOAD_SECONDS = int(os.environ["IDLE_UNLOAD_SECONDS"])
-else:
-    IDLE_UNLOAD_SECONDS = 0 if LOCAL_DEV else 120
-
-# Remote artifacts
 GCS_BASE = os.getenv(
     "GCS_BASE",
     "https://storage.googleapis.com/mouse-rag/demo/v2",
@@ -62,20 +17,18 @@ MANIFEST_URL = f"{GCS_BASE}/patch_manifest.parquet"
 VECTORS_URL = f"{GCS_BASE}/patch_vectors.npy"
 PATCH_PNG_BASE = f"{GCS_BASE}/patch_png"
 
-# Optional local overrides for development
-FAISS_LOCAL_OVERRIDE = os.getenv("FAISS_LOCAL_OVERRIDE")
-MANIFEST_LOCAL_OVERRIDE = os.getenv("MANIFEST_LOCAL_OVERRIDE")
-VECTORS_LOCAL_OVERRIDE = os.getenv("VECTORS_LOCAL_OVERRIDE")
-PATCH_PNG_BASE_OVERRIDE = os.getenv("PATCH_PNG_BASE_OVERRIDE")
-
-# Hugging Face reranker repo
 RERANKER_REPO_ID = os.getenv("RERANKER_REPO_ID", "FrankHUP/mousebrain-reranker")
 RERANKER_FILENAME = os.getenv("RERANKER_FILENAME", "reranker_listwise.pt")
-
-# DINO model id
 DINO_MODEL_ID = os.getenv("DINO_MODEL_ID", "facebook/dinov3-vitb16-pretrain-lvd1689m")
 
-# Cached artifact paths
 FAISS_LOCAL = CACHE_DIR / "patch_index.faiss"
 MANIFEST_LOCAL = CACHE_DIR / "patch_manifest.parquet"
-VECTORS_LOCAL = CACHE_DIR / "patch_vectors.npy"
+VECTORS_LOCAL = CACHE_DIR / "patch_vectors.npy"
+
+TOPK_DEFAULT = int(os.getenv("TOPK_DEFAULT", "12"))
+K_PER_ANGLE_DEFAULT = int(os.getenv("K_PER_ANGLE_DEFAULT", "100"))
+IDLE_UNLOAD_SECONDS = int(os.getenv("IDLE_UNLOAD_SECONDS", "120"))
+
+DEFAULT_ANGLES = (0.0, 90.0, 180.0, 270.0)
+DEFAULT_RERANK_TOPK = 100
+DEFAULT_RERANK_ALPHA = 1.0

device_utils.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import os
+
+
+def resolve_device() -> str:
+    requested = os.getenv("MOUSE_DEMO_DEVICE", "auto").strip().lower()
+
+    if requested == "cpu":
+        return "cpu"
+
+    if requested == "cuda":
+        import torch
+        return "cuda" if torch.cuda.is_available() else "cpu"
+
+    # auto
+    import torch
+    return "cuda" if torch.cuda.is_available() else "cpu"

env.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+import os
+
+
+def _env_flag(name: str) -> bool:
+    value = os.getenv(name)
+    if value is None:
+        return False
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+IS_HF_SPACE = bool(os.getenv("SPACE_ID")) or _env_flag("HF_SPACE")
+
+
+def gpu_decorator(*args, **kwargs):
+    if IS_HF_SPACE:
+        import spaces
+        return spaces.GPU(*args, **kwargs)
+
+    def decorator(fn):
+        return fn
+
+    return decorator

runtime.py

@@ -10,7 +10,6 @@ from huggingface_hub import hf_hub_download
 
 from config import (
     CACHE_DIR,
-    DEVICE,
     FAISS_LOCAL,
     FAISS_URL,
     IDLE_UNLOAD_SECONDS,
@@ -87,10 +86,10 @@ class RuntimeState:
 
     def status_text(self) -> str:
         if self.loading:
-            return f"Loading assets on {DEVICE}..."
+            return "Loading assets..."
         if self.store is None:
-            return f"Cold state ({DEVICE})"
-        return f"Ready ({DEVICE})"
+            return "Cold state"
+        return "Ready"
 
 
 RUNTIME = RuntimeState()

search_service.py

@@ -3,8 +3,9 @@ from __future__ import annotations
 import numpy as np
 import pandas as pd
 
-from config import DEVICE, PATCH_PNG_BASE, PATCH_PNG_BASE_OVERRIDE
+from config import PATCH_PNG_BASE
 from runtime import RUNTIME
+from device_utils import resolve_device
 
 
 def _angles_from_flags(a0: bool, a90: bool, a180: bool, a270: bool) -> tuple[float, ...]:
@@ -36,6 +37,8 @@ def run_search(
     if image is None:
         return [], pd.DataFrame(), "Please upload an image."
 
+    device = resolve_device()
+
     RUNTIME.ensure_loaded()
     assert RUNTIME.store is not None
 
@@ -50,9 +53,9 @@ def run_search(
         use_reranker=bool(use_reranker),
         rerank_topk=int(rerank_topk),
         reranker_model_path=RUNTIME.reranker_model_path,
-        reranker_device=DEVICE,
+        reranker_device=device,
         reranker_batch_size=128,
-        reranker_use_fp16=(DEVICE == "cuda"),
+        reranker_use_fp16=(device == "cuda"),
         reranker_compile=False,
         rerank_alpha=float(rerank_alpha),
     )
@@ -62,14 +65,14 @@ def run_search(
     df = searcher.to_dataframe(hits).copy()
 
     if df.empty:
-        return [], df, f"No results ({DEVICE})."
+        return [], df, f"No results ({device})."
 
-    patch_base = PATCH_PNG_BASE_OVERRIDE or PATCH_PNG_BASE
-    df["image_url"] = df["patch_id"].apply(lambda pid: f"{patch_base}/{int(pid):08d}.png")
+    df["image_url"] = df["patch_id"].apply(lambda pid: f"{PATCH_PNG_BASE}/{int(pid):08d}.png")
 
     gallery = [
         (row["image_url"], f"patch_id={int(row['patch_id'])} score={row['score']:.4f}")
         for _, row in df.iterrows()
     ]
 
-    return gallery, df, f"Completed on {DEVICE}"
+    mode = "GPU (ZeroGPU)" if device == "cuda" else "CPU"
+    return gallery, df, f"Completed on {mode}"

spaces_compat.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
-from config import IS_HF_SPACE
+import os
 
-if IS_HF_SPACE:
-    import spaces  # type: ignore
+IS_HF_SPACE = bool(os.getenv("SPACE_ID"))
 
-    GPU = spaces.GPU
+if IS_HF_SPACE:
+    import spaces as _spaces
+    GPU = _spaces.GPU
 else:
     def GPU(*args, **kwargs):
         def decorator(fn):