Update setup_blender.sh

setup_blender.sh

@@ -18,9 +18,14 @@ 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
+
+# --- Set Environment Variables for Build ---
 # 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}"
+# Add CUDA bin directory to PATH to help find nvcc
+export PATH="${CUDA_HOME}/bin:${PATH}"
+echo "Prepending CUDA bin to PATH: ${PATH}"
 
 
 # --- Download and Extract Blender ---
@@ -54,12 +59,13 @@ if [ ! -f "${UNIRIG_REQS_FILE}" ]; then
 fi
 
 # Upgrade pip, setuptools, and wheel within Blender's Python environment
+# Run this with the modified PATH and CUDA_HOME set
 echo "Upgrading pip for Blender Python..."
 "${BLENDER_PY_EXEC}" -m pip install --upgrade pip setuptools wheel
 
 # 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
+# Ensure CUDA_HOME and updated PATH are available for this step too
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
     torch==${TORCH_VERSION} \
     torchvision==${TORCHVISION_VERSION} \
@@ -67,7 +73,7 @@ echo "Installing PyTorch ${TORCH_VERSION} and Torchvision ${TORCHVISION_VERSION}
 
 # 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
+# CUDA_HOME and updated PATH 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}