install_wan.sh

#!/bin/bash

# Script d'installation automatique WAN Video ComfyUI
# Compatible RunPod et installation locale

set -e

echo "🚀 Installation automatique WAN Video ComfyUI..."
echo "================================================"

# Configuration des variables
WORKSPACE_DIR="/workspace"
COMFY_DIR="$WORKSPACE_DIR/comfywan"
VENV_DIR="$WORKSPACE_DIR/venv"

# Fonction de log
log() {
    echo "✅ $1"
}

error() {
    echo "❌ $1"
    exit 1
}

# Vérification des prérequis
check_requirements() {
    log "Vérification des prérequis..."
    
    # Vérifier Python 3.10
    if ! python --version | grep -q "3.10"; then
        error "Python 3.10 requis. Version actuelle: $(python --version)"
    fi
    
    # Vérifier CUDA
    if ! nvidia-smi &>/dev/null; then
        error "NVIDIA GPU/CUDA non détecté"
    fi
    
    log "Prérequis OK"
}

# Installation des dépendances système
install_system_deps() {
    log "Installation des dépendances système..."
    
    apt update -qq
    apt install -y --no-install-recommends \
        aria2 \
        git \
        git-lfs \
        curl \
        wget \
        ffmpeg \
        build-essential \
        ninja-build \
        libgl1 \
        libglib2.0-0 \
        vim
    
    log "Dépendances système installées"
}

# Clonage de ComfyUI
setup_comfyui() {
    log "Configuration de ComfyUI..."
    
    if [ -d "$COMFY_DIR" ]; then
        log "ComfyUI déjà présent, mise à jour..."
        cd "$COMFY_DIR"
        git pull
    else
        log "Clonage de ComfyUI..."
        cd "$WORKSPACE_DIR"
        git clone --depth=1 https://github.com/comfyanonymous/ComfyUI.git comfywan
    fi
    
    cd "$COMFY_DIR"
    log "ComfyUI configuré"
}

# Configuration de l'environnement Python
setup_python_env() {
    log "Configuration de l'environnement Python..."
    
    if [ ! -d "$VENV_DIR" ]; then
        python -m venv "$VENV_DIR"
    fi
    
    source "$VENV_DIR/bin/activate"
    
    # Mise à jour de pip
    pip install --upgrade pip setuptools wheel
    
    log "Environnement Python configuré"
}

# Installation de PyTorch avec CUDA
install_pytorch() {
    log "Installation de PyTorch avec CUDA 12.8..."
    
    source "$VENV_DIR/bin/activate"
    
    pip install --no-cache-dir torch==2.7.0+cu128 \
        --index-url https://download.pytorch.org/whl/cu128 --no-deps
    
    pip install --no-cache-dir \
        torchvision==0.22.0+cu128 \
        torchaudio==2.7.0 \
        --index-url https://download.pytorch.org/whl/cu128
    
    log "PyTorch installé"
}

# Installation des dépendances ComfyUI
install_comfyui_deps() {
    log "Installation des dépendances ComfyUI..."
    
    source "$VENV_DIR/bin/activate"
    cd "$COMFY_DIR"
    
    pip install -r requirements.txt
    
    log "Dépendances ComfyUI installées"
}

# Installation des nodes personnalisés
install_custom_nodes() {
    log "Installation des nodes personnalisés..."
    
    source "$VENV_DIR/bin/activate"
    cd "$COMFY_DIR/custom_nodes"
    
    # ComfyUI Manager
    if [ ! -d "ComfyUI-Manager" ]; then
        git clone https://github.com/ltdrdata/ComfyUI-Manager.git
    fi
    cd ComfyUI-Manager
    pip install -r requirements.txt
    cd ..
    
    # KJNodes (requis pour WAN)
    if [ ! -d "ComfyUI-KJNodes" ]; then
        git clone https://github.com/kijai/ComfyUI-KJNodes.git
    fi
    cd ComfyUI-KJNodes
    pip install -r requirements.txt
    cd ..
    
    # TeaCache pour l'optimisation
    if [ ! -d "ComfyUI-TeaCache" ]; then
        git clone https://github.com/welltop-cn/ComfyUI-TeaCache.git
    fi
    cd ComfyUI-TeaCache
    pip install -r requirements.txt
    cd ..
    
    log "Nodes personnalisés installés"
}

