Update setup_blender.sh

setup_blender.sh

@@ -20,13 +20,12 @@ UNIRIG_REPO_CLONE_DIR="${APP_DIR}/UniRig"
 
 TORCH_VERSION="2.3.1"
 TORCHVISION_VERSION="0.18.1"
-# ** MODIFIED: Reverted to cu121 as PyTorch 2.3.1 wheels are readily available for it **
-TARGET_CUDA_VERSION_SHORT="cu121"
+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}"
 
 # --- Set Environment Variables for Build ---
-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 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}"
@@ -34,10 +33,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}"
 
-# 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).
 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}"
@@ -45,15 +41,6 @@ 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}"
-echo "Expected Python include directory: ${PYTHON_INCLUDE_DIR}"
-
-if [ -f "${PYTHON_INCLUDE_DIR}/Python.h" ]; then
-    echo "Found Python.h at ${PYTHON_INCLUDE_DIR}/Python.h"
-else
-    echo "WARNING: Python.h NOT FOUND at ${PYTHON_INCLUDE_DIR}/Python.h. Ensure python${BLENDER_PYTHON_VERSION#python}-dev is in packages.txt and installed."
-fi
 
 # --- Download and Extract Blender ---
 mkdir -p "${BLENDER_INSTALL_BASE}"
@@ -73,7 +60,7 @@ if [ ! -d "${INSTALL_DIR}" ] || [ -z "$(ls -A "${INSTALL_DIR}")" ]; then
     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}"
+        echo "ERROR: Blender extraction failed. Expected: ${INSTALL_DIR}"
         ls -la "${BLENDER_INSTALL_BASE}"
         exit 1
     fi
@@ -87,7 +74,7 @@ if [ -f "${INSTALL_DIR}/blender" ]; then
     ln -sf "${INSTALL_DIR}/blender" "${LOCAL_BIN_DIR}/blender"
     echo "Local symlink created at ${LOCAL_BIN_DIR}/blender."
 else
-    echo "WARNING: Blender executable not found at ${INSTALL_DIR}/blender. Cannot create symlink."
+    echo "WARNING: Blender executable not found at ${INSTALL_DIR}/blender."
 fi
 
 # --- Install Dependencies into Blender's Python ---
@@ -104,21 +91,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} (for CUDA ${TARGET_CUDA_VERSION_SHORT}) 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}..."
 "${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."
 
-# 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}..."
+echo "Step 2: Installing remaining dependencies from ${UNIRIG_REQS_FILE_IN_SPACE}..."
+# ** MODIFIED: Added --only-binary flash-attn to prevent source builds for flash-attn **
 # 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., +pt23cu121) will be found due to PyTorch being cu121.
-"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir -r "${UNIRIG_REQS_FILE_IN_SPACE}" -vvv
+# Ensure 'flash-attn==2.5.9.post1' (or similar compatible version) is in unirig_requirements.txt
+"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
+    --only-binary flash-attn \
+    -r "${UNIRIG_REQS_FILE_IN_SPACE}" -vvv
 
 echo "Dependency installation for Blender's Python complete."
 
@@ -127,14 +113,14 @@ UNIRIG_SRC_DIR="${UNIRIG_REPO_CLONE_DIR}/src"
 INIT_PY_PATH="${UNIRIG_SRC_DIR}/__init__.py"
 if [ -d "${UNIRIG_SRC_DIR}" ]; then
   if [ ! -f "${INIT_PY_PATH}" ]; then
-    echo "Creating missing __init__.py in ${UNIRIG_SRC_DIR} to make it a package..."
+    echo "Creating missing __init__.py in ${UNIRIG_SRC_DIR}..."
     touch "${INIT_PY_PATH}"
     echo "__init__.py created."
   else
     echo "${INIT_PY_PATH} already exists."
   fi
 else
-  echo "WARNING: UniRig src directory not found at ${UNIRIG_SRC_DIR}. Cannot create __init__.py."
+  echo "WARNING: UniRig src directory not found at ${UNIRIG_SRC_DIR}."
 fi
 
 # (Optional) VRM Addon installation
@@ -146,7 +132,7 @@ if [ -f "${ABSOLUTE_ADDON_PATH}" ]; then
      "${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."
 else
-    echo "VRM addon zip not found at ${ABSOLUTE_ADDON_PATH}, skipping addon installation."
+    echo "VRM addon zip not found at ${ABSOLUTE_ADDON_PATH}, skipping."
 fi
 
 # --- Cleanup ---