Fix nvdiffrast SM mismatch: add +PTX arch + GL context fallback

cudaErrorNoKernelImageForDevice (209) means the compiled kernels don't
match the actual GPU (likely Blackwell SM 10.0 on ZeroGPU's CUDA 13.0
upgrade). Fix:
- TORCH_CUDA_ARCH_LIST="8.0;8.9;9.0+PTX": +PTX lets CUDA 13.0 runtime
JIT-compile for any newer arch (SM 10.0) at first use
- Marker file (/tmp/nvdiffrast_arch.txt) avoids redundant rebuilds on
warm container restarts; force-reinstalls when arch string changes
- All four bake stages fall back to RasterizeGLContext if CUDA context
creation fails

app.py

@@ -53,78 +53,87 @@ import gradio as gr
 import spaces  # HF ZeroGPU
 
 
+# GPU arch list used to compile nvdiffrast.
+# +PTX on the last entry generates portable PTX code that the CUDA runtime
+# JIT-compiles for any newer GPU (e.g. Blackwell SM 10.0) at first use.
+_NVDR_ARCH = "8.0;8.9;9.0+PTX"
+_NVDR_MARKER = Path("/tmp/nvdiffrast_arch.txt")
+
+
 def _install_nvdiffrast() -> None:
     """Install nvdiffrast from source at startup.
 
     Must run here (not requirements.txt) because pip's build isolation hides
     torch from nvdiffrast's build system.
 
-    Problem: torch 2.9.1 is compiled with CUDA 12.8, but the ZeroGPU machine
-    has CUDA 13.0 toolkit. PyTorch's cpp_extension._check_cuda_version compares
-    `nvcc --version` against `torch.version.cuda` and raises RuntimeError on
-    any mismatch.
+    Shim strategy: ZeroGPU machine has CUDA 13.0 toolkit but torch 2.9.1 was
+    compiled with CUDA 12.8. PyTorch's _check_cuda_version raises on mismatch.
+    We put a shim nvcc on PATH that reports 12.8 for --version but delegates
+    real compilation to the actual CUDA 13.0 nvcc.
 
-    Fix: put a shim `nvcc` earlier on PATH that reports torch's CUDA version
-    for `--version` but delegates all real compilation to the actual nvcc 13.0.
-    CUDA 13.0 is backward-compatible and compiles CUDA 12.x code without issues.
-    """
-    try:
-        import nvdiffrast  # noqa: F401
-        return
-    except (ImportError, OSError):
-        pass
+    Arch strategy: compile native kernels for SM 8.0 (A100) and 8.9 (L4/L40),
+    plus PTX for SM 9.0 (H100/H200). The +PTX entry lets the CUDA 13.0 runtime
+    JIT-compile for any newer architecture (e.g. SM 10.0 Blackwell) at first use.
 
