setup_blender.sh

#!/bin/bash

# Exit immediately if a command exits with a non-zero status.
set -e

# --- Configuration ---
BLENDER_VERSION="4.2.0"
BLENDER_MAJOR_MINOR="4.2"
BLENDER_PYTHON_VERSION="python3.11"
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
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

# --- 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 PATH="${CUDA_HOME}/bin:${PATH}"
echo "Prepending CUDA bin to PATH: ${PATH}"


# --- Download and Extract Blender ---
echo "Downloading Blender ${BLENDER_VERSION}..."
wget -nv -O /tmp/${BLENDER_TARBALL} ${BLENDER_URL}
echo "Download complete."

echo "Extracting Blender to ${INSTALL_DIR}..."
mkdir -p /opt
tar -xJf /tmp/${BLENDER_TARBALL} -C /opt
echo "Extraction complete."

# --- Create Blender Symlink ---
echo "Creating symlink for Blender executable..."
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
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}!"
    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

# 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
"${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 "Dependency installation complete."

# --- Cleanup ---
echo "Cleaning up downloaded tarball..."
rm /tmp/${BLENDER_TARBALL}
echo "Cleanup complete."

echo "Blender setup finished successfully."