Upload 2 files

post_install.py

@@ -325,18 +325,23 @@ def install_pytorch3d():
     
     # Try the official PyTorch3D installation command for current CUDA version
     if PYTORCH_VERSION and CUDA_VERSION:
-        # Try CUDA 11.7 wheels first (compatible with PyTorch 2.0.1)
-        if CUDA_VERSION == "cu117":
-            print("Trying PyTorch3D installation for CUDA 11.7...")
-            if run_command("pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu117_pyt201/download.html"):
-                print("Successfully installed PyTorch3D for CUDA 11.7")
-                return True
-        # Try CUDA 12.6 wheels for newer PyTorch versions
-        elif CUDA_VERSION == "cu126":
+        # Try CUDA 12.6 wheels first (current system CUDA)
+        if CUDA_VERSION == "cu126":
             print("Trying PyTorch3D installation for CUDA 12.6...")
             if run_command("pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu126_pyt271/download.html"):
                 print("Successfully installed PyTorch3D for CUDA 12.6")
                 return True
+        # Try CUDA 11.7 wheels for older PyTorch versions
+        elif CUDA_VERSION == "cu117":
+            print("Trying PyTorch3D installation for CUDA 11.7...")
+            if run_command("pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu117_pyt201/download.html"):
+                print("Successfully installed PyTorch3D for CUDA 11.7")
+                return True
+    
+    # Try to build PyTorch3D from source with system CUDA
+    print("Trying PyTorch3D source installation with system CUDA...")
+    if install_pytorch3d_from_source():
+        return True
     
     # Fallback to CUDA 11.7 wheels
     print("Trying official PyTorch3D installation (CUDA 11.7)...")
@@ -411,6 +416,38 @@ def install_pytorch3d():
     
     return False
 
+def install_pytorch3d_from_source():
+    """Install PyTorch3D from source with system CUDA"""
+    print("Installing PyTorch3D from source with system CUDA...")
+    
+    packages_dir = Path("packages")
+    
+    # Clone PyTorch3D
+    if not (packages_dir / "pytorch3d").exists():
+        if not run_command("git clone https://github.com/facebookresearch/pytorch3d.git", cwd=packages_dir):
+            return False
+    
+    # Install build dependencies
+    print("Installing build dependencies...")
+    run_command("pip install wheel setuptools ninja")
+    
+    # Install PyTorch3D with system CUDA
+    pytorch3d_dir = packages_dir / "pytorch3d"
+    
+    # Set environment variables to use system CUDA
+    env = os.environ.copy()
+    env['FORCE_CUDA'] = '1'
+    env['CUDA_HOME'] = '/usr/local/cuda'
+    env['CUDA_VERSION'] = '12.6'
+    env['TORCH_CUDA_ARCH_LIST'] = '8.9'  # For L4 GPU
+    
+    print("Building PyTorch3D with system CUDA...")
+    if run_command("pip install . --no-build-isolation", cwd=pytorch3d_dir, env=env):
+        print("Successfully installed PyTorch3D from source with system CUDA")
+        return True
+    
+    return False
+
 def install_pytorch_dependencies():
     """Install PyTorch-related dependencies"""
     print("Installing PyTorch dependencies...")