# Installation des dépendances serverless
install_serverless_deps() {
    log "Installation des dépendances serverless..."
    
    source "$VENV_DIR/bin/activate"
    
    pip install \
        requests \
        runpod==1.7.9 \
        websocket-client \
        onnxruntime-gpu \
        triton
    
    log "Dépendances serverless installées"
}

# Installation de SageAttention
install_sage_attention() {
    log "Installation de SageAttention..."
    
    source "$VENV_DIR/bin/activate"
    cd "$WORKSPACE_DIR"
    
    if [ ! -d "SageAttention" ]; then
        git clone https://github.com/thu-ml/SageAttention.git
    fi
    
    cd SageAttention
    python setup.py install
    
    log "SageAttention installé"
}

# Téléchargement des modèles WAN
download_models() {
    log "Téléchargement des modèles WAN (cela peut prendre du temps)..."
    
    cd "$COMFY_DIR"
    
    # Créer les dossiers
    mkdir -p models/{diffusion_models,text_encoders,clip_vision,vae,loras,upscale_models}
    
    # Modèle principal WAN 2.1 VACE (GGUF)
    if [ ! -f "models/diffusion_models/Wan2.1-VACE-14B-Q6_K.gguf" ]; then
        log "Téléchargement du modèle principal WAN..."
        aria2c -x16 -s16 -d models/diffusion_models \
            -o Wan2.1-VACE-14B-Q6_K.gguf --continue=true \
            https://huggingface.co/QuantStack/Wan2.1-VACE-14B-GGUF/resolve/main/Wan2.1-VACE-14B-Q6_K.gguf
    fi
    
    # Text Encoder
    if [ ! -f "models/text_encoders/umt5-xxl-encoder-Q8_0.gguf" ]; then
        log "Téléchargement du text encoder..."
        aria2c -x16 -s16 -d models/text_encoders \
            -o umt5-xxl-encoder-Q8_0.gguf --continue=true \
            https://huggingface.co/city96/umt5-xxl-encoder-gguf/resolve/main/umt5-xxl-encoder-Q8_0.gguf
    fi
    
    # CLIP Vision
    if [ ! -f "models/clip_vision/clip_vision_h.safetensors" ]; then
        log "Téléchargement du CLIP Vision..."
        aria2c -x16 -s16 -d models/clip_vision \
            -o clip_vision_h.safetensors --continue=true \
            https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/resolve/main/split_files/clip_vision/clip_vision_h.safetensors
    fi
    
    # VAE
    if [ ! -f "models/vae/wan_2.1_vae.safetensors" ]; then
        log "Téléchargement du VAE..."
        aria2c -x16 -s16 -d models/vae \
            -o wan_2.1_vae.safetensors --continue=true \
            https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/resolve/main/split_files/vae/wan_2.1_vae.safetensors
    fi
    
    # LoRA essentiel (lightx2v)
    if [ ! -f "models/loras/Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32.safetensors" ]; then
        log "Téléchargement du LoRA lightx2v..."
        aria2c -x16 -s16 -d models/loras \
            -o Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32.safetensors --continue=true \
            https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32.safetensors
    fi
    
    # Fun Reward LoRA
    if [ ! -f "models/loras/Wan2.1-Fun-14B-InP-MPS.safetensors" ]; then
        log "Téléchargement du Fun Reward LoRA..."
        aria2c -x16 -s16 -d models/loras \
            -o Wan2.1-Fun-14B-InP-MPS.safetensors --continue=true \
            https://huggingface.co/alibaba-pai/Wan2.1-Fun-Reward-LoRAs/resolve/main/Wan2.1-Fun-14B-InP-MPS.safetensors
    fi
    
    # Upscaler
    if [ ! -f "models/upscale_models/4xLSDIR.pth" ]; then
        log "Téléchargement de l'upscaler..."
        aria2c -x16 -s16 -d models/upscale_models \
            -o 4xLSDIR.pth --continue=true \
            https://github.com/Phhofm/models/raw/main/4xLSDIR/4xLSDIR.pth
    fi
    
    log "Modèles téléchargés"
}

