build.sh

#!/usr/bin/env bash
# Build UniRelight Blackwell image
set -euo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "${ROOT}"

IMAGE_NAME="${IMAGE_NAME:-unirelight-blackwell:cu129}"
CUDA_IMAGE="${CUDA_IMAGE:-nvidia/cuda:12.9.0-devel-ubuntu24.04}"
TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-12.0}"
UNIRELIGHT_REF="${UNIRELIGHT_REF:-main}"
MAX_JOBS="${MAX_JOBS:-$(nproc 2>/dev/null || echo 8)}"
USE_COMPOSE="${USE_COMPOSE:-0}"

ARGS=(
  --build-arg "CUDA_IMAGE=${CUDA_IMAGE}"
  --build-arg "TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}"
  --build-arg "UNIRELIGHT_REF=${UNIRELIGHT_REF}"
  --build-arg "MAX_JOBS=${MAX_JOBS}"
  --build-arg "INSTALL_FLASH_ATTN=${INSTALL_FLASH_ATTN:-1}"
  --build-arg "INSTALL_XFORMERS=${INSTALL_XFORMERS:-1}"
  --build-arg "INSTALL_APEX=${INSTALL_APEX:-1}"
  --build-arg "INSTALL_NVDIFFRAST=${INSTALL_NVDIFFRAST:-1}"
  --build-arg "UNIRELIGHT_PORT=8080"
)

echo "==> Building ${IMAGE_NAME}"
if [[ "${USE_COMPOSE}" == "1" ]]; then
  docker compose build "${ARGS[@]}"
else
  docker build -t "${IMAGE_NAME}" "${ARGS[@]}" -f Dockerfile .
fi
echo "==> Done. Launch: docker compose up -d → http://localhost:8080"