Update setup_blender.sh

setup_blender.sh

@@ -8,40 +8,44 @@ BLENDER_PYTHON_VERSION="python3.11" # Should match the -dev package (e.g., pytho
 BLENDER_TARBALL="blender-${BLENDER_VERSION}-linux-x64.tar.xz"
 BLENDER_URL="https://download.blender.org/release/Blender${BLENDER_MAJOR_MINOR}/blender-${BLENDER_VERSION}-linux-x64.tar.xz"
 
-# ** MODIFIED FOR LOCAL BLENDER INSTALLATION **
-# Assuming this script is run from /home/user/app (standard for HF Spaces app.py)
-APP_DIR="/home/user/app" # Or use $(pwd) if script is guaranteed to be in app root
-BLENDER_INSTALL_BASE="${APP_DIR}/blender_installation" # Base directory for Blender versions
-INSTALL_DIR="${BLENDER_INSTALL_BASE}/blender-${BLENDER_VERSION}-linux-x64" # Specific version install
-LOCAL_BIN_DIR="${APP_DIR}/local_bin" # For local symlinks, if used
+APP_DIR="/home/user/app"
+BLENDER_INSTALL_BASE="${APP_DIR}/blender_installation"
+INSTALL_DIR="${BLENDER_INSTALL_BASE}/blender-${BLENDER_VERSION}-linux-x64"
+LOCAL_BIN_DIR="${APP_DIR}/local_bin"
 
 BLENDER_PY_EXEC="${INSTALL_DIR}/${BLENDER_MAJOR_MINOR}/python/bin/${BLENDER_PYTHON_VERSION}"
 
 UNIRIG_REQS_FILE_IN_SPACE="${APP_DIR}/unirig_requirements.txt"
 UNIRIG_REPO_CLONE_DIR="${APP_DIR}/UniRig"
 
-# PyTorch specific versions and URLs for pre-installation
 TORCH_VERSION="2.3.1"
 TORCHVISION_VERSION="0.18.1"
-TARGET_CUDA_VERSION_SHORT="cu121" # Ensure this matches ZeroGPU environment
+TARGET_CUDA_VERSION_SHORT="cu121"
 TORCH_INDEX_URL="https://download.pytorch.org/whl/${TARGET_CUDA_VERSION_SHORT}"
-# PYG_FIND_LINKS_URL is used when installing from requirements.txt later
-# PYG_FIND_LINKS_URL="https://data.pyg.org/whl/torch-${TORCH_VERSION}+${TARGET_CUDA_VERSION_SHORT}.html"
 
 # --- Set Environment Variables for Build ---
-export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda} # Keep for system CUDA if needed by compilations
-export PATH="${CUDA_HOME}/bin:${LOCAL_BIN_DIR}:${PATH}" # Add local_bin to PATH early
-export MAX_JOBS=${MAX_JOBS:-4} # For compilation processes
+export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
+export PATH="${CUDA_HOME}/bin:${LOCAL_BIN_DIR}:${PATH}"
+export MAX_JOBS=${MAX_JOBS:-4}
 
-PYTHON_INCLUDE_DIR="/usr/include/python${BLENDER_PYTHON_VERSION#python}" # e.g., /usr/include/python3.11
+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}"
 
+# ** ADDED FOR FLASH-ATTN: Specify CUDA Architectures **
+# 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.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 "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}"
 echo "Initial PATH: ${PATH}"
 echo "CPATH set to: ${CPATH}"
 echo "Expected Python include directory: ${PYTHON_INCLUDE_DIR}"
@@ -53,12 +57,9 @@ else
 fi
 
 # --- Download and Extract Blender ---
-# Create the local installation base directory
 mkdir -p "${BLENDER_INSTALL_BASE}"
-mkdir -p "${LOCAL_BIN_DIR}" # Ensure local_bin exists
+mkdir -p "${LOCAL_BIN_DIR}"
 
