scripts/setup_trellis_env.sh

#!/usr/bin/env bash
# Installs microsoft/TRELLIS.2 into the pre-existing `trellis2` conda env
# (python 3.10, torch==2.6.0+cu124 / torchvision==0.21.0+cu124 already
# installed -- that is TRELLIS.2's own tested combination per its setup.sh).
#
# This script deliberately does NOT create the conda env or touch torch: both
# are assumed done already. It only adds TRELLIS.2's own dependencies and
# compiled CUDA extensions on top.
#
# Idempotent: safe to re-run after a partial failure. Existing clones under
# .trellis_ext_build/ are reused rather than re-cloned; pip installs are
# re-run every time (cheap once wheels/build artifacts are cached) so a retry
# after a mid-way failure naturally resumes.
#
# Never touches the `fpgm` conda env.
set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
TRELLIS_DIR="${REPO_ROOT}/third_party/TRELLIS.2"
ENV_NAME="trellis2"
CONDA_ROOT="${CONDA_ROOT:-/home/quang/miniconda3}"
ENV_PY="${CONDA_ROOT}/envs/${ENV_NAME}/bin/python"
UV_BIN="${UV_BIN:-/home/quang/.local/bin/uv}"
EXT_DIR="${REPO_ROOT}/.trellis_ext_build"
LOG_DIR="${REPO_ROOT}/logs/trellis_setup"
BUILD_SHIM_DIR="${TRELLIS_DIR}/.build_shim"
ATTN_BACKEND_FILE="${REPO_ROOT}/.trellis_attn_backend"

UTILS3D_COMMIT="9a4eb15e4021b67b12c460c7057d642626897ec8"  # pinned in TRELLIS.2/setup.sh
FLASH_ATTN_VERSION="2.7.3"                                   # pinned in TRELLIS.2/setup.sh
FLASH_ATTN_WHEEL_URL="https://github.com/Dao-AILab/flash-attention/releases/download/v${FLASH_ATTN_VERSION}/flash_attn-${FLASH_ATTN_VERSION}%2Bcu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"
XFORMERS_VERSION="0.0.29.post2"                               # pinned to torch==2.6.0 exactly (checked on PyPI)

mkdir -p "${EXT_DIR}" "${LOG_DIR}" "${BUILD_SHIM_DIR}"

echo "==> [0/9] sanity checks"
if [ ! -x "${ENV_PY}" ]; then
    echo "ERROR: conda env '${ENV_NAME}' python not found at ${ENV_PY}." >&2
    echo "       This script assumes the env already exists (python 3.10) with" >&2
    echo "       torch==2.6.0+cu124 / torchvision==0.21.0+cu124 installed." >&2
    exit 1
fi
if [ ! -d "${TRELLIS_DIR}/.git" ]; then
    echo "ERROR: TRELLIS.2 not cloned at ${TRELLIS_DIR}." >&2
    echo "       Run: git clone -b main https://github.com/microsoft/TRELLIS.2.git --recursive ${TRELLIS_DIR}" >&2
    exit 1
fi

PY() { "${ENV_PY}" "$@"; }
UV_INSTALL() { "${UV_BIN}" pip install --python "${ENV_PY}" "$@"; }

