Update setup_blender.sh

setup_blender.sh

@@ -14,6 +14,10 @@ BLENDER_PY_EXEC="${INSTALL_DIR}/${BLENDER_MAJOR_MINOR}/python/bin/${BLENDER_PYTH
 # 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
 
 # --- Download and Extract Blender ---
 echo "Downloading Blender ${BLENDER_VERSION}..."
@@ -48,8 +52,19 @@ fi
 echo "Upgrading pip for Blender Python..."
 "${BLENDER_PY_EXEC}" -m pip install --upgrade pip setuptools wheel
 
-# Install packages from unirig_requirements.txt
-echo "Installing packages from ${UNIRIG_REQS_FILE}..."
+# --- 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..."
+"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
+    torch==${TORCH_VERSION} \
+    torchvision==${TORCHVISION_VERSION} \
+    --index-url ${TORCH_INDEX_URL}
+
+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
+echo "Installing remaining packages from ${UNIRIG_REQS_FILE}..."
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir -r "${UNIRIG_REQS_FILE}"
 
 echo "UniRig dependency installation for Blender Python complete."