#!/usr/bin/env bash
# Download all weights for Flow-CoPD experiments (base model + teachers + reward models).
# Codebase: third_party/flow_grpo (Flow-GRPO). See SETUP.md.
#
# Usage:
#   bash scripts/download_weights.sh            # public weights only (teachers + reward models)
#   bash scripts/download_weights.sh --all      # also the GATED base model (needs HF token + license)
set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"   # repo root (portable)
export HF_HOME="${HF_HOME:-$ROOT/.hf_cache}"          # big models cached here, loaded by repo-id in flow_grpo configs
mkdir -p "$HF_HOME" "$ROOT/checkpoints/teachers"
echo "HF_HOME=$HF_HOME"

dl() { echo ">>> $1"; huggingface-cli download "$@"; }

echo "==================================================================="
echo "[1/3] Reward models (PUBLIC, no token) -> HF cache, loaded by repo-id"
echo "==================================================================="
dl laion/CLIP-ViT-H-14-laion2B-s32B-b79K     # PickScore image/text backbone (pickscore_scorer.py)
dl yuvalkirstain/PickScore_v1                # PickScore reward head        (pickscore_scorer.py)
dl openai/clip-vit-large-patch14            # Aesthetic + CLIPScore backbone (aesthetic_scorer.py / clip_scorer.py)
# (Aesthetic MLP head sac+logos+ava1-l14-linearMSE.pth already ships inside the repo's flow_grpo/assets/)
# (OCR reward = PaddleOCR, installed + predownloaded in setup_env.sh, not an HF model)

echo "==================================================================="
echo "[2/3] TEACHERS = Flow-GRPO task experts (PUBLIC LoRA adapters)"
echo "      loaded via config.train.lora_path -> PeftModel.from_pretrained"
echo "==================================================================="
dl jieliu/SD3.5M-FlowGRPO-Text      --local-dir "$ROOT/checkpoints/teachers/text"       # T2 teacher: text rendering / OCR
dl jieliu/SD3.5M-FlowGRPO-PickScore --local-dir "$ROOT/checkpoints/teachers/pickscore"  # T1 teacher: human preference / aesthetic
# optional 3rd teacher for the 3-teacher ablation (Block C5):
# dl jieliu/SD3.5M-FlowGRPO-GenEval --local-dir "$ROOT/checkpoints/teachers/geneval"

if [[ "${1:-}" == "--all" ]]; then
  echo "==================================================================="
  echo "[3/3] BASE model stabilityai/stable-diffusion-3.5-medium (GATED)"
  echo "==================================================================="
  if ! huggingface-cli whoami >/dev/null 2>&1; then
    echo "!! Not logged in to Hugging Face. The base model is GATED."
    echo "   1) Accept the license:  https://huggingface.co/stabilityai/stable-diffusion-3.5-medium"
    echo "   2) huggingface-cli login   (or: export HF_TOKEN=hf_xxx)"
    echo "   3) re-run: bash scripts/download_weights.sh --all"
    exit 1
  fi
  dl stabilityai/stable-diffusion-3.5-medium   # ~20GB; loaded by repo-id in every flow_grpo config
else
  echo ""
  echo "[3/3] SKIPPED base model (gated). Re-run with --all after HF login:"
  echo "      accept license -> huggingface-cli login -> bash scripts/download_weights.sh --all"
fi

echo "DONE. Public weights in $HF_HOME ; teachers in $ROOT/checkpoints/teachers/"