[Admin maintenance] Support new ZeroGPU hardware

app.py

@@ -13,15 +13,57 @@ from einops import rearrange, repeat
 from tqdm import tqdm
 from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
 
-import subprocess, sys
+import subprocess, sys, glob, tempfile, ctypes
+
+CUDA_HOME = "/cuda-image/usr/local/cuda-13.0"
+CUDA_LIBDIR = os.path.join(CUDA_HOME, "lib64")
 
 @spaces.GPU(duration=210)
 def first_gpu_setup():
-    subprocess.check_call([
-        sys.executable, "-m", "pip", "install", "--no-build-isolation",
-        "git+https://github.com/NVlabs/nvdiffrast/",
-    ])
+    try:
+        import nvdiffrast  # noqa: F401
+        print("nvdiffrast already installed; skipping build.")
+        return
+    except ImportError:
+        pass
+
+    if not os.path.exists(os.path.join(CUDA_HOME, "bin", "nvcc")):
+        raise RuntimeError(
+            f"nvcc not found at {CUDA_HOME}/bin/nvcc on the ZeroGPU worker. "
+            "The new-hardware CUDA path may have moved; please update CUDA_HOME."
+        )
+
+    patch_dir = tempfile.mkdtemp(prefix="torch_cuda_patch_")
+    with open(os.path.join(patch_dir, "sitecustomize.py"), "w") as f:
+        f.write(
+            "try:\n"
+            "    import torch.utils.cpp_extension as _c\n"
+            "    _c._check_cuda_version = lambda *a, **k: None\n"
+            "except Exception:\n"
+            "    pass\n"
+        )
+
+    env = os.environ.copy()
+    env["CUDA_HOME"] = CUDA_HOME
+    env["CUDA_PATH"] = CUDA_HOME
+    env["PATH"] = os.path.join(CUDA_HOME, "bin") + os.pathsep + env.get("PATH", "")
+    env["PYTHONPATH"] = patch_dir + os.pathsep + env.get("PYTHONPATH", "")
+    env["TORCH_CUDA_ARCH_LIST"] = "12.0"
+
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "install", "--no-deps",
+         "setuptools", "wheel", "ninja"],
+    )
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "install", "--no-build-isolation",
+         "git+https://github.com/NVlabs/nvdiffrast/"],
+        env=env,
+    )
 first_gpu_setup()
+
+ctypes.CDLL(os.path.join(CUDA_LIBDIR, "libcudart.so.13"), mode=ctypes.RTLD_GLOBAL)
+os.environ["LD_LIBRARY_PATH"] = CUDA_LIBDIR + os.pathsep + os.environ.get("LD_LIBRARY_PATH", "")
+
 from src.utils.train_util import instantiate_from_config
 from src.utils.camera_util import (
     FOV_to_intrinsics,