[Admin maintenance] Support new ZeroGPU hardware

networks/op/_cudart_preload.py

@@ -0,0 +1,36 @@
+"""Preload libcudart.so.13 with RTLD_GLOBAL.
+
+The ZeroGPU (Blackwell) image's nvcc is CUDA 13, so from-source JIT-compiled
+CUDA extensions (built via torch.utils.cpp_extension.load) link against
+libcudart.so.13. torch's own pip wheel (cu12x) does not ship that version,
+so the compiled extension fails to dlopen with:
+    libcudart.so.13: cannot open shared object file: No such file or directory
+`nvidia-cuda-runtime>=13,<14` (in requirements.txt) provides the .so inside
+its package dir, but it isn't on the default loader search path — preload it
+explicitly before any extension that needs it is imported.
+"""
+import ctypes
+import glob
+import os
+import sys
+
+
+def preload_libcudart():
+    patterns = []
+    for p in sys.path:
+        patterns.append(os.path.join(p, "nvidia", "**", "libcudart.so.13*"))
+    patterns += [
+        "/usr/local/cuda*/targets/*/lib/libcudart.so.13*",
+        "/usr/local/cuda*/lib64/libcudart.so.13*",
+    ]
+    for pattern in patterns:
+        for path in glob.glob(pattern, recursive=True):
+            try:
+                ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)
+                return path
+            except OSError:
+                continue
+    return None
+
+
+preload_libcudart()

networks/op/fused_act.py

@@ -6,18 +6,21 @@ from torch.nn import functional as F
 from torch.autograd import Function
 from torch.utils.cpp_extension import load
 
-
-if os.getenv("SPACE_ID"):
-    import networks.op.fused as fused
-else:
-    module_path = os.path.dirname(__file__)
-    fused = load(
-        "fused",
-        sources=[
-            os.path.join(module_path, "fused_bias_act.cpp"),
-            os.path.join(module_path, "fused_bias_act_kernel.cu"),
-        ],
-    )
+from . import _cudart_preload  # noqa: F401 - preloads libcudart.so.13 for the JIT build below
+
+# The prebuilt networks/op/fused.so was compiled against an old torch C++ ABI
+# and fails with "undefined symbol" on newer torch / Blackwell (sm_120)
+# runtimes. Always rebuild the extension from source at runtime instead of
+# importing the stale binary.
+os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
+module_path = os.path.dirname(__file__)
+fused = load(
+    "fused",
+    sources=[
+        os.path.join(module_path, "fused_bias_act.cpp"),
+        os.path.join(module_path, "fused_bias_act_kernel.cu"),
+    ],
+)
 
 
 class FusedLeakyReLUFunctionBackward(Function):

networks/op/upfirdn2d.py

@@ -6,18 +6,21 @@ from torch.nn import functional as F
 from torch.autograd import Function
 from torch.utils.cpp_extension import load
 
-
-if os.getenv("SPACE_ID"):
-    import networks.op.upfirdn2d_op as upfirdn2d_op
-else:
-    module_path = os.path.dirname(__file__)
-    upfirdn2d_op = load(
-        "upfirdn2d",
-        sources=[
-            os.path.join(module_path, "upfirdn2d.cpp"),
-            os.path.join(module_path, "upfirdn2d_kernel.cu"),
-        ],
-    )
+from . import _cudart_preload  # noqa: F401 - preloads libcudart.so.13 for the JIT build below
+
+# The prebuilt networks/op/upfirdn2d_op.so was compiled against an old torch
+# C++ ABI and fails with "undefined symbol" on newer torch / Blackwell
+# (sm_120) runtimes. Always rebuild the extension from source at runtime
+# instead of importing the stale binary.
+os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
+module_path = os.path.dirname(__file__)
+upfirdn2d_op = load(
+    "upfirdn2d",
+    sources=[
+        os.path.join(module_path, "upfirdn2d.cpp"),
+        os.path.join(module_path, "upfirdn2d_kernel.cu"),
+    ],
+)
 
 
 class UpFirDn2dBackward(Function):

requirements.txt

@@ -4,4 +4,5 @@ imageio==2.37.0
 imageio-ffmpeg==0.6.0
 ninja==1.11.1.4
 torch==2.8.0
-torchvision==0.23
+torchvision==0.23
+nvidia-cuda-runtime>=13,<14