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}"
 
+# flash-attn version known to have compatible wheels (Python 3.11, PyTorch 2.3.1+cu121)
+FLASH_ATTN_VERSION_TO_INSTALL="2.5.8"
+
 # --- Set Environment Variables for Build ---
 export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
 export PATH="${CUDA_HOME}/bin:${LOCAL_BIN_DIR}:${PATH}"
@@ -33,7 +36,7 @@ export CPATH="${PYTHON_INCLUDE_DIR}:${CPATH}"
 export C_INCLUDE_PATH="${PYTHON_INCLUDE_DIR}:${C_INCLUDE_PATH}"
 export CPLUS_INCLUDE_PATH="${PYTHON_INCLUDE_DIR}:${CPLUS_INCLUDE_PATH}"
 
-# TORCH_CUDA_ARCH_LIST might still be useful for other compiled packages if any
+# TORCH_CUDA_ARCH_LIST might still be useful if any part of flash-attn or other packages try to compile kernels.
 export TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"
 
 echo "--- Setup Script Start ---"
@@ -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 "NOTE: flash-attn installation is SKIPPED."
+echo "Attempting to install flash-attn version: ${FLASH_ATTN_VERSION_TO_INSTALL}"
 
 # --- Download and Extract Blender ---
 mkdir -p "${BLENDER_INSTALL_BASE}"
@@ -100,11 +103,19 @@ echo "Step 1: Installing PyTorch ${TORCH_VERSION} (for CUDA ${TARGET_CUDA_VERSIO
     --index-url ${TORCH_INDEX_URL} -vvv
 echo "PyTorch and Torchvision installation attempted."
 
-# Step 2 for flash-attn has been removed.
+echo "Step 2: Installing flash-attn==${FLASH_ATTN_VERSION_TO_INSTALL} from PyPI (binary only)..."
+# Install flash-attn from PyPI, specifying version and forcing binary only.
+# This command does NOT use the PyTorch index-url, allowing pip to search PyPI normally for flash-attn.
+# 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_TO_INSTALL} -vvv
+echo "flash-attn installation attempted from PyPI."
 
-echo "Step 2 (formerly Step 3): Installing remaining dependencies from ${UNIRIG_REQS_FILE_IN_SPACE}..."
-# Ensure flash-attn is REMOVED from unirig_requirements.txt.
+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
 
@@ -130,8 +141,6 @@ VRM_ADDON_REL_PATH="blender/add-on-vrm-v2.20.77_modified.zip" # Relative to UniR
 ABSOLUTE_ADDON_PATH="${UNIRIG_REPO_CLONE_DIR}/${VRM_ADDON_REL_PATH}"
 if [ -f "${ABSOLUTE_ADDON_PATH}" ]; then
     echo "Attempting to install optional VRM addon for Blender..."
-    # Try to install and enable the addon. If it fails, print a warning and continue.
-    # The || true ensures that if the python command fails, the script doesn't exit due to set -e
     (cd "${UNIRIG_REPO_CLONE_DIR}" && \
      "${BLENDER_PY_EXEC}" -c "import bpy, os; print(f'Attempting to install addon from: {os.path.abspath(\"${VRM_ADDON_REL_PATH}\")}'); bpy.ops.preferences.addon_install(overwrite=True, filepath=os.path.abspath('${VRM_ADDON_REL_PATH}')); print('Addon installation script executed. Attempting to enable...'); bpy.ops.preferences.addon_enable(module='io_scene_vrm'); print('VRM Addon enabled successfully.')") \
     || echo "WARNING: VRM addon installation or enabling failed. This is an optional addon. Continuing setup..."