Update setup_blender.sh

setup_blender.sh

@@ -29,7 +29,7 @@ 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}"
-export MAX_JOBS=${MAX_JOBS:-4}
+export MAX_JOBS=${MAX_JOBS:-4} # Already set, good for potential compilations
 
 PYTHON_INCLUDE_DIR="/usr/include/python${BLENDER_PYTHON_VERSION#python}"
 export CPATH="${PYTHON_INCLUDE_DIR}:${CPATH}"
@@ -96,6 +96,9 @@ 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)..."
+"${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}..."
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
     torch==${TORCH_VERSION} \
@@ -103,12 +106,11 @@ 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)..."
-# 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.
+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.
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
-    --only-binary flash-attn \
+    --no-build-isolation \
     flash-attn==${FLASH_ATTN_VERSION_TO_INSTALL} -vvv
 echo "flash-attn installation attempted from PyPI."