echo "==> [1/9] environment setup (CUDA_HOME, PATH, nvidia shim)"
export CUDA_HOME="${CONDA_ROOT}/envs/cuda128"
export PATH="${CONDA_ROOT}/envs/buildtools/bin:${CUDA_HOME}/bin:${PATH}"
# This conda cuda128 env's $CUDA_HOME/include is NOT cuda.h's home (it's
# binutils headers dropped in by another package); the real CUDA headers
# (cuda.h, cuda_runtime.h, ...) live under targets/x86_64-linux/include.
# nvcc/gcc both honor CPATH, so add it there rather than relying on
# torch.utils.cpp_extension's default -I$CUDA_HOME/include.
export CPATH="${CUDA_HOME}/targets/x86_64-linux/include:${CPATH:-}"
# Several CUDA math-library headers/libs (cusparse, cusolver, cublas, ...)
# aren't present in the cuda128 conda env at all (it's a slim nvcc+cudart
# env) but ARE bundled in the pip nvidia-*-cu12 packages that came in as
# torch's own dependencies. Fold every such package's include/lib dir into
# the build search paths so extensions that reach for e.g. cusparse.h find
# it without needing a full separate CUDA toolkit install.
NVIDIA_PIP_ROOT="$("${ENV_PY}" -c "import nvidia, os; print(os.path.dirname(nvidia.__file__))" 2>/dev/null || true)"
if [ -n "${NVIDIA_PIP_ROOT}" ] && [ -d "${NVIDIA_PIP_ROOT}" ]; then
    for pkg_dir in "${NVIDIA_PIP_ROOT}"/*/; do
        [ -d "${pkg_dir}include" ] && export CPATH="${pkg_dir}include:${CPATH}"
        [ -d "${pkg_dir}lib" ] && export LIBRARY_PATH="${pkg_dir}lib:${LIBRARY_PATH:-}" && export LD_LIBRARY_PATH="${pkg_dir}lib:${LD_LIBRARY_PATH:-}"
    done
    echo "    folded nvidia-*-cu12 pip package include/lib dirs from ${NVIDIA_PIP_ROOT} into CPATH/LIBRARY_PATH"
fi
# H200 NVL everywhere on this host -- pin the arch list so extension builds
# don't waste time on other archs and so nvdiffrec's setup.py (which sets
# TORCH_CUDA_ARCH_LIST='' when unset, breaking auto-detection) gets a real
# value.
export TORCH_CUDA_ARCH_LIST="9.0"
export MAX_JOBS="${MAX_JOBS:-32}"
# shellcheck source=/dev/null
source "${REPO_ROOT}/scripts/nvidia_lib_shim.sh"
# conda's compiler_compat/ld does not search /usr/lib/x86_64-linux-gnu by
# default, so `-lcuda` (needed by nvdiffrec's link step) fails to resolve
# even though libcuda.so is right there. The shim dir just sourced already
# has a correctly-versioned libcuda.so symlink; make it usable at link time
# too, not just at runtime via LD_LIBRARY_PATH.
export LIBRARY_PATH="${REPO_ROOT}/.nvshim:${LIBRARY_PATH:-}"

# cuda128 (nvcc-only conda env) ships no libnvrtc, but nvdiffrec's setup.py
# link-flags `-lnvrtc`. The trellis2 env's own `nvidia-cuda-nvrtc-cu12` pip
# package (pulled in as a torch dependency) has libnvrtc.so.12 but no
# unversioned dev symlink, which `-lnvrtc` needs at link time. Provide one in
# a local shim dir instead of touching the pip package or system dirs.
NVRTC_LIB_DIR="$(PY -c "import os,glob; import nvidia.cuda_nvrtc as m; print(os.path.dirname(glob.glob(os.path.join(os.path.dirname(m.__file__), 'lib', 'libnvrtc.so.*'))[0]))" 2>/dev/null || true)"
if [ -n "${NVRTC_LIB_DIR}" ]; then
    NVRTC_SO="$(ls "${NVRTC_LIB_DIR}"/libnvrtc.so.[0-9]* 2>/dev/null | head -1)"
    if [ -n "${NVRTC_SO}" ]; then
        ln -sf "${NVRTC_SO}" "${BUILD_SHIM_DIR}/libnvrtc.so"
        echo "    shimmed libnvrtc.so -> ${NVRTC_SO}"
    fi
    export LIBRARY_PATH="${BUILD_SHIM_DIR}:${NVRTC_LIB_DIR}:${LIBRARY_PATH:-}"
    export LD_LIBRARY_PATH="${NVRTC_LIB_DIR}:${LD_LIBRARY_PATH:-}"
else
    echo "    WARNING: could not locate nvidia-cuda-nvrtc-cu12's libnvrtc; nvdiffrec build may fail to link." >&2
fi

echo "==> [2/9] verify torch + CUDA in '${ENV_NAME}'"
PY - <<'PYEOF'
import torch
print("torch", torch.__version__, "| built for cuda", torch.version.cuda)
assert torch.__version__.startswith("2.6.0"), f"unexpected torch version {torch.__version__}"
assert torch.cuda.is_available(), "torch.cuda.is_available() is False even with the nvidia shim sourced"
print("CUDA OK:", torch.cuda.get_device_name(0))
PYEOF

echo "==> [3/9] basic python dependencies"
UV_INSTALL \
    imageio imageio-ffmpeg tqdm easydict opencv-python-headless ninja trimesh \
    transformers tensorboard pandas lpips zstandard kornia timm rembg onnxruntime \
    huggingface_hub plyfile \
    2>&1 | tee "${LOG_DIR}/01_basics.log"

echo "==> [4/9] utils3d (pinned commit ${UTILS3D_COMMIT})"
UV_INSTALL "git+https://github.com/EasternJournalist/utils3d.git@${UTILS3D_COMMIT}" \
    2>&1 | tee "${LOG_DIR}/02_utils3d.log"

clone_or_reuse() {
    local url="$1" dir="$2" branch="${3:-}"
    if [ -d "${dir}/.git" ]; then
        echo "    reusing existing clone at ${dir}"
    else
        rm -rf "${dir}"
        if [ -n "${branch}" ]; then
            git clone --recursive -b "${branch}" "${url}" "${dir}"
        else
            git clone --recursive "${url}" "${dir}"
        fi
    fi
}

echo "==> [5/9] compiled extensions: FlexGEMM, CuMesh, nvdiffrast, nvdiffrec, o-voxel"

echo "--> FlexGEMM"
clone_or_reuse https://github.com/JeffreyXiang/FlexGEMM.git "${EXT_DIR}/FlexGEMM"
UV_INSTALL --no-build-isolation "${EXT_DIR}/FlexGEMM" 2>&1 | tee "${LOG_DIR}/03_flexgemm.log"

echo "--> CuMesh"
clone_or_reuse https://github.com/JeffreyXiang/CuMesh.git "${EXT_DIR}/CuMesh"
UV_INSTALL --no-build-isolation "${EXT_DIR}/CuMesh" 2>&1 | tee "${LOG_DIR}/04_cumesh.log"

echo "--> nvdiffrast (v0.4.0 -- GL removed upstream, CUDA rasterizer only, which is all TRELLIS.2 uses)"
clone_or_reuse https://github.com/NVlabs/nvdiffrast.git "${EXT_DIR}/nvdiffrast" v0.4.0
UV_INSTALL --no-build-isolation "${EXT_DIR}/nvdiffrast" 2>&1 | tee "${LOG_DIR}/05_nvdiffrast.log"

echo "--> nvdiffrec (renderutils branch)"
clone_or_reuse https://github.com/JeffreyXiang/nvdiffrec.git "${EXT_DIR}/nvdiffrec" renderutils
UV_INSTALL --no-build-isolation "${EXT_DIR}/nvdiffrec" 2>&1 | tee "${LOG_DIR}/06_nvdiffrec.log"

echo "--> o-voxel (repo-local, --no-deps: cumesh/flex_gemm already installed above)"
UV_INSTALL --no-build-isolation --no-deps "${TRELLIS_DIR}/o-voxel" 2>&1 | tee "${LOG_DIR}/07_ovoxel.log"

echo "==> [6/9] attention backend (flash-attn preferred, xformers fallback)"
ATTN_BACKEND=""

echo "--> trying prebuilt flash-attn wheel for torch2.6/cu12/cxx11abiFALSE/cp310"
if UV_INSTALL "${FLASH_ATTN_WHEEL_URL}" 2>&1 | tee "${LOG_DIR}/08a_flash_attn_wheel.log"; then
    if PY -c "import flash_attn" 2>/dev/null; then
        ATTN_BACKEND="flash_attn"
        echo "    flash-attn installed from prebuilt wheel."
    fi
fi

if [ -z "${ATTN_BACKEND}" ]; then
    echo "--> prebuilt wheel unavailable/failed; trying flash-attn==${FLASH_ATTN_VERSION} (may source-build, bounded to 40 min)"
    if timeout 2400 env FLASH_ATTENTION_FORCE_BUILD=FALSE bash -c \
        "'${UV_BIN}' pip install --python '${ENV_PY}' --no-build-isolation 'flash-attn==${FLASH_ATTN_VERSION}'" \
        2>&1 | tee "${LOG_DIR}/08b_flash_attn_build.log"; then
        if PY -c "import flash_attn" 2>/dev/null; then
            ATTN_BACKEND="flash_attn"
            echo "    flash-attn installed (build or fetched by its own setup.py)."
        fi
    fi
fi

if [ -z "${ATTN_BACKEND}" ]; then
    echo "--> flash-attn unavailable; falling back to xformers==${XFORMERS_VERSION} (pinned to torch==2.6.0 exactly)"
    if UV_INSTALL "xformers==${XFORMERS_VERSION}" 2>&1 | tee "${LOG_DIR}/09_xformers.log"; then
        TORCH_VER_AFTER="$(PY -c 'import torch; print(torch.__version__)')"
        if [[ "${TORCH_VER_AFTER}" != 2.6.0* ]]; then
            echo "ERROR: installing xformers changed torch to ${TORCH_VER_AFTER} -- aborting to protect the pinned torch build." >&2
            exit 1
        fi
        if PY -c "import xformers" 2>/dev/null; then
            ATTN_BACKEND="xformers"
            echo "    xformers installed; torch still ${TORCH_VER_AFTER}."
        fi
    fi
fi

if [ -z "${ATTN_BACKEND}" ]; then
    echo "ERROR: neither flash-attn nor xformers could be installed/imported. Stopping" >&2
    echo "       per instructions rather than silently degrading to a slow/broken attention path." >&2
    echo "       See ${LOG_DIR}/08a_flash_attn_wheel.log, 08b_flash_attn_build.log, 09_xformers.log." >&2
    exit 1
fi

echo "${ATTN_BACKEND}" > "${ATTN_BACKEND_FILE}"
echo "    attention backend selected: ${ATTN_BACKEND} (recorded in ${ATTN_BACKEND_FILE})"

echo "==> [7/9] smoke test: import trellis2, o_voxel, cumesh, flex_gemm, nvdiffrast.torch"
# trellis2's repo root has no pyproject.toml -- it's meant to be run with the
# checkout itself on sys.path (exactly how TRELLIS.2's own example.py works),
# not pip-installed. scripts/trellis_generate.py does the same sys.path
# insertion at runtime.
SPARSE_ATTN_BACKEND="${ATTN_BACKEND}" ATTN_BACKEND="${ATTN_BACKEND}" PY - <<PYEOF
import sys
sys.path.insert(0, "${TRELLIS_DIR}")
import importlib
mods = ["torch", "trellis2", "o_voxel", "cumesh", "flex_gemm", "nvdiffrast.torch", "nvdiffrec_render.renderutils"]
for name in mods:
    m = importlib.import_module(name)
    ver = getattr(m, "__version__", None)
    print(f"{name}: OK" + (f" (version {ver})" if ver else ""))
import torch
print("torch.cuda.is_available():", torch.cuda.is_available())
PYEOF

echo "==> [8/9] verify fpgm env's torch is untouched"
FPGM_PY="${CONDA_ROOT}/envs/fpgm/bin/python"
if [ -x "${FPGM_PY}" ]; then
    FPGM_TORCH="$("${FPGM_PY}" -c 'import torch; print(torch.__version__)')"
    echo "    fpgm torch: ${FPGM_TORCH}"
    if [[ "${FPGM_TORCH}" != 2.10.0* ]]; then
        echo "ERROR: fpgm's torch is no longer 2.10.0+cu128 (found ${FPGM_TORCH})! This must not happen." >&2
        exit 1
    fi
else
    echo "    WARNING: fpgm env python not found at ${FPGM_PY}, skipping check." >&2
fi

echo "==> [9/9] done."
echo "    Attention backend: ${ATTN_BACKEND} (export SPARSE_ATTN_BACKEND=${ATTN_BACKEND} before running the pipeline)"
echo "    Generate a mesh with:"
echo "      source ${REPO_ROOT}/scripts/nvidia_lib_shim.sh"
echo "      ${ENV_PY} ${REPO_ROOT}/scripts/trellis_generate.py --image <rgba.png> --out <out.glb>"