Update setup_blender.sh

setup_blender.sh

@@ -16,7 +16,12 @@ TORCH_VERSION="2.3.1"
 TORCHVISION_VERSION="0.18.1"
 TORCH_CUDA_SUFFIX="cu121"
 TORCH_INDEX_URL="https://download.pytorch.org/whl/${TORCH_CUDA_SUFFIX}"
+# Define the PyG wheel index URL (use the one compatible with torch 2.3.1, even if named 2.3.0)
 PYG_WHEEL_INDEX="https://data.pyg.org/whl/torch-2.3.0+cu121.html"  # Compatible with 2.3.1
+# Define CUDA Home (adjust if the path is different in the Space environment)
+export CUDA_HOME=/usr/local/cuda
+echo "Setting CUDA_HOME=${CUDA_HOME}"
+
 
 # --- Download and Extract Blender ---
 echo "Downloading Blender ${BLENDER_VERSION}..."
@@ -36,28 +41,33 @@ echo "Symlink created."
 # --- Install Dependencies into Blender's Python ---
 echo "Installing dependencies into Blender's Python (${BLENDER_PYTHON_VERSION})..."
 
+# Check if Blender Python executable exists
 if [ ! -f "${BLENDER_PY_EXEC}" ]; then
     echo "ERROR: Blender Python executable not found at ${BLENDER_PY_EXEC}!"
     exit 1
 fi
 
+# Check if the requirements file exists
 if [ ! -f "${UNIRIG_REQS_FILE}" ]; then
     echo "ERROR: Requirements file not found at ${UNIRIG_REQS_FILE}!"
     exit 1
 fi
 
-# Upgrade pip
+# Upgrade pip, setuptools, and wheel within Blender's Python environment
 echo "Upgrading pip for Blender Python..."
 "${BLENDER_PY_EXEC}" -m pip install --upgrade pip setuptools wheel
 
-# Install PyTorch and Torchvision first
+# Install PyTorch and Torchvision first using the specified index URL
 echo "Installing PyTorch ${TORCH_VERSION} and Torchvision ${TORCHVISION_VERSION} (CUDA: ${TORCH_CUDA_SUFFIX})..."
+# Ensure CUDA_HOME is available for this step too if needed by torch install hooks
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
     torch==${TORCH_VERSION} \
     torchvision==${TORCHVISION_VERSION} \
     --index-url ${TORCH_INDEX_URL}
 
-# Install remaining dependencies with PyG wheel index
+# Install remaining dependencies from the requirements file
+# Use the -f flag to point pip to the PyG wheel index page
+# CUDA_HOME should be inherited by this pip process
 echo "Installing remaining packages from ${UNIRIG_REQS_FILE} with PyG wheel index..."
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir -r "${UNIRIG_REQS_FILE}" -f ${PYG_WHEEL_INDEX}
 
@@ -68,4 +78,4 @@ echo "Cleaning up downloaded tarball..."
 rm /tmp/${BLENDER_TARBALL}
 echo "Cleanup complete."
 
-echo "Blender setup finished successfully."
+echo "Blender setup finished successfully."