#!/usr/bin/env bash
set -Eeuo pipefail

# Downloads distilled FLUX.2 Klein 9B model files into current directory.
#
# Usage:
#   HF_TOKEN=hf_xxx ./download_models.sh

if [[ -z "${HF_TOKEN:-}" ]]; then
  echo "ERROR: HF_TOKEN is required."
  exit 1
fi

# Reduce open-files/Xet issues seen in some containers.
export HF_HUB_DISABLE_XET=1
export HF_HUB_ENABLE_HF_TRANSFER=0
ulimit -n 65535 2>/dev/null || true

if ! command -v uv >/dev/null 2>&1; then
  echo "uv not found, installing..."
  curl -LsSf https://astral.sh/uv/install.sh | sh
  export PATH="$HOME/.local/bin:$PATH"
fi

if ! command -v uv >/dev/null 2>&1; then
  echo "ERROR: uv not available after install."
  exit 1
fi

TMP_DIR="${PWD}/.hf_tmp_downloads"
mkdir -p "$TMP_DIR"

download_to_cwd() {
  local repo="$1"
  local file="$2"
  local out_name
  local downloaded_path
  local tries=3
  local i

  out_name="$(basename "$file")"

  for i in $(seq 1 "$tries"); do
    echo "Downloading ($i/$tries): $repo :: $file"
    if downloaded_path="$(uvx --from "huggingface_hub[cli]>=0.31.0" hf download "$repo" "$file" --token "$HF_TOKEN" --local-dir "$TMP_DIR")"; then
      cp -f "$downloaded_path" "${PWD}/${out_name}"
      echo "Saved: ${PWD}/${out_name}"
      return 0
    fi
    sleep 5
  done

  echo "ERROR: Failed after $tries attempts: $repo / $file"
  return 1
}

download_to_cwd "black-forest-labs/FLUX.2-klein-9B" "flux-2-klein-9b.safetensors"
download_to_cwd "Comfy-Org/vae-text-encorder-for-flux-klein-9b" "split_files/text_encoders/qwen_3_8b.safetensors"
download_to_cwd "Comfy-Org/vae-text-encorder-for-flux-klein-9b" "split_files/vae/flux2-vae.safetensors"

echo
echo "Done. Files in cwd:"
ls -lh "${PWD}/flux-2-klein-9b.safetensors" "${PWD}/qwen_3_8b.safetensors" "${PWD}/flux2-vae.safetensors"