scripts/setup_env.sh

#!/usr/bin/env bash
# Creates the `fpgm` conda env and installs everything needed to run the
# pipeline, including two third_party checkouts that pip metadata alone can't
# express cleanly: sam3's own editable install, and tapnet's install *without*
# its declared deps (see step 7 below for why that's essential, not optional).
#
# Idempotent: safe to re-run after a partial failure. Never installs, deletes,
# or touches conda/pip config outside this one env.
set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
ENV_NAME="fpgm"
PYTHON_VERSION="3.12"
CONDA_ROOT="${CONDA_ROOT:-/home/quang/miniconda3}"
UV_BIN="${UV_BIN:-/home/quang/.local/bin/uv}"

# The kernel driver on this box is 570.x, which tops out at CUDA 12.8: a plain
# `pip install torch` would resolve a cu13x wheel that either fails at import
# or silently runs CPU-only. The index must be pinned explicitly.
TORCH_INDEX_URL="https://download.pytorch.org/whl/cu128"
TORCH_VERSION="2.10.0"

SAM3_REPO_URL="https://github.com/facebookresearch/sam3"
TAPNET_REPO_URL="https://github.com/google-deepmind/tapnet"
THIRD_PARTY_DIR="${REPO_ROOT}/third_party"
PINNED_COMMITS_FILE="${THIRD_PARTY_DIR}/PINNED_COMMITS.txt"

echo "==> [1/8] sourcing conda from ${CONDA_ROOT}"
# shellcheck source=/dev/null
source "${CONDA_ROOT}/etc/profile.d/conda.sh"

echo "==> [2/8] creating conda env '${ENV_NAME}' (python=${PYTHON_VERSION}) if missing"
if ! conda env list | awk '{print $1}' | grep -qx "${ENV_NAME}"; then
    conda create -y -n "${ENV_NAME}" "python=${PYTHON_VERSION}"
else
    echo "    env '${ENV_NAME}' already exists, skipping creation"
fi
conda activate "${ENV_NAME}"
PYTHON_BIN="$(command -v python)"

echo "==> [3/8] installing torch ${TORCH_VERSION} + torchvision from ${TORCH_INDEX_URL}"
"${UV_BIN}" pip install --python "${PYTHON_BIN}" \
    "torch==${TORCH_VERSION}" torchvision \
    --index-url "${TORCH_INDEX_URL}"

echo "==> [4/8] installing project dependencies (editable, with dev extras)"
"${UV_BIN}" pip install --python "${PYTHON_BIN}" -e "${REPO_ROOT}[dev]"

echo "==> [5/8] cloning third_party checkouts and pinning commit SHAs"
mkdir -p "${THIRD_PARTY_DIR}"
pins_tmp="${PINNED_COMMITS_FILE}.tmp"
: > "${pins_tmp}"

clone_and_pin() {
    local repo_url="$1" dest_name="$2"
    local dest="${THIRD_PARTY_DIR}/${dest_name}"
    if [ ! -d "${dest}/.git" ]; then
        git clone "${repo_url}" "${dest}"
    else
        echo "    ${dest_name} already cloned, skipping (not pulling -- keeps the pin reproducible)"
    fi
    local sha
    sha="$(git -C "${dest}" rev-parse HEAD)"
    echo "${dest_name} ${repo_url} ${sha}" >> "${pins_tmp}"
}

clone_and_pin "${SAM3_REPO_URL}" "sam3"
clone_and_pin "${TAPNET_REPO_URL}" "tapnet"
mv "${pins_tmp}" "${PINNED_COMMITS_FILE}"
echo "    pinned commits recorded in ${PINNED_COMMITS_FILE}"

echo "==> [6/8] installing sam3 (editable, full deps)"
"${UV_BIN}" pip install --python "${PYTHON_BIN}" -e "${THIRD_PARTY_DIR}/sam3"

echo "==> [7/8] installing tapnet (editable, --no-deps)"
# tapnet's pyproject.toml unconditionally lists jax/dm-haiku/jaxline/optax --
# the whole JAX training stack -- even though the TAPNext++ checkpoint this
# project uses runs through tapnet's torch reimplementation, which never
# imports any of them. Installing with deps would drag in a second,
# incompatible accelerator stack (and likely a CUDA/jaxlib version fight
# against the torch install from step 3) purely for code paths that are never
# executed here. --no-deps is essential, not an optimization.
"${UV_BIN}" pip install --python "${PYTHON_BIN}" --no-deps -e "${THIRD_PARTY_DIR}/tapnet"

echo "==> [8/8] done."
echo "    Activate with: conda activate ${ENV_NAME}"
echo "    If nvidia-smi / torch.cuda.is_available() reports a driver/library"
echo "    version mismatch, run: source ${REPO_ROOT}/scripts/nvidia_lib_shim.sh"