Update setup_blender.sh

setup_blender.sh

@@ -23,8 +23,9 @@ 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}"
 
-# Specific flash-attn version to try, known for good wheel coverage
-FLASH_ATTN_VERSION_TO_INSTALL="2.5.8"
+# Direct URL for the compatible flash-attn wheel for v2.5.8
+# 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.8/flash_attn-2.5.8+pt23cu121-cp311-cp311-linux_x86_64.whl"
 
 # --- Set Environment Variables for Build ---
 export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda} # This might be nominal if nvcc isn't actually used
@@ -45,7 +46,7 @@ echo "Blender Python Executable: ${BLENDER_PY_EXEC}"
 echo "Using CUDA_HOME=${CUDA_HOME}"
 echo "Targeting PyTorch for CUDA: ${TARGET_CUDA_VERSION_SHORT}"
 echo "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}"
-echo "Attempting to install flash-attn version: ${FLASH_ATTN_VERSION_TO_INSTALL}"
+echo "Attempting to install flash-attn from direct wheel URL: ${FLASH_ATTN_WHEEL_URL}"
 
 # --- Download and Extract Blender ---
 mkdir -p "${BLENDER_INSTALL_BASE}"
@@ -106,17 +107,13 @@ 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 flash-attn==${FLASH_ATTN_VERSION_TO_INSTALL} (binary only, no build isolation, checking PyTorch index)..."
-# Try to install flash-attn, strongly preferring a binary wheel.
-# --no-build-isolation allows it to see the already installed PyTorch.
-# --only-binary ":all:" tells pip to not attempt source builds for any package in this command.
-# --find-links adds PyTorch's index for this specific command, in case compatible wheels are there.
+echo "Step 2: Installing flash-attn from direct wheel URL..."
+# Install flash-attn from a direct wheel URL to ensure compatibility and avoid source build.
+# Using --no-deps as we manage other dependencies separately and assume the wheel is self-contained or relies on PyTorch.
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
-    --no-build-isolation \
-    --only-binary ":all:" \
-    --find-links ${TORCH_INDEX_URL} \
-    flash-attn==${FLASH_ATTN_VERSION_TO_INSTALL} -vvv
-echo "flash-attn installation attempted."
+    --no-deps \
+    "${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.