Update setup_blender.sh

setup_blender.sh

@@ -23,9 +23,6 @@ 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}"
 
-# flash-attn version known to have compatible wheels (Python 3.11, PyTorch 2.3.1+cu121)
-FLASH_ATTN_VERSION_TO_INSTALL="2.5.9.post1" # This version has pt23cu121 wheels
-
 # --- 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}"
@@ -45,7 +42,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 latest compatible flash-attn binary."
 
 # --- Download and Extract Blender ---
 mkdir -p "${BLENDER_INSTALL_BASE}"
@@ -106,15 +103,14 @@ 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} from PyPI (binary only, no build isolation)..."
-# Try to install flash-attn, strongly preferring a binary wheel.
+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.
 # The --no-build-isolation flag allows it to see the already installed PyTorch.
-# The --only-binary :all: flag (or --only-binary flash-attn) tells pip to not attempt source builds.
-# We use :all: here for this specific command to be absolutely sure pip doesn't try to build flash-attn or its immediate small deps from source.
+# The --only-binary ":all:" flag tells pip to not attempt source builds for any package in this command.
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
     --no-build-isolation \
     --only-binary ":all:" \
-    flash-attn==${FLASH_ATTN_VERSION_TO_INSTALL} -vvv
+    flash-attn -vvv
 echo "flash-attn installation attempted from PyPI."
 
 echo "Step 3: Installing remaining dependencies from ${UNIRIG_REQS_FILE_IN_SPACE}..."