#!/usr/bin/env bash
# /users/jw02425/models/SSL/vit/download.sh
# Download iBOT checkpoints and save as: ibot_{arch}_in1k(.pth)
set -euo pipefail

OUT_DIR=${1:-.}
mkdir -p "$OUT_DIR"

declare -A urls
urls["ibot_vits_in1k"]="https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/archive/2022/ibot/vits_16/checkpoint_teacher.pth"
urls["ibot_vitb_in1k"]="https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/archive/2022/ibot/vitb_16/checkpoint_teacher.pth"
# urls["ibot_vitl_in1k"]="https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/archive/2022/ibot/vitl_16/checkpoint_teacher.pth"
urls["ibot_vitb_in1krand"]="https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/archive/2022/ibot/vitb_16_rand_mask/checkpoint_teacher.pth"
# urls["ibot_vitl_in1krand"]="https://lf3-nlp-opensource.bytetos.com/obj/nlp-opensource/archive/2022/ibot/vitl_16_rand_mask/checkpoint_teacher.pth"

# moco
urls["moco_vits_in1k"]="https://dl.fbaipublicfiles.com/moco-v3/vit-s-300ep/vit-s-300ep.pth.tar"
urls["moco_vitb_in1k"]="https://dl.fbaipublicfiles.com/moco-v3/vit-b-300ep/vit-b-300ep.pth.tar"



for name in "${!urls[@]}"; do
    url="${urls[$name]}"
    out_path="$OUT_DIR/${name}.pth"
    if [[ ! -f "$out_path" ]]; then
        echo "File $out_path already exists, skipping download."
        continue
    fi
    echo "Downloading $url -> $out_path"
    wget -c -O "$out_path" "$url"
done