#!/bin/bash

set -euo pipefail
TMP_CACHE=$(mktemp -d)
export HF_HUB_CACHE="$TMP_CACHE"
TMP=$(mktemp -d)
MODELS_DIR="ComfyUI/models"
LORAS_DIR="$MODELS_DIR/loras/krea2"

hf_cli_install() {
    if ! command -v hf &> /dev/null; then
        echo "hf not found, installing..."
        curl -LsSf https://hf.co/cli/install.sh | bash
    else
        echo "hf is already installed."
    fi
}

hf_download_flat() {
    local repo="$1" file="$2" dest="$3"
    hf download "$repo" "$file" --local-dir "$TMP"
    mkdir -p "$dest"
    mv "$TMP/$file" "$dest/$(basename "$file")"
}

hf_cli_install
hf_download_flat "Comfy-Org/Krea-2" "diffusion_models/krea2_raw_int8_convrot.safetensors" "$MODELS_DIR/diffusion_models"
hf_download_flat "Comfy-Org/Krea-2" "text_encoders/qwen3vl_4b_fp8_scaled.safetensors" "$MODELS_DIR/text_encoders"
hf_download_flat "Comfy-Org/Wan_2.1_ComfyUI_repackaged" "split_files/vae/wan_2.1_vae.safetensors" "$MODELS_DIR/vae"

hf_download_flat "Comfy-Org/Krea-2" "loras/krea2_turbo_lora_rank_64_bf16.safetensors" "$LORAS_DIR"

hf_download_flat "cdkkkk/setup" "krea2/MysticXXX_KREA2_v2_stripped.safetensors" "$LORAS_DIR"
hf_download_flat "cdkkkk/setup" "krea2/lenovo_krea2.safetensors" "$LORAS_DIR"
hf_download_flat "cdkkkk/setup" "krea2/realism_engine_krea2_v2_stripped.safetensors" "$LORAS_DIR"
hf_download_flat "cdkkkk/setup" "krea2/snofs_krea_v1_nostrip.safetensors" "$LORAS_DIR"
hf_download_flat "cdkkkk/setup" "krea2/HMCum_krea2_epoch30.safetensors" "$LORAS_DIR"
hf_download_flat "cdkkkk/setup" "krea2/Krea2_RAW_RoughV2_000005000.safetensors" "$LORAS_DIR"
hf_download_flat "cdkkkk/setup" "krea2/snofs_krea_v1.safetensors" "$LORAS_DIR"

rm -rf "$TMP_CACHE"
rm -rf "$TMP"