Fix torch.compile AssertionError in ZeroGPU workers on torch 2.13

Rebuilding the Space picked up torch 2.13.0 from the ZeroGPU base image.
In a GPU worker, spaces 0.51.3 restores torch.cuda._maybe_exchange_device
to the _exchange_device builtin; torch 2.13's dynamo registers one handler
per function and fails on the duplicate ("Handler already registered for
<built-in function _cuda_exchangeDevice>"), so every viewer session died
during renderer warm-up.

Restore the real builtin at the start of both @spaces.GPU functions,
before anything compiles. No-op when the two are already distinct.

Co-Authored-By: Claude Fable 5.2 <noreply@anthropic.com>

app.py

@@ -56,6 +56,7 @@ import gradio as gr
 from viser_app.scene_loaders import (
     IMAGE_EXTS, VIDEO_EXTS, extract_video_frames,
 )
+from viser_app.zerogpu_compat import fix_cuda_exchange_device
 
 # ---------------------------------------------------------------------------
 # Config
@@ -151,6 +152,7 @@ def reconstruct_gpu(img_dir: str, target_h: int, target_w: int, max_views: int)
     from viser_app.pi3_estimator import _list_images, pi3_to_scene_data, run_pi3
     from viser_app.recon_scene import reconstruct, save_artifact
 
+    fix_cuda_exchange_device()  # torch.compile vs ZeroGPU, see zerogpu_compat
     image_paths = _list_images(Path(img_dir))
     if max_views and len(image_paths) > max_views:
         idxs = np.linspace(0, len(image_paths) - 1, max_views).round().astype(int)
@@ -202,6 +204,7 @@ def _live_session_worker(artifact_path: str):
     from viser_app.viewer import install_viser_handlers
     from viser_app.viewer_optim import install_optimizations
 
+    fix_cuda_exchange_device()  # torch.compile vs ZeroGPU, see zerogpu_compat
     yield ("prep", 0.15, "Loading scene…")
     model = DVSM_MODEL_CPU.cuda().eval()
     artifact = torch.load(artifact_path, weights_only=False, map_location="cpu")

viser_app/zerogpu_compat.py

@@ -0,0 +1,28 @@
+"""Workarounds for the ZeroGPU (`spaces`) runtime."""
+
+
+def fix_cuda_exchange_device() -> bool:
+    """Make torch.compile usable in a ZeroGPU worker on torch >= 2.13.
+
+    When a GPU worker starts, spaces.zero.torch.patching.unpatch() (seen in
+    spaces 0.51.3) restores torch.cuda._maybe_exchange_device to the
+    *_exchange_device* builtin, so both names point at one function. torch 2.13's
+    dynamo registers one handler per function and dies on the duplicate with
+    "AssertionError: Handler already registered for <built-in function
+    _cuda_exchangeDevice>" the first time a compiled function is traced.
+
+    Call this inside the @spaces.GPU function before anything compiles. It is a
+    no-op when the two functions are already distinct. Returns True if it had to
+    repair the alias.
+    """
+    import torch
+
+    exchange = getattr(torch.cuda, "_exchange_device", None)
+    if exchange is None or getattr(torch.cuda, "_maybe_exchange_device", None) is not exchange:
+        return False
+    real = getattr(torch._C, "_cuda_maybeExchangeDevice", None)
+    if real is None or real is exchange:
+        def real(device):  # distinct object, same behaviour as the alias
+            return exchange(device)
+    torch.cuda._maybe_exchange_device = real
+    return True