Upload post_install.py

post_install.py

@@ -10,6 +10,10 @@ import subprocess
 import shutil
 from pathlib import Path
 
+# Global variables for PyTorch version detection
+PYTORCH_VERSION = None
+CUDA_VERSION = None
+
 def run_command(command, cwd=None, env=None):
     """Run a shell command and return the result"""
     print(f"Running: {command}")
@@ -47,27 +51,32 @@ def install_torch_sparse():
     except ImportError:
         pass
     
-    # Check if PyTorch is already installed with the correct version
+    # Check current PyTorch version and adapt to it
     try:
         import torch
         print(f"Current PyTorch version: {torch.__version__}")
         
-        # Check if we need to update PyTorch
-        if torch.__version__.startswith("2.0.1") and "+cu117" in torch.__version__:
-            print(f"PyTorch {torch.__version__} already installed with correct CUDA version")
+        # Extract PyTorch version info for compatibility
+        if "+" in torch.__version__:
+            pytorch_base = torch.__version__.split("+")[0]
+            cuda_version = torch.__version__.split("+")[1]
         else:
-            print(f"PyTorch {torch.__version__} installed, but need to update to 2.0.1+cu117")
-            # Uninstall current PyTorch and reinstall with correct version
-            print("Uninstalling current PyTorch...")
-            run_command("pip uninstall torch torchvision torchaudio -y")
-            print("Installing compatible PyTorch version...")
-            if not run_command("pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu117"):
-                return False
+            pytorch_base = torch.__version__
+            cuda_version = "cpu"
+        
+        print(f"PyTorch base version: {pytorch_base}, CUDA: {cuda_version}")
+        
+        # Store version info for later use
+        global PYTORCH_VERSION, CUDA_VERSION
+        PYTORCH_VERSION = pytorch_base
+        CUDA_VERSION = cuda_version
+        
     except ImportError:
-        # First, install a compatible PyTorch version with CUDA 11.7 (as expected by PyTorch3D)
-        print("Installing compatible PyTorch version...")
+        print("PyTorch not found, installing default version...")
         if not run_command("pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu117"):
             return False
+        PYTORCH_VERSION = "2.0.1"
+        CUDA_VERSION = "cu117"
     
     # Check PyTorch installation
     print("Checking PyTorch installation...")
@@ -83,8 +92,13 @@ def install_torch_sparse():
         print(f"Error checking PyTorch version: {e}")
     
     # Now install torch-sparse with the compatible version
-    print("Installing torch-sparse with PyTorch 2.0.1...")
-    if run_command("pip install torch-sparse -f https://data.pyg.org/whl/torch-2.0.1+cu117.html"):
+    print(f"Installing torch-sparse with PyTorch {PYTORCH_VERSION}...")
+    
+    # Try to find compatible torch-sparse wheel
+    wheel_url = f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+{CUDA_VERSION}.html"
+    print(f"Trying wheel URL: {wheel_url}")
+    
+    if run_command(f"pip install torch-sparse -f {wheel_url}"):
         print("Successfully installed torch-sparse")
         
         # Verify torch-sparse is compatible
@@ -97,6 +111,22 @@ def install_torch_sparse():
         
         return True
     
+    # If the specific wheel fails, try alternative versions
+    print("Trying alternative torch-sparse versions...")
+    alternative_versions = [
+        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+{CUDA_VERSION}.html",
+        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+cu118.html",
+        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+cu117.html",
+        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+cu116.html"
+    ]
+    
+    for url in alternative_versions:
+        if url != wheel_url:  # Skip the one we already tried
+            print(f"Trying alternative URL: {url}")
+            if run_command(f"pip install torch-sparse -f {url}"):
+                print("Successfully installed torch-sparse with alternative version")
+                return True
+    
     return False
 
 def install_torch_scatter():
@@ -112,8 +142,13 @@ def install_torch_scatter():
         pass
     
     # Install torch-scatter with the compatible PyTorch version
