Revert cu130 forcing; use ZeroGPU-native cu128 torch + pip cu128 nvcc

Root cause found via [GPU-DIAG]: 42 GiB free yet allocation fails --
not a real OOM. The forced cu130 torch is incompatible with ZeroGPU's
fork-based GPU memory management (ZeroGPU targets the cu128 builds of
its supported torch versions). Stop forcing cu130: drop the startup
torch reinstall and the allocator-env band-aids, and run the cu128
torch the base image provides. diff-gaussian then needs a CUDA 12.8
compiler (container only has 13.0), so install nvidia-cuda-nvcc-cu12 /
cccl and point the build's CUDA_HOME at it (arch list incl. sm_120).
Net: fewer hacks than before, aligned with ZeroGPU's supported stack.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

app.py

@@ -22,59 +22,15 @@ if os.environ.get("HF_HUB_ENABLE_HF_TRANSFER") == "1":
         print("⚠️  hf_transfer not available, disabling fast download")
         os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
 
-# ZeroGPU attaches the GPU inside a forked subprocess where NVML can't reinit.
-# PyTorch's expandable_segments allocator uses CUDA VMM, whose NVML query then
-# asserts ("NVML_SUCCESS == r" in CUDACachingAllocator.cpp). Disabling ONLY
-# expandable segments avoids that NVML path while keeping the normal caching
-# allocator (so memory is still pooled/reused -- disabling caching entirely OOMs).
-# Must be set before torch initializes CUDA, i.e. before "import torch".
-os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:False")
-
 import json
 from pathlib import Path
 from datetime import datetime
 import random
 
-# -----------------------------
-# Ensure a CUDA-13.0 (cu130) torch BEFORE torch is imported.
-# The ZeroGPU container ships the CUDA 13.0 toolkit, but the image's torch is a
-# cu128 build (torch.version.cuda == "12.8"), so the runtime compile of
-# diff-gaussian-rasterization fails with a 13.0-vs-12.8 mismatch. Build-time pip
-# kept resolving the cu128 wheel regardless of index config, so fix it here: if
-# the installed torch isn't a +cu130 build, reinstall the cu130 stack. This runs
-# before "import torch" so the whole process and the build subprocess use cu130.
-def _ensure_cu130_torch():
-    import subprocess as _sp
-    import sys as _sys
-    import importlib.metadata as _md
-    try:
-        installed = _md.version("torch")
-    except _md.PackageNotFoundError:
-        installed = ""
-    print(f"[CU130-CHECK] installed torch == {installed or '<none>'}")
-    if "+cu130" in installed:
-        print("[CU130-CHECK] cu130 torch already present, skipping reinstall")
-        return
-    print("[CU130-CHECK] reinstalling torch/torchvision/xformers from the cu130 index...")
-    # NOT --no-deps: torch 2.11.0+cu130 must pull its matching nvidia-*-cu13 CUDA
-    # runtime libs. With --no-deps the container kept the base image's nvidia-*-cu12
-    # (12.8) libs, so cu130 torch ran on cu12 runtime -- a broken mix that made
-    # allocations fail (surfacing as the NVML/OOM assert). --extra-index-url lets
-    # torch's non-NVIDIA deps still resolve from PyPI.
-    _sp.check_call([
-        _sys.executable, "-m", "pip", "install", "--no-cache-dir",
-        "--index-url", "https://download.pytorch.org/whl/cu130",
-        "--extra-index-url", "https://pypi.org/simple",
-        "torch==2.11.0+cu130", "torchvision==0.26.0+cu130", "xformers==0.0.35",
-    ])
-    print("[CU130-CHECK] cu130 stack installed")
-
-_ensure_cu130_torch()
-
 import gradio as gr
 import numpy as np
 import torch
-print(f"[CU130-CHECK] torch {torch.__version__} / CUDA {torch.version.cuda}")
+print(f"[GPU-DIAG] torch {torch.__version__} / CUDA {torch.version.cuda}")
 import yaml
 from easydict import EasyDict as edict
 from einops import rearrange
@@ -142,20 +98,31 @@ try:
 except ImportError:
     print("Installing diff-gaussian-rasterization (compiling for detected CUDA arch)...")
     env = os.environ.copy()
