#!/bin/bash
# Script de téléchargement des modèles Stable Diffusion
# Pour ComfyUI avec RX 7800 XT

set -e

MODELS_DIR="$(dirname "$0")/models/checkpoints"
VAE_DIR="$(dirname "$0")/models/vae"

mkdir -p "$MODELS_DIR" "$VAE_DIR"

echo "🎨 Téléchargement des modèles pour ComfyUI"
echo "=========================================="
echo ""

# Fonction de téléchargement avec barre de progression
download_model() {
    local url=$1
    local output=$2
    local name=$3
    
    if [ -f "$output" ]; then
        echo "✓ $name déjà présent"
        return 0
    fi
    
    echo "📥 Téléchargement de $name..."
    echo "   URL: $url"
    echo "   Destination: $output"
    echo ""
    
    wget --progress=bar:force -O "$output" "$url" 2>&1 | tail -f -n +6
    
    if [ $? -eq 0 ]; then
        echo "✓ $name téléchargé avec succès"
    else
        echo "✗ Erreur lors du téléchargement de $name"
        rm -f "$output"
        return 1
    fi
    echo ""
}

# Menu de sélection
echo "Quel modèle souhaitez-vous télécharger ?"
echo ""
echo "1) SDXL Base 1.0 (6.9GB) - Recommandé, usage général"
echo "2) SDXL Base + VAE (7.2GB total) - Avec meilleur VAE"
echo "3) Annuler"
echo ""
read -p "Votre choix (1-3): " choice

case $choice in
    1)
        echo ""
        download_model \
            "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors" \
            "$MODELS_DIR/sd_xl_base_1.0.safetensors" \
            "SDXL Base 1.0"
        ;;
    2)
        echo ""
        download_model \
            "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors" \
            "$MODELS_DIR/sd_xl_base_1.0.safetensors" \
            "SDXL Base 1.0"
        
        download_model \
            "https://huggingface.co/stabilityai/sdxl-vae/resolve/main/diffusion_pytorch_model.safetensors" \
            "$VAE_DIR/sdxl_vae.safetensors" \
            "SDXL VAE"
        ;;
    3)
        echo "Annulation..."
        exit 0
        ;;
    *)
        echo "Choix invalide"
        exit 1
        ;;
esac

echo ""
echo "=========================================="
echo "✅ Téléchargement terminé !"
echo ""
echo "Prochaines étapes :"
echo "  1. Lancer ComfyUI : docker compose up -d"
echo "  2. Vérifier les logs : docker logs -f comfyui"
echo "  3. Accéder à l'interface : http://192.168.1.11:8188"
echo ""
echo "Espace disque utilisé :"
du -sh "$MODELS_DIR" "$VAE_DIR" 2>/dev/null | head -2