Update setup_blender.sh

setup_blender.sh

@@ -1,6 +1,4 @@
 #!/bin/bash
-
-# Exit immediately if a command exits with a non-zero status.
 set -e
 
 # --- Configuration ---
@@ -11,31 +9,43 @@ 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"
 INSTALL_DIR="/opt/blender-${BLENDER_VERSION}-linux-x64"
 BLENDER_PY_EXEC="${INSTALL_DIR}/${BLENDER_MAJOR_MINOR}/python/bin/${BLENDER_PYTHON_VERSION}"
-UNIRIG_REQS_FILE="/home/user/app/unirig_requirements.txt"  # Adjust path if different
+
+# Path to your Gradio Space's unirig_requirements.txt (copy of UniRig's official)
+UNIRIG_REQS_FILE_IN_SPACE="/home/user/app/unirig_requirements.txt"
+# Path to the cloned UniRig repository in the Space
+UNIRIG_REPO_CLONE_DIR="/home/user/app/UniRig" # Standard path
+
 TORCH_VERSION="2.3.1"
 TORCHVISION_VERSION="0.18.1"
-TORCH_CUDA_SUFFIX="cu121"
-TORCH_INDEX_URL="https://download.pytorch.org/whl/${TORCH_CUDA_SUFFIX}"
-# Define the PyG wheel index URL (use the one compatible with torch 2.3.1, even if named 2.3.0)
-PYG_WHEEL_INDEX="https://data.pyg.org/whl/torch-2.3.0+cu121.html"  # Compatible with 2.3.1
+NUMPY_VERSION="1.26.4"
+TARGET_CUDA_VERSION_SHORT="cu121" # Verify ZeroGPU environment.
+
+TORCH_INDEX_URL="https://download.pytorch.org/whl/${TARGET_CUDA_VERSION_SHORT}"
+PYG_FIND_LINKS_URL="https://data.pyg.org/whl/torch-${TORCH_VERSION}+${TARGET_CUDA_VERSION_SHORT}.html"
+SPCONV_PACKAGE="spconv-${TARGET_CUDA_VERSION_SHORT}"
 
 # --- Set Environment Variables for Build ---
-# Define CUDA Home (adjust if the path is different in the Space environment)
-export CUDA_HOME=/usr/local/cuda
-echo "Setting CUDA_HOME=${CUDA_HOME}"
-# Add CUDA bin directory to PATH to help find nvcc
+export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
 export PATH="${CUDA_HOME}/bin:${PATH}"
-echo "Prepending CUDA bin to PATH: ${PATH}"
-
+export MAX_JOBS=4
+echo "Using CUDA_HOME=${CUDA_HOME}"
+echo "Updated PATH: ${PATH}"
 
 # --- Download and Extract Blender ---
 echo "Downloading Blender ${BLENDER_VERSION}..."
-wget -nv -O /tmp/${BLENDER_TARBALL} ${BLENDER_URL}
-echo "Download complete."
+if [ ! -f "/tmp/${BLENDER_TARBALL}" ]; then
+    wget -nv -O /tmp/${BLENDER_TARBALL} ${BLENDER_URL}
+else
+    echo "Blender tarball /tmp/${BLENDER_TARBALL} already downloaded."
+fi
 
 echo "Extracting Blender to ${INSTALL_DIR}..."
-mkdir -p /opt
-tar -xJf /tmp/${BLENDER_TARBALL} -C /opt
+if [ ! -d "${INSTALL_DIR}/${BLENDER_MAJOR_MINOR}" ]; then
+    mkdir -p /opt
+    tar -xJf /tmp/${BLENDER_TARBALL} -C /opt
+else
+    echo "Blender already extracted to ${INSTALL_DIR}."
+fi
 echo "Extraction complete."
 
 # --- Create Blender Symlink ---
@@ -44,44 +54,73 @@ ln -sf ${INSTALL_DIR}/blender /usr/local/bin/blender
 echo "Symlink created."
 
 # --- Install Dependencies into Blender's Python ---
-echo "Installing dependencies into Blender's Python (${BLENDER_PYTHON_VERSION})..."
-
-# Check if Blender Python executable exists
+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}!"
     exit 1
 fi
-
-# Check if the requirements file exists
-if [ ! -f "${UNIRIG_REQS_FILE}" ]; then
-    echo "ERROR: Requirements file not found at ${UNIRIG_REQS_FILE}!"
+if [ ! -f "${UNIRIG_REQS_FILE_IN_SPACE}" ]; then
+    echo "ERROR: UniRig requirements file not found at ${UNIRIG_REQS_FILE_IN_SPACE}!"
     exit 1
 fi
 
-# Upgrade pip, setuptools, and wheel within Blender's Python environment
-# Run this with the modified PATH and CUDA_HOME set
 echo "Upgrading pip for Blender Python..."
-"${BLENDER_PY_EXEC}" -m pip install --upgrade pip setuptools wheel
+"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir --upgrade pip setuptools wheel
 
-# Install PyTorch and Torchvision first using the specified index URL
-echo "Installing PyTorch ${TORCH_VERSION} and Torchvision ${TORCHVISION_VERSION} (CUDA: ${TORCH_CUDA_SUFFIX})..."
-# Ensure CUDA_HOME and updated PATH are available for this step too
+echo "1. Installing PyTorch ${TORCH_VERSION} and Torchvision ${TORCHVISION_VERSION}..."
 "${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
     torch==${TORCH_VERSION} \
     torchvision==${TORCHVISION_VERSION} \
     --index-url ${TORCH_INDEX_URL}
 
-# Install remaining dependencies from the requirements file
-# Use the -f flag to point pip to the PyG wheel index page
-# CUDA_HOME and updated PATH should be inherited by this pip process
-echo "Installing remaining packages from ${UNIRIG_REQS_FILE} with PyG wheel index..."
-"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir -r "${UNIRIG_REQS_FILE}" -f ${PYG_WHEEL_INDEX}
+echo "2. Installing dependencies from ${UNIRIG_REQS_FILE_IN_SPACE} (includes bpy, flash_attn)..."
+"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir -r "${UNIRIG_REQS_FILE_IN_SPACE}"
 
-echo "Dependency installation complete."
+echo "3. Installing ${SPCONV_PACKAGE}..."
+"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir ${SPCONV_PACKAGE}
+
+echo "4. Installing torch-scatter and torch-cluster..."
+"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir \
+    torch-scatter \
+    torch-cluster \
+    -f "${PYG_FIND_LINKS_URL}"
+
+echo "5. Installing numpy==${NUMPY_VERSION}..."
+"${BLENDER_PY_EXEC}" -m pip install --no-cache-dir numpy==${NUMPY_VERSION}
+
+echo "Dependency installation for Blender's Python complete."
+
+# --- FIX: Ensure UniRig/src is treated as a package ---
+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..."
+    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."
+fi
+# --- END FIX ---
+
+# (Optional) VRM Addon installation
+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..."
+    (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."
+else
+    echo "VRM addon zip not found at ${ABSOLUTE_ADDON_PATH}, skipping addon installation."
+fi
 
 # --- Cleanup ---
-echo "Cleaning up downloaded tarball..."
-rm /tmp/${BLENDER_TARBALL}
+echo "Cleaning up downloaded Blender tarball..."
+rm -f /tmp/${BLENDER_TARBALL}
 echo "Cleanup complete."
-
 echo "Blender setup finished successfully."
+