#!/usr/bin/env bash
# Pre-download all three model components, sequentially, in the foreground.
#
# Total: ~58 GB
#   - microsoft/Lens-Turbo     : 16.4 GB  (gated — must Agree on HF web)
#   - openai/gpt-oss-20b       : 41.3 GB  (Apache-2.0, no agreement)
#   - black-forest-labs/FLUX.2-dev/vae/ : ~0.6 GB  (gated — must Agree)
#
# Strategy: download to default HF cache (NOT --local-dir), because Lens's
# model_index.json references text_encoder/vae by HF repo id. The cache
# layout `$HF_HOME/hub/models--ORG--REPO/snapshots/...` is what
# LensPipeline.from_pretrained auto-discovers.
#
# `hf` has built-in resume — interrupt with Ctrl-C and re-run
# the same command to continue. Run inside tmux/screen if you want to
# detach.
#
# Optional skip flags:
#   --skip-lens       skip microsoft/Lens-Turbo
#   --skip-gpt-oss    skip openai/gpt-oss-20b
#   --skip-vae        skip FLUX.2-dev/vae/

set -uo pipefail

SKIP_LENS=0
SKIP_GPT=0
SKIP_VAE=0
for arg in "$@"; do
    case "$arg" in
        --skip-lens)     SKIP_LENS=1 ;;
        --skip-gpt-oss)  SKIP_GPT=1 ;;
        --skip-vae)      SKIP_VAE=1 ;;
        -h|--help)
            sed -n '2,25p' "$0"
            exit 0 ;;
        *)
            echo "unknown arg: $arg" >&2; exit 2 ;;
    esac
done

HF_HOME="${HF_HOME:-/mnt/modelall/hf_cache}"
export HF_HOME
mkdir -p "$HF_HOME"

# Env hygiene -----------------------------------------------------------------
# `HF_HUB_ENABLE_HF_TRANSFER` is deprecated in huggingface_hub >= 1.x — the
# hf_transfer Rust binary is no longer used. Suppress the warning by
# unsetting it (matches v2 README §1).
unset HF_HUB_ENABLE_HF_TRANSFER
# To opt INTO the new Xet-based fast transfer, uncomment:
#   export HF_XET_HIGH_PERFORMANCE=1
# Default-off here for conservative behavior; benchmark before enabling.

echo "[download] HF_HOME=$HF_HOME"
echo "[download] hf version:"
hf version 2>/dev/null | sed 's/^/  /' || \
    { echo "  ERROR: 'hf' command not found. pip install -U huggingface_hub" >&2; exit 1; }

# Pre-flight: disk + RAM sanity check
echo
echo "[download] disk space at HF_HOME:"
df -h "$(dirname "$HF_HOME")" | tail -1
echo "[download] RAM:"
free -g | head -2
echo
echo "Need ~80 GB free disk and ≥48 GB RAM."
read -t 10 -p "Continue? [Y/n] (auto-yes in 10s) " ans || ans=Y
[ "${ans:-Y}" = "n" ] && exit 1

# Check login
echo
echo "[download] checking HF login..."
python3 -c "
from huggingface_hub import whoami
try:
    info = whoami()
    print(f'  logged in as: {info[\"name\"]}')
except Exception as e:
    print(f'  NOT LOGGED IN: {e}')
    print('  -> run: hf auth login')
    import sys; sys.exit(1)
" || exit 1

echo
echo "[download] gated-repo reminder (click Agree on these pages BEFORE running):"
echo "  - https://huggingface.co/microsoft/Lens-Turbo"
echo "  - https://huggingface.co/black-forest-labs/FLUX.2-dev"
echo "  (openai/gpt-oss-20b is Apache-2.0, no agreement needed)"
echo

# ---------------------------------------------------------------------------
# Step 1/3: Lens-Turbo (16.4 GB)
# ---------------------------------------------------------------------------
if [ "$SKIP_LENS" = "0" ]; then
    echo
    echo "============================================================"
    echo "[download] 1/3: microsoft/Lens-Turbo  (16.4 GB, gated)"
    echo "============================================================"
    hf download microsoft/Lens-Turbo --max-workers 8
    echo "[download] 1/3 done"
else
    echo "[download] 1/3 SKIPPED (--skip-lens)"
fi

# ---------------------------------------------------------------------------
# Step 2/3: GPT-OSS-20B text encoder (~41 GB)
# ---------------------------------------------------------------------------
if [ "$SKIP_GPT" = "0" ]; then
    echo
    echo "============================================================"
    echo "[download] 2/3: openai/gpt-oss-20b   (41 GB, Apache-2.0)"
    echo "============================================================"
    hf download openai/gpt-oss-20b --max-workers 8
    echo "[download] 2/3 done"
else
    echo "[download] 2/3 SKIPPED (--skip-gpt-oss)"
fi

# ---------------------------------------------------------------------------
# Step 3/3: FLUX.2 VAE (~600 MB, only the vae/ subdir)
# ---------------------------------------------------------------------------
if [ "$SKIP_VAE" = "0" ]; then
    echo
    echo "============================================================"
    echo "[download] 3/3: black-forest-labs/FLUX.2-dev/vae/  (~600 MB, gated)"
    echo "============================================================"
    hf download black-forest-labs/FLUX.2-dev \
        --include "vae/*" --max-workers 4
    echo "[download] 3/3 done"
else
    echo "[download] 3/3 SKIPPED (--skip-vae)"
fi

echo
echo "[download] cache sizes:"
du -sh "$HF_HOME"/hub/models--*/ 2>/dev/null

echo
echo "[download] verifying transformer shards exist:"
ls -lh "$HF_HOME"/hub/models--microsoft--Lens-Turbo/snapshots/*/transformer/*.safetensors 2>/dev/null \
    || echo "  (no transformer shards yet — Lens-Turbo may not be downloaded)"

echo
echo "[download] done. Restart ComfyUI, then run the CLI:"
echo "  python cli/lens_turbo_cli.py 'your prompt' --ip <SERVER> --gpu 0"