minor chnages in app

app.py

@@ -55,12 +55,21 @@ _VIEWER_MAX_SIDE = 1024
 _PAYLOAD_SLOT = "__DW_EMBED_JSON__"
 
 
-def _load_model_once(device: torch.device):
+def _load_model_once():
+    """Build the model **on the CPU**, once, outside any GPU window.
+
+    On ZeroGPU the GPU is attached only for the duration of a `@spaces.GPU`
+    call and the whole reservation is billed against the caller's daily quota,
+    so reading a checkpoint off /data and pulling the DINOv3 encoder from the
+    hub must not happen in there.  CUDA also cannot be initialised in the main
+    process, hence `torch.device("cpu")` here and `.to("cuda")` inside the
+    worker.
+    """
     global MODEL_PATH, MODEL, MODEL_SPEC, MODEL_CFG
     from infer.predict import load_model
 
     with MODEL_LOCK:
-        if MODEL is None or next(MODEL.parameters()).device != device:
+        if MODEL is None:
             if not MODEL_PATH_ON_BUCKET.is_file():
                 raise RuntimeError(
                     f"Model checkpoint not found at {MODEL_PATH_ON_BUCKET}. "
@@ -74,7 +83,7 @@ def _load_model_once(device: torch.device):
                 )
             MODEL_PATH = str(MODEL_PATH_ON_BUCKET)
             MODEL, MODEL_SPEC, MODEL_CFG = load_model(
-                MODEL_PATH, device, hf_token
+                MODEL_PATH, torch.device("cpu"), hf_token
             )
     return MODEL, MODEL_SPEC, MODEL_CFG
 
@@ -255,32 +264,69 @@ def _bundle(out_dir: Path, files: list[str]) -> list[str]:
 
 
 # ---------------------------------------------------------------------
-# TTA is roughly 4x the work, and a 1600 px scene tiles into a lot of windows.
-@spaces.GPU(duration=180)
+# ZeroGPU bookkeeping
+# ---------------------------------------------------------------------
+# A `@spaces.GPU(duration=N)` call reserves N seconds against the caller's
+# daily GPU allowance *up front*: if their remaining quota is below N the
+# request is refused with "GPU time exhausted" no matter how many GPUs are
+# idle.  The old code asked for 180 s around the whole handler — checkpoint
+# load, hub download, meshing, PNG encoding and zipping included — which put
+# the demo out of reach of anonymous visitors (~60 s) and burned a signed-in
+# user's quota on work that never touches the GPU.
+#
+# So: only the forward pass is decorated, the durations are sized to it, and
+# TTA (~4x the work) gets its own longer window instead of everyone paying for
+# the worst case.
+_GPU_DURATION = 45
+_GPU_DURATION_TTA = 120
+
+
+def _run_model(rgb, gsd_m, tta: bool):
+    """The only GPU-bound step.  Runs inside a ZeroGPU reservation."""
+    from infer.engine import predict_scene
+
+    model, spec, cfg = _load_model_once()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device).eval()
+    amp_dtype = (
+        (torch.bfloat16 if cfg.amp_dtype == "bf16" else torch.float16)
+        if device.type == "cuda"
+        else None
+    )
+    try:
+        height, _ = predict_scene(
+            model, rgb, gsd_m, spec, device,
+            tta=bool(tta), amp_dtype=amp_dtype, overlap=0.5, batch_tiles=4,
+        )
+    finally:
+        # hand the card back clean; the next reservation is a fresh attach
+        model.to("cpu")
+        if device.type == "cuda":
+            torch.cuda.empty_cache()
+    return height
+
+
+_run_model_gpu = spaces.GPU(duration=_GPU_DURATION)(_run_model)
+_run_model_gpu_tta = spaces.GPU(duration=_GPU_DURATION_TTA)(_run_model)
+
+
 def predict(image_path: str | None, gsd: float, tta: bool = False):
     if not image_path:
         raise gr.Error("Upload an image first.")
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     try:
         from dwdata.preprocess import read_scene
-        from infer.engine import predict_scene
 
-        model, spec, cfg = _load_model_once(device)
+        # CPU work, deliberately outside the reservation
+        _, spec, _ = _load_model_once()
         rgb, scene_meta = read_scene(
             image_path,
             user_gsd_m=float(gsd or 0),
             assumed_gsd_m=spec.canonical_gsd_m,
             max_side=1600,
         )
-        amp_dtype = (
-            (torch.bfloat16 if cfg.amp_dtype == "bf16" else torch.float16)
-            if device.type == "cuda"
-            else None
-        )
-        height, _ = predict_scene(
-            model, rgb, scene_meta.gsd_m, spec, device,
-            tta=bool(tta), amp_dtype=amp_dtype, overlap=0.5, batch_tiles=4,
-        )
+
+        runner = _run_model_gpu_tta if tta else _run_model_gpu
+        height = runner(rgb, scene_meta.gsd_m, bool(tta))
 
         height = _sanitise(height)
         low, high = _height_bounds(height)
@@ -341,7 +387,18 @@ def predict(image_path: str | None, gsd: float, tta: bool = False):
             (str(out_dir / "hillshade.png"), "Hillshade (315° / 45°)"),
         ]
         return gallery, height16, viewer_html, _bundle(out_dir, files), info
+    except gr.Error:
+        raise
     except Exception as exc:
+        text = str(exc).lower()
+        if "quota" in text or "gpu task aborted" in text or "gpu time" in text:
+            # not a model failure: the caller's ZeroGPU allowance ran out
+            raise gr.Error(
+                "Your ZeroGPU quota for today is used up. Sign in to Hugging "
+                "Face for a larger allowance, turn TTA off (it asks for a "
+                f"{_GPU_DURATION_TTA} s window instead of {_GPU_DURATION} s), "
+                "or retry after the quota resets."
+            ) from exc
         raise gr.Error(f"Inference failed: {exc}") from exc
 
 
@@ -389,5 +446,15 @@ with gr.Blocks(title="Single View Height Estimator") as demo:
     )
 
 
+# Build the model on the CPU while the Space boots, so the first visitor's GPU
+# reservation pays for the forward pass only and not for the checkpoint read
+# and the encoder download.  A failure here (no bucket, no HF_TOKEN) is
+# reported per-request as before, so it must not take the Space down.
+try:
+    _load_model_once()
+except Exception as _warm_exc:  # pragma: no cover - startup convenience only
+    print(f"[app] model warm-up deferred: {_warm_exc}")
+
+
 if __name__ == "__main__":
     demo.launch()