arbor/setup.sh at main · Stability-AI/arbor · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
#!/usr/bin/env bash
set -euo pipefail

HELP=false
NEW_ENV=false
BASIC=false
XFORMERS=false
FLASHATTN=false
SPCONV=false
OVOXEL=true

while [[ $# -gt 0 ]]; do
  case "$1" in
    -h|--help) HELP=true; shift ;;
    --new-env) NEW_ENV=true; shift ;;
    --basic) BASIC=true; shift ;;
    --xformers) XFORMERS=true; shift ;;
    --flash-attn) FLASHATTN=true; shift ;;
    --spconv) SPCONV=true; shift ;;
    --no-o-voxel) OVOXEL=false; shift ;;
    *) echo "Unknown argument: $1" >&2; HELP=true; shift ;;
  esac
done

if [[ "$HELP" == true ]]; then
  cat <<'HELP_TEXT'
Usage: setup.sh [OPTIONS]
  --new-env       Create and activate conda env arbor with Python 3.10 and PyTorch CUDA 11.8.
  --basic         Install Python inference dependencies and this package in editable mode.
  --xformers      Install xformers matching PyTorch/CUDA when supported.
  --flash-attn    Install flash-attn.
  --spconv        Install spconv matching CUDA.
  --no-o-voxel    Skip building the vendored o_voxel extension.
HELP_TEXT
  exit 0
fi

if [[ "$NEW_ENV" == true ]]; then
  conda create -n arbor python=3.10 -y
  source "$(conda info --base)/etc/profile.d/conda.sh"
  conda activate arbor
  conda install pytorch==2.4.0 torchvision==0.19.0 pytorch-cuda=11.8 -c pytorch -c nvidia -y
fi

PYTORCH_VERSION=$(python -c "import torch; print(torch.__version__)")
PLATFORM=$(python -c "import torch; print('cuda' if torch.cuda.is_available() and torch.version.cuda else 'cpu')")
CUDA_VERSION=$(python -c "import torch; print(torch.version.cuda or '')")
CUDA_MAJOR_VERSION=${CUDA_VERSION%%.*}
CUDA_MINOR_VERSION=$(python -c "import torch; v=torch.version.cuda or ''; print(v.split('.')[1] if '.' in v else '')")

echo "[SYSTEM] PyTorch=${PYTORCH_VERSION} CUDA=${CUDA_VERSION:-none}"

if [[ "$BASIC" == true ]]; then
  pip install --upgrade pip
  pip install -r requirements.txt
  pip install git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8
  pip install nvdiffrast==0.3.3 || pip install git+https://github.com/NVlabs/nvdiffrast.git@v0.3.3
  pip install -e .
fi

if [[ "$XFORMERS" == true ]]; then
  if [[ "$PLATFORM" == "cuda" && "$CUDA_VERSION" == "11.8" && "$PYTORCH_VERSION" == 2.4.0* ]]; then
    pip install xformers==0.0.27.post2 --index-url https://download.pytorch.org/whl/cu118
  else
    echo "[XFORMERS] Install manually for PyTorch=${PYTORCH_VERSION} CUDA=${CUDA_VERSION}."
  fi
fi

if [[ "$FLASHATTN" == true ]]; then
  pip install flash-attn
fi

if [[ "$SPCONV" == true ]]; then
  if [[ "$PLATFORM" == "cuda" && "$CUDA_MAJOR_VERSION" == "11" ]]; then
    pip install spconv-cu118
  elif [[ "$PLATFORM" == "cuda" && "$CUDA_MAJOR_VERSION" == "12" ]]; then
    if [[ "$CUDA_MINOR_VERSION" == "1" ]]; then
      pip install spconv-cu121
    elif [[ "$CUDA_MINOR_VERSION" == "4" ]]; then
      pip install spconv-cu124
    else
      pip install spconv-cu120
    fi
  else
    echo "[SPCONV] CUDA is required for Arbor inference."
  fi
fi

if [[ "$OVOXEL" == true && -d "third_party/o_voxel" ]]; then
  echo "[GEOM] Building o_voxel..."
  pip install --no-build-isolation --no-deps ./third_party/o_voxel
fi