Update setup_blender.sh

setup_blender.sh

@@ -24,20 +24,20 @@ 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"
+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}
+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}"
-export MAX_JOBS=${MAX_JOBS:-4} # Already set, good for potential compilations
+export MAX_JOBS=${MAX_JOBS:-4} # For compilation jobs if any occur
 
 PYTHON_INCLUDE_DIR="/usr/include/python${BLENDER_PYTHON_VERSION#python}"
 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 if any part of flash-attn or other packages try to compile kernels.
-export TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"
+# TORCH_CUDA_ARCH_LIST is important if flash-attn *does* try to compile parts of itself.
+export TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0" # Common GPU architectures
 
 echo "--- Setup Script Start ---"
 echo "Target Blender Installation Directory: ${INSTALL_DIR}"
@@ -96,7 +96,7 @@ fi
 echo "Upgrading pip for Blender Python..."
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir --upgrade pip setuptools wheel -vvv
 
-echo "Installing packaging and ninja for flash-attn build (if needed)..."
+echo "Installing packaging and ninja (recommended for flash-attn)..."
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir packaging ninja -vvv
 
 echo "Step 1: Installing PyTorch ${TORCH_VERSION} (for CUDA ${TARGET_CUDA_VERSION_SHORT}) and Torchvision ${TORCHVISION_VERSION}..."
@@ -106,18 +106,20 @@ 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 with --no-build-isolation..."
-# Install flash-attn from PyPI, specifying version and using --no-build-isolation
-# This allows flash-attn to use the already installed PyTorch and other packages.
+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.
+# 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.
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
     --no-build-isolation \
+    --only-binary ":all:" \
     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}..."
-# Ensure flash-attn is REMOVED from unirig_requirements.txt to avoid conflicts or reinstallation attempts.
+# Ensure flash-attn is REMOVED from unirig_requirements.txt.
 # 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
 
@@ -139,7 +141,7 @@ else
 fi
 
 # (Optional) VRM Addon installation
-VRM_ADDON_REL_PATH="blender/add-on-vrm-v2.20.77_modified.zip" # Relative to UniRig repo root
+VRM_ADDON_REL_PATH="blender/add-on-vrm-v2.20.77_modified.zip"
 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..."