-# Check if the *specific version's directory* inside INSTALL_DIR exists and is not empty.
-# Blender extracts to a folder like "blender-4.2.0-linux-x64"
 if [ ! -d "${INSTALL_DIR}" ] || [ -z "$(ls -A "${INSTALL_DIR}")" ]; then
     echo "Blender not found or directory empty at ${INSTALL_DIR}. Proceeding with download and extraction."
     echo "Downloading Blender ${BLENDER_VERSION}..."
@@ -69,14 +70,11 @@ if [ ! -d "${INSTALL_DIR}" ] || [ -z "$(ls -A "${INSTALL_DIR}")" ]; then
     fi
 
     echo "Extracting Blender to ${BLENDER_INSTALL_BASE}..."
-    # Extracts to a subdirectory like blender-4.2.0-linux-x64 inside BLENDER_INSTALL_BASE
     tar -xJf "/tmp/${BLENDER_TARBALL}" -C "${BLENDER_INSTALL_BASE}"
-    # The tarball extracts into a folder named "blender-4.2.0-linux-x64", so INSTALL_DIR should be correct.
     if [ -d "${INSTALL_DIR}" ]; then
         echo "Blender extracted successfully to ${INSTALL_DIR}"
     else
         echo "ERROR: Blender extraction failed or extracted to an unexpected path. Expected: ${INSTALL_DIR}"
-        # List contents of BLENDER_INSTALL_BASE for debugging
         ls -la "${BLENDER_INSTALL_BASE}"
         exit 1
     fi
@@ -85,21 +83,18 @@ else
 fi
 echo "Extraction complete."
 
-# --- Create Blender Symlink (Optional - to local_bin) ---
 if [ -f "${INSTALL_DIR}/blender" ]; then
     echo "Creating local symlink for Blender executable in ${LOCAL_BIN_DIR}..."
     ln -sf "${INSTALL_DIR}/blender" "${LOCAL_BIN_DIR}/blender"
-    echo "Local symlink created at ${LOCAL_BIN_DIR}/blender. Make sure ${LOCAL_BIN_DIR} is in your PATH."
+    echo "Local symlink created at ${LOCAL_BIN_DIR}/blender."
 else
     echo "WARNING: Blender executable not found at ${INSTALL_DIR}/blender. Cannot create symlink."
 fi
 
-
 # --- Install Dependencies into Blender's Python ---
 echo "Installing dependencies into Blender's Python (${BLENDER_PY_EXEC})..."
 if [ ! -f "${BLENDER_PY_EXEC}" ]; then
     echo "ERROR: Blender Python executable not found at ${BLENDER_PY_EXEC}!"
-    echo "Check Blender extraction and paths."
     exit 1
 fi
 if [ ! -f "${UNIRIG_REQS_FILE_IN_SPACE}" ]; then
@@ -111,19 +106,30 @@ 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..."
-# Install torch and torchvision first, as other packages (torch-scatter, etc.) depend on it being present for their setup.
-# Ensure your unirig_requirements.txt also lists these versions so pip can resolve correctly later if needed,
-# but this direct install ensures they are present for build steps of other packages.
 "${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."
 
-echo "Step 2: Installing remaining dependencies from ${UNIRIG_REQS_FILE_IN_SPACE}..."
-# Now install the rest of the packages.
-# unirig_requirements.txt should contain torch-scatter, torch-cluster, spconv, flash_attn (if needed), bpy, etc.
-# with their respective find-links URLs or direct package names.
+# ** 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.
+
+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.
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir -r "${UNIRIG_REQS_FILE_IN_SPACE}" -vvv
 
 echo "Dependency installation for Blender's Python complete."
@@ -144,11 +150,10 @@ 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 "Installing optional VRM addon for Blender..."
-    # Ensure CWD is UniRig repo for relative path in addon script if any
     (cd "${UNIRIG_REPO_CLONE_DIR}" && \
      "${BLENDER_PY_EXEC}" -c "import bpy, os; bpy.ops.preferences.addon_install(overwrite=True, filepath=os.path.abspath('${VRM_ADDON_REL_PATH}')); bpy.ops.preferences.addon_enable(module='io_scene_vrm'); print('VRM Addon installed and enabled.')")
     echo "VRM addon installation attempted."