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

set -e

# Détection robuste du répertoire du script
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MODELS_DIR="$SCRIPT_DIR/models/checkpoints"
VAE_DIR="$SCRIPT_DIR/models/vae"

mkdir -p "$MODELS_DIR" "$VAE_DIR"

echo "🎨 Téléchargement des modèles SDXL recommandés"
echo "================================================"
echo ""
echo "Espace disque requis : ~25-35 GB au total"
echo ""

# Fonction de téléchargement avec barre de progression
download_model() {
    local url=$1
    local output=$2
    local name=$3
    local size=$4
    
    if [ -f "$output" ]; then
        echo "✓ $name déjà présent ($(du -h "$output" | cut -f1))"
        return 0
    fi
    
    echo ""
    echo "📥 Téléchargement de $name ($size)..."
    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
}

# Menu de sélection
echo "Quels modèles souhaitez-vous télécharger ?"
echo ""
echo "=== MODÈLES PHOTORÉALISTES ==="
echo "1) RealVisXL V4.0 (6.9GB) - Portraits ultra-réalistes, peau détaillée"
echo "2) Juggernaut XL (6.9GB) - Polyvalent, excellent pour tout type d'image"
echo "3) Proteus XL (6.9GB) - Photoréalisme amélioré, excellente qualité"
echo ""
echo "=== MODÈLES ART/STYLISÉS ==="
echo "4) DreamShaper XL (6.9GB) - Art fantastique, peinture, illustration"
echo "5) SDXL Base 1.0 (6.9GB) - Modèle de base (déjà installé)"
echo ""
echo "=== MODÈLES RAPIDES ==="
echo "6) SDXL Lightning (6.9GB) - Génération 4x plus rapide, qualité légèrement inférieure"
echo ""
echo "=== VAE (Optionnel mais recommandé) ==="
echo "7) SDXL VAE (335MB) - Meilleure qualité de couleurs et détails"
echo ""
echo "=== INSTALLATION COMPLETE ==="
echo "8) Tout télécharger (⚠️ 35GB) - Tous les modèles ci-dessus"
echo ""
echo "0) Annuler"
echo ""
read -p "Votre choix (0-8): " choice

