#!/bin/bash
# Script de téléchargement des modèles ACE-Step 1.5 COMPLETS
# À lancer sur Folzo Corp

set -e

BASE_DIR="/home/folzo/Documents/OpenCode/Projets/Effi/comfyui/models"

echo "🎵 ACE-Step 1.5 - Téléchargement COMPLET des modèles"
echo "===================================================="
echo ""

# 1. VAE
echo "📦 1/4 - Téléchargement du VAE..."
mkdir -p "$BASE_DIR/vae"
cd "$BASE_DIR/vae"
if [ ! -f "ace_1.5_vae.safetensors" ]; then
    wget -q --progress=bar:force "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/main/vae/diffusion_pytorch_model.safetensors" -O "ace_1.5_vae.safetensors"
    echo "✅ VAE téléchargé"
else
    echo "✅ VAE déjà présent"
fi
ls -lh "ace_1.5_vae.safetensors"

echo ""

# 2. Diffusion Model
echo "📦 2/4 - Téléchargement du modèle de diffusion..."
mkdir -p "$BASE_DIR/diffusion_models"
cd "$BASE_DIR/diffusion_models"
if [ ! -f "acestep_v1.5_turbo.safetensors" ]; then
    wget -q --progress=bar:force "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/main/acestep-v15-turbo/model.safetensors" -O "acestep_v1.5_turbo.safetensors"
    echo "✅ Diffusion model téléchargé"
else
    echo "✅ Diffusion model déjà présent"
fi
ls -lh "acestep_v1.5_turbo.safetensors"

echo ""

# 3. Text Encoder 0.6B
echo "📦 3/4 - Téléchargement du text encoder 0.6B..."
mkdir -p "$BASE_DIR/text_encoders"
cd "$BASE_DIR/text_encoders"
if [ ! -f "qwen_0.6b_ace15.safetensors" ]; then
    wget -q --progress=bar:force "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/main/acestep-5Hz-lm-0.6B/model.safetensors" -O "qwen_0.6b_ace15.safetensors"
    echo "✅ Text encoder 0.6B téléchargé"
else
    echo "✅ Text encoder 0.6B déjà présent"
fi
ls -lh "qwen_0.6b_ace15.safetensors"

echo ""

# 4. Text Encoder 4B
echo "📦 4/4 - Téléchargement du text encoder 4B..."
cd "$BASE_DIR/text_encoders"
if [ ! -f "qwen_4b_ace15.safetensors" ]; then
    wget -q --progress=bar:force "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/main/acestep-5Hz-lm-4B/model.safetensors" -O "qwen_4b_ace15.safetensors"
    echo "✅ Text encoder 4B téléchargé"
else
    echo "✅ Text encoder 4B déjà présent"
fi
ls -lh "qwen_4b_ace15.safetensors"

echo ""
echo "===================================================="
echo "🎉 Téléchargement terminé !"
echo ""
echo "Prochaine étape : redémarrer ComfyUI"
echo "docker restart comfyui"
echo "===================================================="