-    try:
-        import torch as _torch
-        if _torch.cuda.is_available():
-            maj, minr = _torch.cuda.get_device_capability()
-            arch = f"{maj}.{minr}"                 # e.g., "9.0" on H100/H200, "8.0" on A100
-            env["TORCH_CUDA_ARCH_LIST"] = f"{arch}+PTX"
-        else:
-            # Build stage may not see a GPU on HF Spaces: compile a cross-arch set
-            env["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6;8.9;9.0+PTX"
-    except Exception:
-        env["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6;8.9;9.0+PTX"
 
-    # (Optional) side-step allocator+NVML quirks in restrictive containers
-    env.setdefault("PYTORCH_NO_CUDA_MEMORY_CACHING", "1")
+    # Compile against the pip-installed CUDA 12.8 toolkit (nvidia-cuda-nvcc-cu12),
+    # NOT the container's CUDA 13.0 toolkit. torch here is the cu128 build ZeroGPU
+    # requires, and torch's build check errors on a major CUDA mismatch (nvcc 13 vs
+    # torch 12.8). Point CUDA_HOME at the pip nvcc and feed the runtime/cub headers
+    # and cudart lib via CPATH/LIBRARY_PATH so nvcc 12.8 matches torch 12.8.
+    import site
+    _site = site.getsitepackages()[0]
+    _nv = os.path.join(_site, "nvidia")
+    _nvcc_home = os.path.join(_nv, "cuda_nvcc")
+    if os.path.isdir(_nvcc_home):
+        env["CUDA_HOME"] = _nvcc_home
+        env["PATH"] = os.path.join(_nvcc_home, "bin") + os.pathsep + env.get("PATH", "")
+        _incs = [os.path.join(_nv, "cuda_runtime", "include"),
+                 os.path.join(_nv, "cuda_cccl", "include")]
+        env["CPATH"] = os.pathsep.join(_incs + ([env["CPATH"]] if env.get("CPATH") else []))
+        _rt_lib = os.path.join(_nv, "cuda_runtime", "lib")
+        env["LIBRARY_PATH"] = _rt_lib + os.pathsep + env.get("LIBRARY_PATH", "")
+        env["LD_LIBRARY_PATH"] = _rt_lib + os.pathsep + env.get("LD_LIBRARY_PATH", "")
+    else:
+        print(f"⚠️  pip CUDA toolkit not found at {_nvcc_home}; falling back to system nvcc")
+
+    # Build a cross-arch set incl. Blackwell sm_120 (the ZeroGPU GPU is cap 12.0);
+    # the build stage may not see a GPU, so don't rely on get_device_capability.
+    env["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6;8.9;9.0;12.0+PTX"
 
     # --no-build-isolation: the package's setup.py imports torch to detect the
     # CUDA arch, but pip's isolated build env has no torch. Reuse the main env

requirements.txt

@@ -1,11 +1,13 @@
-# torch 2.11.0 is ZeroGPU-supported (bare version for the config validator) and is
-# the version line where torchvision 0.26.0 and xformers 0.0.35 agree (xformers
-# 0.0.33 hard-pinned torch==2.9.0, which ZeroGPU rejects; 0.0.35 relaxed to >=2.10).
-# The CUDA build matters: the ZeroGPU container ships CUDA 13.0, so torch must be
-# cu130 for diff-gaussian-rasterization to compile. Build-time pip resolves the
-# cu128 wheel here, so the cu130 reinstall is forced at startup in app.py.
+# torch 2.11.0 is ZeroGPU-supported and the version line where torchvision 0.26.0
+# and xformers 0.0.35 agree (xformers 0.0.33 hard-pinned torch==2.9.0, which ZeroGPU
+# rejects; 0.0.35 relaxed to >=2.10). This is the cu128 (CUDA 12.8) build ZeroGPU
+# ships -- required for ZeroGPU's GPU memory management to work.
 torch==2.11.0
 torchvision==0.26.0
+# CUDA 12.8 compiler + headers so diff-gaussian-rasterization compiles against a
+# toolkit matching cu128 torch (the container only ships CUDA 13.0). See app.py.
+nvidia-cuda-nvcc-cu12~=12.8.0
+nvidia-cuda-cccl-cu12~=12.8.0
 numpy==1.26.4
 pillow==10.4.0
 diffusers[torch]==0.30.3