Update setup_blender.sh

setup_blender.sh

@@ -23,6 +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"
+
 # --- Set Environment Variables for Build ---
 export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda} # This might be nominal if nvcc isn't actually used
 export PATH="${CUDA_HOME}/bin:${LOCAL_BIN_DIR}:${PATH}"
@@ -42,7 +45,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 latest compatible flash-attn binary."
+echo "Attempting to install flash-attn version: ${FLASH_ATTN_VERSION_TO_INSTALL}"
 
 # --- Download and Extract Blender ---
 mkdir -p "${BLENDER_INSTALL_BASE}"
@@ -103,14 +106,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 from PyPI (binary only, no build isolation, latest compatible version)..."
-# Try to install flash-attn, strongly preferring a binary wheel, letting pip choose the version.
+echo "Step 2: Installing flash-attn==${FLASH_ATTN_VERSION_TO_INSTALL} from PyPI (with --no-build-isolation)..."
+# Try to install flash-attn, letting pip search PyPI.
 # The --no-build-isolation flag allows it to see the already installed PyTorch.
-# The --only-binary ":all:" flag tells pip to not attempt source builds for any package in this command.
+# We are NOT using --only-binary here to see if it attempts a source build and what error occurs.
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
     --no-build-isolation \
-    --only-binary ":all:" \
-    flash-attn -vvv
+    flash-attn==${FLASH_ATTN_VERSION_TO_INSTALL} -vvv
 echo "flash-attn installation attempted from PyPI."
 
 echo "Step 3: Installing remaining dependencies from ${UNIRIG_REQS_FILE_IN_SPACE}..."