-    print("Installing torch-scatter with PyTorch 2.0.1...")
-    if run_command("pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.1+cu117.html"):
+    print(f"Installing torch-scatter with PyTorch {PYTORCH_VERSION}...")
+    
+    # Try to find compatible torch-scatter wheel
+    wheel_url = f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+{CUDA_VERSION}.html"
+    print(f"Trying wheel URL: {wheel_url}")
+    
+    if run_command(f"pip install torch-scatter -f {wheel_url}"):
         print("Successfully installed torch-scatter")
         
         # Verify torch-scatter is compatible
@@ -126,6 +161,22 @@ def install_torch_scatter():
         
         return True
     
+    # If the specific wheel fails, try alternative versions
+    print("Trying alternative torch-scatter versions...")
+    alternative_versions = [
+        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+{CUDA_VERSION}.html",
+        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+cu118.html",
+        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+cu117.html",
+        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+cu116.html"
+    ]
+    
+    for url in alternative_versions:
+        if url != wheel_url:  # Skip the one we already tried
+            print(f"Trying alternative URL: {url}")
+            if run_command(f"pip install torch-scatter -f {url}"):
+                print("Successfully installed torch-scatter with alternative version")
+                return True
+    
     return False
 
 def install_nvdiffrast():
@@ -187,35 +238,21 @@ def install_pytorch3d():
         return True
     
     # Try different wheel URLs for different Python/PyTorch versions
-    # First, try to determine the correct wheel URL based on current PyTorch version
-    try:
-        import torch
-        torch_version = torch.__version__
-        print(f"Determining PyTorch3D wheel URL for PyTorch {torch_version}")
-        
-        # Extract PyTorch major.minor version
-        if "+" in torch_version:
-            pytorch_base = torch_version.split("+")[0]
-        else:
-            pytorch_base = torch_version
-        
-        # Extract CUDA version
-        cuda_version = "cu117"  # default
-        if "+cu" in torch_version:
-            cuda_version = torch_version.split("+")[1]
-        
-        # Try specific wheel URL for current PyTorch version
-        specific_url = f"https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_{cuda_version}_pyt{pytorch_base.replace('.', '')}/download.html"
-        print(f"Trying specific wheel URL: {specific_url}")
-        if run_command(f"pip install pytorch3d -f {specific_url}"):
-            print(f"Successfully installed PyTorch3D with specific wheel URL")
-            return True
-    except Exception as e:
-        print(f"Error determining specific wheel URL: {e}")
+    # Use the detected PyTorch version
+    print(f"Determining PyTorch3D wheel URL for PyTorch {PYTORCH_VERSION}+{CUDA_VERSION}")
+    
+    # Try specific wheel URL for current PyTorch version
+    specific_url = f"https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_{CUDA_VERSION}_pyt{PYTORCH_VERSION.replace('.', '')}/download.html"
+    print(f"Trying specific wheel URL: {specific_url}")
+    if run_command(f"pip install pytorch3d -f {specific_url}"):
+        print(f"Successfully installed PyTorch3D with specific wheel URL")
+        return True
     
     # Fallback to known working wheel URLs
     wheel_urls = [
+        f"https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_{CUDA_VERSION}_pyt{PYTORCH_VERSION.replace('.', '')}/download.html",
         "https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu117_pyt201/download.html",
+        "https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu126_pyt271/download.html",
         "https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py39_cu117_pyt201/download.html",
         "https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py38_cu117_pyt201/download.html"
     ]
@@ -316,8 +353,9 @@ def main():
         print(f"✓ PyTorch {torch.__version__} - CUDA: {torch.cuda.is_available()}")
         
         # Check if PyTorch version is compatible
-        if not torch.__version__.startswith("2.0.1") or "+cu117" not in torch.__version__:
+        if PYTORCH_VERSION and not torch.__version__.startswith(PYTORCH_VERSION):
             print(f"⚠ Warning: PyTorch version {torch.__version__} may not be compatible with installed extensions")
+            print(f"Expected version: {PYTORCH_VERSION}")
         
         import torch_sparse
         print("✓ torch-sparse")