#!/bin/sh
# Preseed ComfyUI's models volume with checkpoints, VAEs, LoRAs, etc.
# Runs once via the comfyui-model-init service (see docker-compose.yml).
# Safe to re-run — already-present files are skipped.
#
# ComfyUI doesn't have a "pull" command of its own, so this is plain curl
# against direct download URLs. For HuggingFace, the direct URL is:
#   https://huggingface.co/<repo>/resolve/main/<file>
# For gated HF repos (Flux-dev, SD3, etc.), set HF_TOKEN in .env — the
# script attaches it as a bearer token automatically.

set -e

apk add --no-cache curl >/dev/null

mkdir -p /models/checkpoints /models/vae /models/loras /models/controlnet \
         /models/clip /models/clip_vision /models/upscale_models /models/embeddings

fetch() {
    dest="$1"; name="$2"; url="$3"
    target="/models/$dest/$name"

    if [ -f "$target" ]; then
        echo "✓ $dest/$name already present"
        return
    fi

    echo "→ Downloading $dest/$name…"
    mkdir -p "/models/$dest"

    if [ -n "$HF_TOKEN" ] && echo "$url" | grep -q huggingface.co; then
        curl -fL -C - --retry 3 -H "Authorization: Bearer $HF_TOKEN" \
             -o "$target.partial" "$url"
    else
        curl -fL -C - --retry 3 -o "$target.partial" "$url"
    fi
    mv "$target.partial" "$target"
}

# ─── Edit the list below to choose what gets preseeded ──────────────────────
# Format: fetch <subdir under /models> <filename to save as> <direct URL>
#
# No checkpoints are downloaded by default — the deployment ships expecting
# you to point at your own model mirror or the public examples below.
# Whatever filename you pick should match the `ckpt_name` field in
# workflows/txt2img.json and workflows/img2img.json (the shipped default
# is CyberRealisticXLPlay_V8.0_FP16.safetensors); update either the
# script or the workflows so they line up.

# Examples — uncomment what you want.

# SDXL Base 1.0 (~6.9 GB)
# fetch checkpoints sd_xl_base_1.0.safetensors \
#     https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors

# SDXL VAE (fixes washed-out colours on some SDXL checkpoints)
# fetch vae sdxl_vae.safetensors \
#     https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors

# Flux.1-dev (~23 GB, gated — needs HF_TOKEN with access to black-forest-labs)
# fetch checkpoints flux1-dev.safetensors \
#     https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors

# 4x-UltraSharp upscaler
# fetch upscale_models 4x-UltraSharp.pth \
#     https://huggingface.co/lokCX/4x-Ultrasharp/resolve/main/4x-UltraSharp.pth

echo "Done."