Include PyTorch CUDA headers for nvdiffrast JIT

app.py

@@ -75,6 +75,29 @@ def _find_site_package_dir(relative_path: str) -> Optional[Path]:
     return None
 
 
+def _find_site_package_dirs(relative_path: str) -> List[Path]:
+    import site
+
+    roots = []
+    for getter in (site.getsitepackages, lambda: [site.getusersitepackages()]):
+        try:
+            roots.extend(getter())
+        except Exception:
+            pass
+    try:
+        roots.append(str(Path(torch.__file__).resolve().parents[1]))
+    except Exception:
+        pass
+
+    found = []
+    for root in dict.fromkeys(roots):
+        parent = Path(root) / relative_path
+        if not parent.exists():
+            continue
+        found.extend(sorted(parent.iterdir()))
+    return [path for path in found if path.is_dir()]
+
+
 def _symlink_children(src_dir: Optional[Path], dst_dir: Path) -> None:
     if src_dir is None or not src_dir.exists():
         return
@@ -97,6 +120,18 @@ def _configure_cuda_extension_runtime() -> None:
     runtime_include = cuda_runtime_root / 'include' if cuda_runtime_root else None
     nvcc_include = cuda_nvcc_root / 'include' if cuda_nvcc_root else None
     nvcc_bin = cuda_nvcc_root / 'bin' if cuda_nvcc_root else None
+    nvidia_package_dirs = _find_site_package_dirs('nvidia')
+    nvidia_include_dirs = [
+        path / 'include'
+        for path in nvidia_package_dirs
+        if (path / 'include').exists()
+    ]
+    nvidia_lib_dirs = [
+        lib_dir
+        for path in nvidia_package_dirs
+        for lib_dir in (path / 'lib', path / 'lib64')
+        if lib_dir.exists()
+    ]
 
     if not runtime_lib or not (runtime_lib / 'libcudart.so.12').exists():
         print('[app_hf] CUDA 12 runtime package not found; using system CUDA_HOME.', flush=True)
@@ -125,16 +160,20 @@ def _configure_cuda_extension_runtime() -> None:
 
     _symlink_children(runtime_include, compat_include)
     _symlink_children(nvcc_include, compat_include)
-    _symlink_children(runtime_lib, compat_lib)
+    for include_dir in nvidia_include_dirs:
+        _symlink_children(include_dir, compat_include)
+    for lib_dir in nvidia_lib_dirs:
+        _symlink_children(lib_dir, compat_lib)
     libcudart_link = compat_lib / 'libcudart.so'
     if not libcudart_link.exists() and not libcudart_link.is_symlink():
         libcudart_link.symlink_to(runtime_lib / 'libcudart.so.12')
 
     os.environ['CUDA_HOME'] = str(compat_root)
     os.environ['CUDA_PATH'] = str(compat_root)
+    os.environ.setdefault('TORCH_EXTENSIONS_DIR', '/tmp/anigen_torch_extensions_cuda12')
     _prepend_env_path('PATH', [compat_bin, nvcc_bin])
-    _prepend_env_path('LIBRARY_PATH', [compat_lib, runtime_lib])
-    _prepend_env_path('LD_LIBRARY_PATH', [compat_lib, runtime_lib])
+    _prepend_env_path('LIBRARY_PATH', [compat_lib] + nvidia_lib_dirs)
+    _prepend_env_path('LD_LIBRARY_PATH', [compat_lib] + nvidia_lib_dirs)
 
     try:
         import ctypes