Update setup_blender.sh

setup_blender.sh

@@ -12,12 +12,22 @@ BLENDER_URL="https://download.blender.org/release/Blender${BLENDER_MAJOR_MINOR}/
 INSTALL_DIR="/opt/blender-${BLENDER_VERSION}-linux-x64"
 BLENDER_PY_EXEC="${INSTALL_DIR}/${BLENDER_MAJOR_MINOR}/python/bin/${BLENDER_PYTHON_VERSION}"
 # Assuming unirig_requirements.txt is in the root directory alongside this script and app.py
-# If UniRig repo is cloned first, adjust path e.g., "UniRig/unirig_requirements.txt"
 UNIRIG_REQS_FILE="unirig_requirements.txt"
 # Define the specific torch version and index URL (matching unirig_requirements.txt)
 TORCH_VERSION="2.3.1"
 TORCHVISION_VERSION="0.18.1"
-TORCH_INDEX_URL="https://download.pytorch.org/whl/cu121" # Make sure this matches your target CUDA
+TORCH_CUDA_SUFFIX="cu121" # CUDA version suffix for wheels
+TORCH_INDEX_URL="https://download.pytorch.org/whl/${TORCH_CUDA_SUFFIX}" # Make sure this matches your target CUDA
+
+# PyTorch Geometric wheel URLs (Update these if versions change)
+PYG_TORCH_VERSION_SUFFIX="2.3.1" # PyTorch version suffix in PyG URLs
+PYG_PYTHON_VERSION="cp311"       # Python version suffix
+PYG_CUDA_SUFFIX="cu121"          # CUDA version suffix
+PYG_BASE_URL="https://data.pyg.org/whl/torch-${PYG_TORCH_VERSION_SUFFIX}+${PYG_CUDA_SUFFIX}"
+# Find exact wheel names from the HTML pages if these links break
+TORCH_SCATTER_WHL_URL="${PYG_BASE_URL}/torch_scatter-2.1.2+pt${PYG_TORCH_VERSION_SUFFIX/./}${PYG_CUDA_SUFFIX}-${PYG_PYTHON_VERSION}-${PYG_PYTHON_VERSION}-linux_x86_64.whl"
+TORCH_CLUSTER_WHL_URL="${PYG_BASE_URL}/torch_cluster-1.6.3+pt${PYG_TORCH_VERSION_SUFFIX/./}${PYG_CUDA_SUFFIX}-${PYG_PYTHON_VERSION}-${PYG_PYTHON_VERSION}-linux_x86_64.whl"
+
 
 # --- Download and Extract Blender ---
 echo "Downloading Blender ${BLENDER_VERSION}..."
@@ -54,7 +64,7 @@ echo "Upgrading pip for Blender Python..."
 
 # --- Install PyTorch and Torchvision FIRST ---
 # This is crucial because some packages (like torch-scatter) need torch during their own setup
-echo "Installing PyTorch ${TORCH_VERSION} and Torchvision ${TORCHVISION_VERSION} first..."
+echo "Installing PyTorch ${TORCH_VERSION} and Torchvision ${TORCHVISION_VERSION} first (CUDA: ${TORCH_CUDA_SUFFIX})..."
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
     torch==${TORCH_VERSION} \
     torchvision==${TORCHVISION_VERSION} \
@@ -62,8 +72,19 @@ echo "Installing PyTorch ${TORCH_VERSION} and Torchvision ${TORCHVISION_VERSION}
 
 echo "PyTorch and Torchvision installation complete."
 
-# --- Install the rest of the packages from unirig_requirements.txt ---
-# Pip will skip torch and torchvision if they are already installed at the correct version
+# --- Install torch-scatter and torch-cluster using direct wheel URLs SECOND ---
+# This avoids build issues by using pre-compiled wheels directly.
+echo "Installing torch-scatter and torch-cluster from specific wheels..."
+echo "Scatter wheel: ${TORCH_SCATTER_WHL_URL}"
+echo "Cluster wheel: ${TORCH_CLUSTER_WHL_URL}"
+"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
+    "${TORCH_SCATTER_WHL_URL}" \
+    "${TORCH_CLUSTER_WHL_URL}"
+
+echo "torch-scatter and torch-cluster installation complete."
+
+# --- Install the rest of the packages from unirig_requirements.txt THIRD ---
+# Pip will skip torch, torchvision, torch-scatter, torch-cluster if already installed.
 echo "Installing remaining packages from ${UNIRIG_REQS_FILE}..."
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir -r "${UNIRIG_REQS_FILE}"
 
@@ -75,3 +96,4 @@ rm /tmp/${BLENDER_TARBALL}
 echo "Cleanup complete."
 
 echo "Blender setup finished successfully."
+