Update setup_blender.sh

setup_blender.sh

@@ -20,12 +20,13 @@ UNIRIG_REPO_CLONE_DIR="${APP_DIR}/UniRig"
 
 TORCH_VERSION="2.3.1"
 TORCHVISION_VERSION="0.18.1"
-TARGET_CUDA_VERSION_SHORT="cu121"
+# ** MODIFIED: Changed to cu122 to align with common flash-attn pre-built wheels **
+TARGET_CUDA_VERSION_SHORT="cu122"
 TORCH_INDEX_URL="https://download.pytorch.org/whl/${TARGET_CUDA_VERSION_SHORT}"
 
 # --- Set Environment Variables for Build ---
-export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
-export PATH="${CUDA_HOME}/bin:${LOCAL_BIN_DIR}:${PATH}"
+export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda} # This might be nominal in ZeroGPU if nvcc isn't present
+export PATH="${CUDA_HOME}/bin:${LOCAL_BIN_DIR}:${PATH}" # Add CUDA_HOME/bin to PATH
 export MAX_JOBS=${MAX_JOBS:-4}
 
 PYTHON_INCLUDE_DIR="/usr/include/python${BLENDER_PYTHON_VERSION#python}"
@@ -33,18 +34,16 @@ 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}"
 
-# ** ADDED FOR FLASH-ATTN: Specify CUDA Architectures **
+# Specify CUDA Architectures for compilation (e.g. flash-attn if it still tries to build parts)
 # For ZeroGPU, common architectures are Ampere (8.0, 8.6) or Hopper (9.0).
-# Adjust if you know the specific GPU architecture of your ZeroGPU instance.
-# Using a common set for broader compatibility.
-export TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"
-# Alternatively, for a more focused build if you know the target:
+export TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0" # General set for compatibility
 # export TORCH_CUDA_ARCH_LIST="8.6" # Example for A10G on ZeroGPU
 
 echo "--- Setup Script Start ---"
 echo "Target Blender Installation Directory: ${INSTALL_DIR}"
 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 "Initial PATH: ${PATH}"
 echo "CPATH set to: ${CPATH}"
@@ -105,31 +104,20 @@ fi
 echo "Upgrading pip for Blender Python..."
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir --upgrade pip setuptools wheel -vvv
 
-echo "Step 1: Installing PyTorch ${TORCH_VERSION} and Torchvision ${TORCHVISION_VERSION} for Blender's Python..."
+echo "Step 1: Installing PyTorch ${TORCH_VERSION} (for CUDA ${TARGET_CUDA_VERSION_SHORT}) and Torchvision ${TORCHVISION_VERSION} for Blender's Python..."
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
     torch==${TORCH_VERSION} \
     torchvision==${TORCHVISION_VERSION} \
     --index-url ${TORCH_INDEX_URL} -vvv
 echo "PyTorch and Torchvision installation attempted."
 
-# ** ADDED: Separate step for flash-attn installation **
-# Ensure 'flash-attn' is listed in your unirig_requirements.txt if you want pip to manage its version,
-# or specify version here e.g., flash-attn==<version>
-# The TORCH_CUDA_ARCH_LIST environment variable set earlier will be used by its build process.
-# Also ensure 'build-essential' is in your packages.txt for g++ compiler.
-echo "Step 2: Installing flash-attn for Blender's Python..."
-# Assuming flash-attn is in unirig_requirements.txt and will be picked up by the next step,
-# OR if it needs to be a specific version not in the file, install it explicitly:
-# "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir flash-attn -vvv
-# For now, we'll rely on it being in unirig_requirements.txt and the env vars being set.
-# If flash-attn continues to fail, you might need to install it explicitly here.
-# The most robust way is to let it be installed by the requirements file if it's listed there
-# with any specific version constraints UniRig might have. The environment variables should help.
+# Step 2 (flash-attn) is now implicitly handled by Step 3, relying on the correct PyTorch CUDA version.
+echo "Step 2: (flash-attn will be installed via requirements file in the next step)"
 
 echo "Step 3: Installing remaining dependencies from ${UNIRIG_REQS_FILE_IN_SPACE}..."
-# This will install torch-scatter, torch-cluster, spconv, flash-attn (if listed), bpy, etc.
-# flash-attn build will use TORCH_CUDA_ARCH_LIST.
-# torch-scatter etc. will use the pre-installed PyTorch.
+# This will install torch-scatter, torch-cluster, spconv, flash-attn (if listed and compatible wheel found), bpy, etc.
+# flash-attn build (if a wheel is not found) would use TORCH_CUDA_ARCH_LIST.
+# Hopefully, a pre-built wheel for flash-attn (e.g., +pt23cu122) will be found due to PyTorch being cu122.
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir -r "${UNIRIG_REQS_FILE_IN_SPACE}" -vvv
 
 echo "Dependency installation for Blender's Python complete."