# Configuration finale
final_setup() {
    log "Configuration finale..."
    
    # Créer le dossier logs
    mkdir -p "$WORKSPACE_DIR/logs"
    
    # Copier les fichiers de configuration s'ils existent
    if [ -f "/workspace/extra_model_paths.yaml" ]; then
        cp /workspace/extra_model_paths.yaml "$COMFY_DIR/"
    fi
    
    log "Configuration terminée"
}

# Test de l'installation
test_installation() {
    log "Test de l'installation..."
    
    source "$VENV_DIR/bin/activate"
    cd "$COMFY_DIR"
    
    # Test import des dépendances critiques
    python -c "
import torch
import comfy
print(f'PyTorch: {torch.__version__}')
print(f'CUDA disponible: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'GPU: {torch.cuda.get_device_name(0)}')
print('✅ Installation réussie!')
"
    
    log "Test réussi !"
}

# Menu d'options
show_menu() {
    echo
    echo "🎯 Options d'installation:"
    echo "1. Installation complète (recommandé)"
    echo "2. Installation sans LoRAs spéciaux (plus rapide)"
    echo "3. Mise à jour uniquement"
    echo "4. Test de l'installation existante"
    echo
    read -p "Votre choix (1-4): " choice
    echo
}

# LoRAs spéciaux (optionnel)
download_special_loras() {
    log "Téléchargement des LoRAs spéciaux..."
    
    cd "$COMFY_DIR"
    
    echo "⚠️  Les LoRAs suivants nécessitent un token CivitAI (optionnel)"
    echo "Vous pouvez les ignorer et les télécharger manuellement plus tard"
    echo
    
    read -p "Avez-vous un token CivitAI ? (y/n): " has_token
    
    if [ "$has_token" = "y" ]; then
        read -p "Entrez votre token CivitAI: " token
        
        # Hulk transformation
        log "Téléchargement Hulk LoRA..."
        aria2c -x16 -s16 -d models/loras \
            -o Hulk_epoch35.safetensors --continue=true \
            "https://civitai.com/api/download/models/1588339?type=Model&format=SafeTensor&token=${token}" || true
        
        # Super Saiyan
        log "Téléchargement Super Saiyan LoRA..."
        aria2c -x16 -s16 -d models/loras \
            -o super_saiyan_35_epochs.safetensors --continue=true \
            "https://civitai.com/api/download/models/1554033?type=Model&format=SafeTensor&token=${token}" || true
    else
        log "LoRAs spéciaux ignorés (vous pouvez les télécharger manuellement)"
    fi
}

# Fonction principale
main() {
    echo "🌟 Installation WAN Video ComfyUI"
    echo "=================================="
    
    show_menu
    
    case $choice in
        1)
            log "🚀 Installation complète..."
            check_requirements
            install_system_deps
            setup_comfyui
            setup_python_env
            install_pytorch
            install_comfyui_deps
            install_custom_nodes
            install_serverless_deps
            install_sage_attention
            download_models
            download_special_loras
            final_setup
            test_installation
            ;;
        2)
            log "⚡ Installation rapide (sans LoRAs spéciaux)..."
            check_requirements
            install_system_deps
            setup_comfyui
            setup_python_env
            install_pytorch
            install_comfyui_deps
            install_custom_nodes
            install_serverless_deps
            install_sage_attention
            download_models
            final_setup
            test_installation
            ;;
        3)
            log "🔄 Mise à jour..."
            source "$VENV_DIR/bin/activate"
            cd "$COMFY_DIR"
            git pull
            pip install --upgrade -r requirements.txt
            test_installation
            ;;
        4)
            log "🧪 Test de l'installation..."
            test_installation
            ;;
        *)
            error "Choix invalide"
            ;;
    esac
    
    echo
    echo "🎉 Installation terminée !"
    echo "========================================="
    echo "📁 Dossier: $COMFY_DIR"
    echo "🐍 Venv: $VENV_DIR"
    echo
    echo "🚀 Pour démarrer ComfyUI:"
    echo "source $VENV_DIR/bin/activate"
    echo "cd $COMFY_DIR"
    echo "python main.py --port 3000 --use-sage-attention"
    echo
    echo "🌐 Interface: http://localhost:3000"
    echo "📖 Workflows: Utilisez le fichier vace-i2v.json fourni"
}

# Gestion des erreurs
trap 'error "Installation interrompue"' INT TERM

# Lancement
main "$@"
Made with