#!/bin/bash
# Download FLUX.1 schnell + accompanying VAE/text encoders into a ComfyUI tree.
# Uses ungated mirrors — the official Black-Forest-Labs repo is gated and
# requires an HF token. See docs/setup-comfyui-mrock.md.

set -euo pipefail

ROOT="${1:-$HOME/dev/comfyui/models}"

if [ ! -d "$ROOT" ]; then
    echo "models root $ROOT does not exist — pass it as the first argument" >&2
    exit 1
fi

mkdir -p "$ROOT/unet" "$ROOT/vae" "$ROOT/clip"

CKPT="https://huggingface.co/Comfy-Org/flux1-schnell/resolve/main/flux1-schnell.safetensors"
VAE="https://huggingface.co/sirorable/flux-ae-vae/resolve/main/ae.safetensors"
CLIP_L="https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/clip_l.safetensors"
T5="https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp8_e4m3fn.safetensors"

dl() {
    local url=$1 dest=$2
    if [ -s "$dest" ]; then
        echo "skip $dest (already present)"
        return
    fi
    echo "downloading $url -> $dest"
    curl -L --fail --retry 3 --retry-delay 5 -C - -o "$dest" "$url"
}

dl "$CKPT"   "$ROOT/unet/flux1-schnell.safetensors"
dl "$VAE"    "$ROOT/vae/ae.safetensors"
dl "$CLIP_L" "$ROOT/clip/clip_l.safetensors"
dl "$T5"     "$ROOT/clip/t5xxl_fp8_e4m3fn.safetensors"

echo "done"