Update setup_blender.sh

setup_blender.sh

@@ -23,10 +23,8 @@ 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
-# Using v2.5.8 as it has a pt23cu121 wheel available.
-# 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"
+# flash-attn version known to have compatible wheels (Python 3.11, PyTorch 2.3.1+cu121)
+FLASH_ATTN_VERSION="2.5.8"
 
 # --- Set Environment Variables for Build ---
 export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
@@ -46,7 +44,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 from: ${FLASH_ATTN_WHEEL_URL}"
+echo "Attempting to install flash-attn version: ${FLASH_ATTN_VERSION}"
 
 # --- Download and Extract Blender ---
 mkdir -p "${BLENDER_INSTALL_BASE}"
@@ -104,14 +102,18 @@ 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 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 2: Installing flash-attn==${FLASH_ATTN_VERSION} from PyPI (binary only)..."
+# Install flash-attn from PyPI, specifying version and forcing binary only
+# The --only-binary flag applies to the package name immediately following it.
+"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
+    --only-binary flash-attn \
+    flash-attn==${FLASH_ATTN_VERSION} -vvv
+echo "flash-attn installation attempted from PyPI."
 
 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.
+# These packages should find the PyTorch and flash-attn already installed.
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
     -r "${UNIRIG_REQS_FILE_IN_SPACE}" -vvv