case $choice in
    1)
        echo ""
        download_model \
            "https://huggingface.co/SG161222/RealVisXL_V4.0/resolve/main/RealVisXL_V4.0.safetensors" \
            "$MODELS_DIR/RealVisXL_V4.0.safetensors" \
            "RealVisXL V4.0" \
            "6.9GB"
        echo ""
        echo "💡 Conseil : Excellent pour les portraits et photos réalistes"
        echo "   Prompt tip: 'professional photo, 8k uhd, dslr, soft lighting'"
        ;;
    2)
        echo ""
        download_model \
            "https://huggingface.co/RunDiffusion/Juggernaut-XL-v9/resolve/main/Juggernaut-XL-v9.safetensors" \
            "$MODELS_DIR/Juggernaut-XL-v9.safetensors" \
            "Juggernaut XL v9" \
            "6.9GB"
        echo ""
        echo "💡 Conseil : Très polyvalent, bon compromis qualité/vitesse"
        echo "   Prompt tip: Fonctionne bien avec tous les styles"
        ;;
    3)
        echo ""
        download_model \
            "https://huggingface.co/dathudehenry/ProteusV0.4/resolve/main/ProteusV0.4.safetensors" \
            "$MODELS_DIR/ProteusV0.4.safetensors" \
            "Proteus V0.4" \
            "6.9GB"
        echo ""
        echo "💡 Conseil : Alternative à RealVisXL, excellent pour la peau"
        echo "   Prompt tip: 'raw photo, realistic, detailed skin texture'"
        ;;
    4)
        echo ""
        download_model \
            "https://huggingface.co/Lykon/DreamShaper/resolve/main/DreamShaperXL1.0.safetensors" \
            "$MODELS_DIR/DreamShaperXL1.0.safetensors" \
            "DreamShaper XL 1.0" \
            "6.9GB"
        echo ""
        echo "💡 Conseil : Parfait pour l'art, la fantasy et les illustrations"
        echo "   Prompt tip: 'digital art, concept art, highly detailed, fantasy'"
        ;;
    5)
        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" \
            "6.9GB"
        echo ""
        echo "💡 Conseil : Modèle de base Stability AI, très stable"
        ;;
    6)
        echo ""
        download_model \
            "https://huggingface.co/ByteDance/SDXL-Lightning/resolve/main/sdxl_lightning_8step.safetensors" \
            "$MODELS_DIR/sdxl_lightning_8step.safetensors" \
            "SDXL Lightning 8-step" \
            "6.9GB"
        echo ""
        echo "💡 Conseil : Utilise seulement 8 étapes (au lieu de 30-40)"
        echo "   ⚠️ Configurez 'steps: 8' dans OpenWebUI pour ce modèle"
        echo "   Parfait pour les tests rapides !"
        ;;
    7)
        echo ""
        download_model \
            "https://huggingface.co/stabilityai/sdxl-vae/resolve/main/diffusion_pytorch_model.safetensors" \
            "$VAE_DIR/sdxl_vae.safetensors" \
            "SDXL VAE" \
            "335MB"
        echo ""
        echo "💡 Conseil : Utilisez ce VAE avec tous les modèles SDXL pour de meilleures couleurs"
        ;;
    8)
        echo ""
        echo "⚠️  Cela va télécharger ~35GB de modèles"
        read -p "Êtes-vous sûr ? (oui/non): " confirm
        if [ "$confirm" = "oui" ]; then
            download_model \
                "https://huggingface.co/SG161222/RealVisXL_V4.0/resolve/main/RealVisXL_V4.0.safetensors" \
                "$MODELS_DIR/RealVisXL_V4.0.safetensors" \
                "RealVisXL V4.0" \
                "6.9GB"
            
            download_model \
                "https://huggingface.co/RunDiffusion/Juggernaut-XL-v9/resolve/main/Juggernaut-XL-v9.safetensors" \
                "$MODELS_DIR/Juggernaut-XL-v9.safetensors" \
                "Juggernaut XL v9" \
                "6.9GB"
            
            download_model \
                "https://huggingface.co/Lykon/DreamShaper/resolve/main/DreamShaperXL1.0.safetensors" \
                "$MODELS_DIR/DreamShaperXL1.0.safetensors" \
                "DreamShaper XL 1.0" \
                "6.9GB"
            
            download_model \
                "https://huggingface.co/ByteDance/SDXL-Lightning/resolve/main/sdxl_lightning_8step.safetensors" \
                "$MODELS_DIR/sdxl_lightning_8step.safetensors" \
                "SDXL Lightning 8-step" \
                "6.9GB"
            
            download_model \
                "https://huggingface.co/stabilityai/sdxl-vae/resolve/main/diffusion_pytorch_model.safetensors" \
                "$VAE_DIR/sdxl_vae.safetensors" \
                "SDXL VAE" \
                "335MB"
            
            echo ""
            echo "✅ Tous les modèles ont été téléchargés !"
        else
            echo "Annulation..."
        fi
        ;;
    0)
        echo "Annulation..."
        exit 0
        ;;
    *)
        echo "Choix invalide"
        exit 1
        ;;
esac

echo ""
echo "=========================================="
echo "✅ Téléchargement terminé !"
echo ""
echo "Prochaines étapes :"
echo "  1. Redémarrer ComfyUI : sudo docker restart comfyui"
echo "  2. Choisir le modèle dans OpenWebUI (Admin → Settings → Images)"
echo "  3. Consulter docs/MODELS.md pour les réglages optimaux"
echo ""
echo "Espace disque utilisé :"
du -sh "$MODELS_DIR" "$VAE_DIR" 2>/dev/null | head -2