#!/usr/bin/env bash
# SPDX-License-Identifier: Apache-2.0
#
# Fetch what tt-MoGe's tests / demo need:
#   1. MoGe-2 ViT-L checkpoint (Ruicheng/moge-2-vitl-normal on HF, ~1.3 GB)
#   2. A couple of sample images for the demo / pretrained eval
#
# Overridable via env:
#   TT_MOGE_WEIGHTS  default ./weights
#   TT_MOGE_DATA     default ./data

set -euo pipefail

WEIGHTS_DIR="${TT_MOGE_WEIGHTS:-./weights}"
DATA_DIR="${TT_MOGE_DATA:-./data}"
mkdir -p "$WEIGHTS_DIR" "$DATA_DIR"

have() { [[ -s "$1" ]]; }

echo "Weights dir: $WEIGHTS_DIR"
echo "Data dir:    $DATA_DIR"

# --- MoGe-2 checkpoint (model_config + weights in a single model.pt)
if have "$WEIGHTS_DIR/model.pt"; then
  echo "ok  model.pt already present"
else
  echo "... model.pt (MoGe-2 ViT-L, ~1.3 GB)"
  curl -fL -o "$WEIGHTS_DIR/model.pt" \
    "https://huggingface.co/Ruicheng/moge-2-vitl-normal/resolve/main/model.pt"
fi

# --- Sample images (any RGB image works; these are just for the demo / eval)
for name in example.jpg; do
  if have "$DATA_DIR/$name"; then
    echo "ok  data/$name already present"
  else
    echo "... data/$name"
    curl -fL -o "$DATA_DIR/$name" \
      "https://raw.githubusercontent.com/microsoft/MoGe/main/assets/example_images/COLOR.jpg" || \
      echo "    (sample image download failed — point MOGE_IMAGE at any RGB image)"
  fi
done

echo
echo "Done."
ls -lh "$WEIGHTS_DIR" "$DATA_DIR" 2>/dev/null || true