+    Marker file: /tmp/nvdiffrast_arch.txt stores the arch string used.  If it
+    matches _NVDR_ARCH, skip rebuild on warm container restarts.
+    """
     import os
     import shutil
     import tempfile
     import torch
 
+    # Check if already installed with the current arch — skip rebuild if so.
+    try:
+        import nvdiffrast  # noqa: F401
+        if _NVDR_MARKER.exists() and _NVDR_MARKER.read_text().strip() == _NVDR_ARCH:
+            print("nvdiffrast: already installed with correct arch.", flush=True)
+            return
+        print("nvdiffrast: arch mismatch or missing marker — forcing reinstall.", flush=True)
+        subprocess.run(
+            [sys.executable, "-m", "pip", "uninstall", "nvdiffrast", "-y"],
+            check=False, capture_output=True,
+        )
+    except (ImportError, OSError):
+        pass
+
     torch_cuda = torch.version.cuda or "12.8"
     print(f"nvdiffrast: torch.version.cuda={torch_cuda}", flush=True)
 
-    # Locate real nvcc — check CUDA_HOME first, then common ZeroGPU paths.
+    # Locate real nvcc.
     real_nvcc: str | None = None
-    cuda_candidates = [
+    real_cuda_home: str | None = None
+    for candidate in [
         os.environ.get("CUDA_HOME", ""),
         "/cuda-image/usr/local/cuda-13.0",
         "/cuda-image/usr/local/cuda-12.8",
         "/cuda-image/usr/local/cuda",
         "/usr/local/cuda-13.0",
-        "/usr/local/cuda-12.8",
         "/usr/local/cuda",
-    ]
-    real_cuda_home: str | None = None
-    for candidate in cuda_candidates:
+    ]:
         if candidate and os.path.isfile(os.path.join(candidate, "bin", "nvcc")):
             real_cuda_home = candidate
             real_nvcc = os.path.join(candidate, "bin", "nvcc")
             break
     if real_nvcc is None:
         real_nvcc = shutil.which("nvcc")
-
     print(f"nvdiffrast: real_nvcc={real_nvcc}", flush=True)
 
     env = os.environ.copy()
 
     if real_nvcc:
-        # Build a fake CUDA home: a shim nvcc that reports torch's CUDA version
-        # on --version (to pass PyTorch's _check_cuda_version) but delegates all
-        # real compilation to the actual nvcc. CUDA 13.0 is backward-compatible
-        # and compiles CUDA 12.x code without issues.
         fake_cuda = Path(tempfile.mkdtemp(prefix="cuda_shim_"))
-        fake_bin = fake_cuda / "bin"
-        fake_bin.mkdir()
+        (fake_cuda / "bin").mkdir()
         if real_cuda_home:
             for sub in ("include", "lib", "lib64"):
                 src = Path(real_cuda_home) / sub
                 if src.exists():
                     (fake_cuda / sub).symlink_to(src)
-        shim = fake_bin / "nvcc"
+        shim = fake_cuda / "bin" / "nvcc"
         shim.write_text(
             "#!/bin/sh\n"
             'case "$1" in\n'
             '  --version)\n'
-            f'    echo "nvcc: NVIDIA (R) Cuda compiler driver"\n'
             f'    echo "Cuda compilation tools, release {torch_cuda}, V{torch_cuda}.0"\n'
             '    ;;\n'
             '  *)\n'
@@ -135,23 +144,17 @@ def _install_nvdiffrast() -> None:
         shim.chmod(0o755)
         env["CUDA_HOME"] = str(fake_cuda)
         env["CUDA_PATH"] = str(fake_cuda)
-        print(f"nvdiffrast: shim active — CUDA_HOME={fake_cuda}", flush=True)
-    else:
-        print("nvdiffrast: nvcc not found — build will likely fail", flush=True)
+        print(f"nvdiffrast: shim CUDA_HOME={fake_cuda}", flush=True)
 
-    # At startup there is no GPU allocated, so PyTorch cannot auto-detect
-    # compute capabilities — _get_cuda_arch_flags returns an empty list and
-    # crashes with IndexError.  Explicitly target A100 (8.0) and H100/H200 (9.0).
-    env.setdefault("TORCH_CUDA_ARCH_LIST", "8.0;9.0")
+    # PTX in the arch list lets the runtime JIT-compile for any newer GPU
+    # (e.g. Blackwell SM 10.0) that lacks a pre-compiled cubin.
+    env["TORCH_CUDA_ARCH_LIST"] = _NVDR_ARCH
 
-    print("Building nvdiffrast from source (~2 min)...", flush=True)
+    print(f"Building nvdiffrast (arch={_NVDR_ARCH}, ~2 min)...", flush=True)
     result = subprocess.run(
         [sys.executable, "-m", "pip", "install", "--no-build-isolation",
          "git+https://github.com/NVlabs/nvdiffrast.git"],
-        env=env,
-        check=False,
-        capture_output=True,
-        text=True,
+        env=env, check=False, capture_output=True, text=True,
     )
     print(result.stdout[-3000:] if result.stdout else "", flush=True)
     if result.returncode != 0:
@@ -159,6 +162,7 @@ def _install_nvdiffrast() -> None:
         print(result.stderr[-2000:] if result.stderr else "", flush=True)
     else:
         print("nvdiffrast build succeeded.", flush=True)
+        _NVDR_MARKER.write_text(_NVDR_ARCH)
 
 
 _install_nvdiffrast()

src/stages/stage2_bake_albedo.py

@@ -66,7 +66,10 @@ def bake_albedo(
     from PIL import Image
 
     device = torch.device("cuda")
-    ctx = dr.RasterizeCudaContext()
+    try:
+        ctx = dr.RasterizeCudaContext()
+    except Exception:
+        ctx = dr.RasterizeGLContext()
 
     hi = trimesh.load(str(high_poly_path), force="mesh")
     lo = trimesh.load(str(unwrapped_glb_path), force="mesh")
@@ -143,7 +146,10 @@ def bake_material(
     from PIL import Image
 
     device = torch.device("cuda")
-    ctx = dr.RasterizeCudaContext()
+    try:
+        ctx = dr.RasterizeCudaContext()
+    except Exception:
+        ctx = dr.RasterizeGLContext()
 
     hi = trimesh.load(str(high_poly_path), force="mesh")
     lo = trimesh.load(str(unwrapped_glb_path), force="mesh")

src/stages/stage2_bake_ao.py

@@ -81,7 +81,10 @@ def bake_ao(
     max_dist = 0.5  # AO search radius (in mesh units)
 
     device = torch.device("cuda")
-    ctx = dr.RasterizeCudaContext()
+    try:
+        ctx = dr.RasterizeCudaContext()
+    except Exception:
+        ctx = dr.RasterizeGLContext()
 
     # Use the unwrapped mesh for both source positions and AO scene
     scene_mesh = trimesh.load(str(mesh_path), force="mesh")

src/stages/stage2_bake_normal.py

@@ -94,7 +94,10 @@ def bake_normal_map(
     from PIL import Image
 
     device = torch.device("cuda")
-    ctx = dr.RasterizeCudaContext()
+    try:
+        ctx = dr.RasterizeCudaContext()
+    except Exception:
+        ctx = dr.RasterizeGLContext()
 
     # ---- Load meshes -------------------------------------------------------
     hi = trimesh.load(str(high_poly_path), force="mesh")