Update setup_blender.sh

setup_blender.sh

@@ -23,6 +23,10 @@ TORCHVISION_VERSION="0.18.1"
 TARGET_CUDA_VERSION_SHORT="cu121" # PyTorch 2.3.1 wheels are available for cu121
 TORCH_INDEX_URL="https://download.pytorch.org/whl/${TARGET_CUDA_VERSION_SHORT}"
 
+# Direct URL for the compatible flash-attn wheel
+# Corresponds to Python 3.11 (cp311), PyTorch 2.3.1+cu121
+FLASH_ATTN_WHEEL_URL="https://github.com/Dao-AILab/flash-attention/releases/download/v2.5.9.post1/flash_attn-2.5.9.post1+pt23cu121-cp311-cp311-linux_x86_64.whl"
+
 # --- Set Environment Variables for Build ---
 export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
 export PATH="${CUDA_HOME}/bin:${LOCAL_BIN_DIR}:${PATH}"
@@ -98,12 +102,15 @@ echo "Step 1: Installing PyTorch ${TORCH_VERSION} (for CUDA ${TARGET_CUDA_VERSIO
     --index-url ${TORCH_INDEX_URL} -vvv
 echo "PyTorch and Torchvision installation attempted."
 
-echo "Step 2: Installing remaining dependencies from ${UNIRIG_REQS_FILE_IN_SPACE}..."
-# ** MODIFIED: Added --only-binary flash-attn to prevent source builds for flash-attn **
-# This will install torch-scatter, torch-cluster, spconv, flash-attn (if listed and compatible wheel found), bpy, etc.
-# Ensure 'flash-attn==2.5.9.post1' (or similar compatible version) is in unirig_requirements.txt
+echo "Step 2: Installing flash-attn directly from wheel URL..."
+# Install flash-attn from a direct wheel URL to ensure compatibility
+"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir "${FLASH_ATTN_WHEEL_URL}" -vvv
+echo "flash-attn installation attempted from wheel."
+
+echo "Step 3: Installing remaining dependencies from ${UNIRIG_REQS_FILE_IN_SPACE}..."
+# Ensure flash-attn is REMOVED from unirig_requirements.txt to avoid conflicts or reinstallation attempts.
+# This will install torch-scatter, torch-cluster, spconv, bpy, etc.
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
-    --only-binary flash-attn \
     -r "${UNIRIG_REQS_FILE_IN_SPACE}" -vvv
 
 echo "Dependency installation for Blender's Python complete."