File size: 1,571 Bytes
0faef79 |
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 |
#!/bin/bash
set -eux
# Default model name
model_name="MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth"
check_docker() {
if ! command -v docker &>/dev/null; then
echo "Docker could not be found. Please install Docker and try again."
exit 1
fi
}
download_model_checkpoint() {
if [ -f "./files/checkpoints/${model_name}" ]; then
echo "Model checkpoint ${model_name} already exists. Skipping download."
return
fi
echo "Downloading model checkpoint ${model_name}..."
wget "https://download.europe.naverlabs.com/ComputerVision/MASt3R/${model_name}" -P ./files/checkpoints
}
set_dcomp() {
if command -v docker-compose &>/dev/null; then
dcomp="docker-compose"
elif command -v docker &>/dev/null && docker compose version &>/dev/null; then
dcomp="docker compose"
else
echo "Docker Compose could not be found. Please install Docker Compose and try again."
exit 1
fi
}
run_docker() {
export MODEL=${model_name}
if [ "$with_cuda" -eq 1 ]; then
$dcomp -f docker-compose-cuda.yml up --build
else
$dcomp -f docker-compose-cpu.yml up --build
fi
}
with_cuda=0
for arg in "$@"; do
case $arg in
--with-cuda)
with_cuda=1
;;
--model_name=*)
model_name="${arg#*=}.pth"
;;
*)
echo "Unknown parameter passed: $arg"
exit 1
;;
esac
done
main() {
check_docker
download_model_checkpoint
set_dcomp
run_docker
}
main
|