AntonioJun/workspace

1#!/usr/bin/env bash2# One-shot setup for this workspace.3#4# Installs packages required by every Python file under /workspace, preserves the5# project layout expected by the code, and downloads/clones the external data,6# repos, and checkpoints used at runtime.7#8# Usage:9#   ./setup.sh10#   HF_TOKEN=hf_xxx ./setup.sh -y11#   ./setup.sh --skip-models12#   ./setup.sh --skip-data13#   ./setup.sh --skip-workspace14#   ./setup.sh --with-caches15#   ./setup.sh --with-segvggt16#   ./setup.sh --force17 18set -euo pipefail19 20WORKSPACE_ROOT="${VSI_WORKSPACE_ROOT:-/workspace}"21DATA_ROOT="${VSI_DATA_ROOT:-/root/data}"22MODELS_ROOT="${VSI_MODELS_ROOT:-/root/models}"23VENV_ROOT="${VSI_VENV_ROOT:-/root/.venv}"24PYTHON_BIN="${VSI_PYTHON_BIN:-python3.11}"25WORKSPACE_REPO="${VSI_WORKSPACE_REPO:-AntonioJun/workspace}"26 27SKIP_MODELS=028SKIP_DATA=029SKIP_WORKSPACE=030WITH_CACHES=031WITH_SEGVGGT=032FORCE=033ASSUME_YES=034HF_TOKEN="${HF_TOKEN:-${HUGGING_FACE_HUB_TOKEN:-}}"35 36for arg in "$@"; do37  case "$arg" in38    --skip-models) SKIP_MODELS=1 ;;39    --skip-data) SKIP_DATA=1 ;;40    --skip-workspace) SKIP_WORKSPACE=1 ;;41    --with-caches) WITH_CACHES=1 ;;42    --with-segvggt) WITH_SEGVGGT=1 ;;43    --force) FORCE=1 ;;44    -y|--yes) ASSUME_YES=1 ;;45    --token=*) HF_TOKEN="${arg#--token=}" ;;46    -h|--help) awk 'NR == 1 {next} /^#/ {sub(/^# ?/, ""); print; next} {exit}' "$0"; exit 0 ;;47    *) echo "unknown argument: $arg" >&2; exit 1 ;;48  esac49done50 51log() { printf '\n==> %s\n' "$1"; }52warn() { printf '!! %s\n' "$1" >&2; }53die() { printf 'XX %s\n' "$1" >&2; exit 1; }54 55 56if [ -z "$HF_TOKEN" ] && { [ "$SKIP_DATA" -eq 0 ] || [ "$SKIP_MODELS" -eq 0 ] || [ "$SKIP_WORKSPACE" -eq 0 ]; }; then57  if [ "$ASSUME_YES" -eq 1 ]; then58    warn "HF_TOKEN is empty; public downloads may work, gated downloads will fail"59  else60    log "Hugging Face token requested for gated/model/dataset downloads"61    echo "Create/read one at https://huggingface.co/settings/tokens."62    echo "You also need access to facebook/sam3 and nyu-visionx/VSI-Bench where applicable."63    read -r -s -p "HF token (input hidden, blank to continue without one): " HF_TOKEN64    echo65  fi66fi67export HF_TOKEN68export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"69export HF_HUB_ENABLE_HF_TRANSFER=170export HF_XET_HIGH_PERFORMANCE=171 72log "Checking system packages"73missing=()74for bin in git curl unzip "$PYTHON_BIN"; do75  command -v "$bin" >/dev/null 2>&1 || missing+=("$bin")76done77if [ "${#missing[@]}" -gt 0 ]; then78  if command -v apt-get >/dev/null 2>&1; then79    apt-get update -qq80    apt-get install -y -qq git curl unzip python3.11 python3.11-venv python3.11-dev build-essential ffmpeg81  else82    die "missing required system tools: ${missing[*]}"83  fi84fi85mkdir -p "$DATA_ROOT" "$DATA_ROOT/caches" "$DATA_ROOT/spatial codes" "$MODELS_ROOT"86 87clone_repo() {88  local url="$1" dest="$2"89  if [ -d "$dest/.git" ] && [ "$FORCE" -eq 0 ]; then90    log "Already cloned: $dest"91    return92  fi93  log "Cloning $url -> $dest"94  rm -rf "$dest"95  git clone --depth 1 "$url" "$dest"96}97 98create_venv() {99  if [ -d "$VENV_ROOT" ] && [ "$FORCE" -eq 0 ]; then100    log "Using existing venv: $VENV_ROOT"101  else102    log "Creating venv: $VENV_ROOT"103    rm -rf "$VENV_ROOT"104    "$PYTHON_BIN" -m venv "$VENV_ROOT"105  fi106  "$VENV_ROOT/bin/pip" install --upgrade -q pip wheel "setuptools<81"107}108 109install_workspace_requirements() {110  log "Installing Python packages used by every /workspace Python file"111  "$VENV_ROOT/bin/pip" install -q -r /dev/stdin <<'REQS'112numpy<2113scipy114pandas115PyYAML116loguru117datasets118Pillow119opencv-python==4.11.0.86120opencv-contrib-python-headless==4.10.0.84121torch122torchvision123accelerate==1.14.0124transformers==5.14.1125huggingface_hub[cli]>=0.24126hf_transfer127safetensors128timm129einops130sentencepiece131protobuf132av133imageio134pycocotools135hydra-core136omegaconf137pytest>=8.3.5138black139REQS140}141 142install_external_repos() {143  log "Cloning model source repositories under $MODELS_ROOT"144  clone_repo "https://github.com/facebookresearch/sam3.git" "$MODELS_ROOT/sam3"145  clone_repo "https://github.com/bytedance-seed/depth-anything-3.git" "$MODELS_ROOT/depth-anything-3"146  if [ "$WITH_SEGVGGT" -eq 1 ]; then147    clone_repo "https://github.com/Seed3D/SegVGGT.git" "$MODELS_ROOT/SegVGGT"148  fi149 150  log "Installing editable model repos where present"151  if [ -f "$MODELS_ROOT/sam3/pyproject.toml" ] || [ -f "$MODELS_ROOT/sam3/setup.py" ]; then152    "$VENV_ROOT/bin/pip" install -q -e "$MODELS_ROOT/sam3"153  else154    warn "SAM3 repo has no pyproject.toml/setup.py at $MODELS_ROOT/sam3; skipped editable install"155  fi156  if [ -f "$MODELS_ROOT/depth-anything-3/pyproject.toml" ] || [ -f "$MODELS_ROOT/depth-anything-3/setup.py" ]; then157    "$VENV_ROOT/bin/pip" install -q -e "$MODELS_ROOT/depth-anything-3"158  else159    warn "DA3 repo has no pyproject.toml/setup.py at $MODELS_ROOT/depth-anything-3; skipped editable install"160  fi161  if [ "$WITH_SEGVGGT" -eq 1 ] && [ -f "$MODELS_ROOT/SegVGGT/requirements.txt" ]; then162    "$VENV_ROOT/bin/pip" install -q -r "$MODELS_ROOT/SegVGGT/requirements.txt"163  fi164}165 166hf_cli() {167  if [ -x "$VENV_ROOT/bin/hf" ]; then168    printf '%s169' "$VENV_ROOT/bin/hf"170  elif [ -x "$VENV_ROOT/bin/huggingface-cli" ]; then171    printf '%s172' "$VENV_ROOT/bin/huggingface-cli"173  else174    command -v hf || command -v huggingface-cli || true175  fi176}177 178hf_download() {179  local repo_id="$1" repo_type="$2" dest="$3" optional="${4:-}"180  if [ -d "$dest" ] && [ "$(ls -A "$dest" 2>/dev/null)" ] && [ "$FORCE" -eq 0 ]; then181    log "Already present: $dest"182    return183  fi184  local cli185  cli="$(hf_cli)"186  [ -n "$cli" ] || die "no hf/huggingface-cli command available"187  log "Downloading $repo_id ($repo_type) -> $dest"188  mkdir -p "$dest"189  token_args=()190  [ -n "$HF_TOKEN" ] && token_args=(--token "$HF_TOKEN")191  if ! "$cli" download "$repo_id" --repo-type "$repo_type" --local-dir "$dest" "${token_args[@]}"; then192    if [ "$optional" = "optional" ]; then193      warn "download failed but marked optional: $repo_id"194    else195      die "download failed: $repo_id"196    fi197  fi198}199 200sync_workspace_from_backup() {201  [ "$SKIP_WORKSPACE" -eq 0 ] || return 0202  log "Syncing workspace files from Hugging Face dataset repo: $WORKSPACE_REPO"203  WITH_CACHES="$WITH_CACHES" WORKSPACE_REPO="$WORKSPACE_REPO" WORKSPACE_ROOT="$WORKSPACE_ROOT" \204    "$VENV_ROOT/bin/python" - <<'PYSYNC'205import os206import tarfile207from pathlib import Path208from huggingface_hub import HfApi, hf_hub_download209 210repo = os.environ["WORKSPACE_REPO"]211root = Path(os.environ["WORKSPACE_ROOT"])212token = os.environ.get("HF_TOKEN") or None213api = HfApi(token=token)214folders = ["harness", "symbolic", "analysis", "corruption", "calibration", "encoder", "inference", "tests"]215wanted_spatial = (216    "data/spatial codes/sam3+depth-anything-3/metric/tracking/selective/32/",217    "data/spatial codes/sam3+depth-anything-3/metric/tracking/uniform/32/",218    "data/spatial codes/ground truth/explicit/",219    "data/spatial codes/ground truth/compact/",220)221try:222    bundle = hf_hub_download(repo, "bundles/spatial-codes.tar.gz", repo_type="dataset", token=token)223    with tarfile.open(bundle) as tar:224        members = [m for m in tar.getmembers() if any(m.name.startswith(w) for w in wanted_spatial)]225        tar.extractall(root, members=members)226    print(f"[data/spatial codes] {len(members)} files extracted from bundle", flush=True)227except Exception as exc:228    print(f"[data/spatial codes] bundle unavailable ({exc}); using per-file fallback", flush=True)229    folders.extend(w.rstrip("/") for w in wanted_spatial)230if os.environ.get("WITH_CACHES") == "1":231    folders.append("data/caches")232 233def fetch(rel):234    dest = root / rel235    if dest.is_file() and dest.stat().st_size > 0:236        return 0237    hf_hub_download(repo, rel, repo_type="dataset", local_dir=str(root), token=token)238    return 1239 240for folder in folders:241    files = [242        e.path243        for e in api.list_repo_tree(repo, repo_type="dataset", path_in_repo=folder, recursive=True)244        if e.__class__.__name__ == "RepoFile"245    ]246    got = sum(fetch(rel) for rel in files)247    print(f"[{folder}] {len(files)} files ({got} downloaded, rest already present)", flush=True)248for rel in ["README.md", "backup.py", "selective_frame_counts.csv"]:249    try:250        fetch(rel)251    except Exception as exc:252        print(f"[{rel}] skipped: {exc}", flush=True)253print("workspace sync complete", flush=True)254PYSYNC255  chmod +x "$WORKSPACE_ROOT/setup.sh" "$WORKSPACE_ROOT/backup.py" 2>/dev/null || true256}257 258install_data() {259  [ "$SKIP_DATA" -eq 0 ] || return 0260  log "Cloning thinking-in-space under $DATA_ROOT"261  clone_repo "https://github.com/vision-x-nyu/thinking-in-space.git" "$DATA_ROOT/thinking-in-space"262  hf_download "nyu-visionx/VSI-Bench" dataset "$DATA_ROOT/VSI-Bench"263  log "Extracting VSI-Bench scene archives"264  for name in scannet arkitscenes scannetpp; do265    zip_path="$DATA_ROOT/VSI-Bench/${name}.zip"266    out_dir="$DATA_ROOT/VSI-Bench/${name}"267    if [ -f "$zip_path" ] && { [ ! -d "$out_dir" ] || [ "$FORCE" -eq 1 ]; }; then268      unzip -q -o "$zip_path" -d "$DATA_ROOT/VSI-Bench"269    fi270  done271}272 273install_models() {274  [ "$SKIP_MODELS" -eq 0 ] || return 0275  log "Downloading model checkpoints under $MODELS_ROOT"276  hf_download "facebook/sam3" model "$MODELS_ROOT/sam3/checkpoints" optional277  hf_download "depth-anything/DA3-LARGE-1.1" model "$MODELS_ROOT/depth-anything-3/checkpoints/DA3-LARGE-1.1" optional278  hf_download "depth-anything/DA3NESTED-GIANT-LARGE-1.1" model "$MODELS_ROOT/depth-anything-3/checkpoints/DA3NESTED-GIANT-LARGE-1.1" optional279  hf_download "Qwen/Qwen3.5-4B" model "$MODELS_ROOT/qwen3.5-4b"280  hf_download "Qwen/Qwen3.5-2B" model "$MODELS_ROOT/qwen3.5-2b"281  hf_download "OpenGVLab/InternVL3_5-4B-HF" model "$MODELS_ROOT/internvl3.5-4b"282  hf_download "OpenGVLab/InternVL3_5-2B-HF" model "$MODELS_ROOT/internvl3.5-2b"283  if [ "$WITH_SEGVGGT" -eq 1 ]; then284    hf_download "Seed3D/SegVGGT" model "$MODELS_ROOT/SegVGGT/checkpoint" optional285  fi286}287 288write_env_file() {289  log "Writing environment helper: /root/vsi-env.sh"290  cat > /root/vsi-env.sh <<ENVEOF291export VSI_WORKSPACE_ROOT="$WORKSPACE_ROOT"292export VSI_DATA_ROOT="$DATA_ROOT"293export VSI_ROOT="$DATA_ROOT/VSI-Bench"294export VSI_CACHE_ROOT="$DATA_ROOT/caches"295export VSI_CODES="$DATA_ROOT/spatial codes"296export VSI_MODELS_ROOT="$MODELS_ROOT"297export VSI_SAM3_ROOT="$MODELS_ROOT/sam3"298export VSI_DA3_ROOT="$MODELS_ROOT/depth-anything-3"299export VSI_SEGVGGT_ROOT="$MODELS_ROOT/SegVGGT"300export HARNESS_OFFICIAL_EVAL="$DATA_ROOT/thinking-in-space/lmms_eval/tasks/vsibench/utils.py"301export SYMBOLIC_OFFICIAL_EVAL="$DATA_ROOT/thinking-in-space/lmms_eval/tasks/vsibench/utils.py"302export PYTHONPATH="$WORKSPACE_ROOT:\$PYTHONPATH"303ENVEOF304  cat > /root/.venv-map.json <<MAPEOF305{306  "mode": "shared",307  "venv": "$VENV_ROOT",308  "workspace": "$WORKSPACE_ROOT",309  "data_root": "$DATA_ROOT",310  "models_root": "$MODELS_ROOT"311}312MAPEOF313}314 315smoke_test_imports() {316  log "Running import smoke test for workspace dependency coverage"317  "$VENV_ROOT/bin/python" - <<'PYSMOKE'318import importlib319mods = [320    "numpy", "scipy", "pandas", "yaml", "loguru", "PIL.Image", "cv2",321    "torch", "torchvision", "transformers", "huggingface_hub", "hydra",322    "pytest", "black", "pycocotools.mask",323]324failed = []325for mod in mods:326    try:327        importlib.import_module(mod)328    except Exception as exc:329        failed.append(f"{mod}: {exc}")330if failed:331    print("IMPORT_SMOKE_FAILED")332    for item in failed:333        print(" -", item)334    raise SystemExit(1)335print("IMPORT_SMOKE_OK")336PYSMOKE337}338 339final_checks() {340  log "Running final workspace checks"341  (cd "$WORKSPACE_ROOT" && "$VENV_ROOT/bin/python" -m black --check .)342  (cd "$WORKSPACE_ROOT" && "$VENV_ROOT/bin/python" -m compileall -q .)343  (cd "$WORKSPACE_ROOT" && "$VENV_ROOT/bin/python" -m pytest -q) || warn "pytest did not fully pass; review output above"344}345 346create_venv347install_workspace_requirements348install_external_repos349sync_workspace_from_backup350install_data351install_models352write_env_file353smoke_test_imports354final_checks355 356log "Setup complete"357cat <<DONEEOF358 359Workspace:  $WORKSPACE_ROOT360Data:       $DATA_ROOT361Models:     $MODELS_ROOT362Venv:       $VENV_ROOT363Env helper: /root/vsi-env.sh364 365Activate with:366  source $VENV_ROOT/bin/activate367  source /root/vsi-env.sh368 369DONEEOF370