#!/bin/bash
set -e

MODELS_DIR="$(pwd)/models"
CHECKPOINTS_DIR="$MODELS_DIR/checkpoints"
UPSCALE_DIR="$MODELS_DIR/upscale_models"
VAE_DIR="$MODELS_DIR/vae"
REFINER_DIR="$MODELS_DIR/refiners"
LORA_DIR="$MODELS_DIR/loras"

mkdir -p "$CHECKPOINTS_DIR" "$UPSCALE_DIR" "$VAE_DIR" "$REFINER_DIR" "$LORA_DIR"

echo "📥 Téléchargement des modèles SDXL dans $CHECKPOINTS_DIR"

# 1️⃣ SDXL Base
if [ ! -f "$CHECKPOINTS_DIR/sd_xl_base_1.0.safetensors" ]; then
    echo "➡️ Téléchargement SDXL Base..."
    curl -L -o "$CHECKPOINTS_DIR/sd_xl_base_1.0.safetensors" \
        https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors
else
    echo "✅ SDXL Base déjà présent"
fi

# 2️⃣ SDXL Refiner
if [ ! -f "$REFINER_DIR/sd_xl_refiner_1.0.safetensors" ]; then
    echo "➡️ Téléchargement SDXL Refiner..."
    curl -L -o "$REFINER_DIR/sd_xl_refiner_1.0.safetensors" \
        https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/resolve/main/sd_xl_refiner_1.0.safetensors
else
    echo "✅ SDXL Refiner déjà présent"
fi

# 3️⃣ VAE
if [ ! -f "$VAE_DIR/sdxl_vae.safetensors" ]; then
    echo "➡️ Téléchargement VAE..."
    curl -L -o "$VAE_DIR/sdxl_vae.safetensors" \
        https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors
else
    echo "✅ VAE déjà présent"
fi

# 4️⃣ Upscaler Real-ESRGAN x4
if [ ! -f "$UPSCALE_DIR/RealESRGAN_x4plus.pth" ]; then
    echo "➡️ Téléchargement Real-ESRGAN x4..."
    
    if [ -z "$HF_TOKEN" ]; then
        echo "❌ Erreur : variable HF_TOKEN non définie."
        echo "👉 Exécute : export HF_TOKEN=ton_token"
        exit 1
    fi

    curl -L -o "$UPSCALE_DIR/RealESRGAN_x4plus.pth" \
        -H "Authorization: Bearer $HF_TOKEN" \
        https://huggingface.co/xinntao/Real-ESRGAN/resolve/main/experiments/pretrained_models/RealESRGAN_x4plus.pth
else
    echo "✅ Real-ESRGAN déjà présent"
fi

# 5️⃣ Exemple LoRA (SDXL Anime style depuis CivitAI ou HuggingFace)
if [ ! -f "$LORA_DIR/animeStyle.safetensors" ]; then
    echo "➡️ Téléchargement LoRA (Anime Style)..."
    curl -L -o "$LORA_DIR/animeStyle.safetensors" \
        https://civitai.com/api/download/models/12345   # ⚠️ remplacer par l’ID réel du modèle
else
    echo "✅ LoRA Anime Style déjà présent"
fi

echo "🎉 Tous les modèles ont été téléchargés avec succès !"

#!/bin/bash
