#!/bin/bash
# Script de téléchargement des modèles ACE-Step 1.5 pour ComfyUI
# Usage: bash download-ace-step.sh

set -e

BASE_DIR="$HOME/ace-models"
COMFYUI_DIR="$HOME/Documents/OpenCode/Projets/Effi/comfyui/models"

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

# Créer le dossier de téléchargement
mkdir -p "$BASE_DIR"
cd "$BASE_DIR"

# 1. ace_step_1.5_turbo_aio (modèle principal AIO)
echo "📦 1/4 - Téléchargement ace_step_1.5_turbo_aio (~9.4GB)..."
if [ ! -f "ace_step_1.5_turbo_aio.safetensors" ]; then
    curl -L --progress-bar \
      "https://huggingface.co/Comfy-Org/ace_step_1.5_ComfyUI_files/resolve/main/checkpoints/ace_step_1.5_turbo_aio.safetensors" \
      -o "ace_step_1.5_turbo_aio.safetensors"
    echo "✅ Téléchargé"
else
    echo "✅ Déjà présent"
fi
ls -lh ace_step_1.5_turbo_aio.safetensors 2>/dev/null || echo "Fichier non trouvé"

echo ""

# 2. qwen_0.6b_ace15 (text encoder)
echo "📦 2/4 - Téléchargement qwen_0.6b_ace15 (~1.2GB)..."
if [ ! -f "qwen_0.6b_ace15.safetensors" ]; then
    curl -L --progress-bar \
      "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/main/Qwen3-Embedding-0.6B/model.safetensors" \
      -o "qwen_0.6b_ace15.safetensors"
    echo "✅ Téléchargé"
else
    echo "✅ Déjà présent"
fi
ls -lh qwen_0.6b_ace15.safetensors 2>/dev/null || echo "Fichier non trouvé"

echo ""

# 3. qwen_1.7b_ace15 (text encoder)
echo "📦 3/4 - Téléchargement qwen_1.7b_ace15 (~3.5GB)..."
if [ ! -f "qwen_1.7b_ace15.safetensors" ]; then
    curl -L --progress-bar \
      "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/main/acestep-5Hz-lm-1.7B/model.safetensors" \
      -o "qwen_1.7b_ace15.safetensors"
    echo "✅ Téléchargé"
else
    echo "✅ Déjà présent"
fi
ls -lh qwen_1.7b_ace15.safetensors 2>/dev/null || echo "Fichier non trouvé"

echo ""

# 4. ace_1.5_vae (VAE)
echo "📦 4/4 - Téléchargement ace_1.5_vae (~322MB)..."
if [ ! -f "ace_1.5_vae.safetensors" ]; then
    curl -L --progress-bar \
      "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/main/vae/diffusion_pytorch_model.safetensors" \
      -o "ace_1.5_vae.safetensors"
    echo "✅ Téléchargé"
else
    echo "✅ Déjà présent"
fi
ls -lh ace_1.5_vae.safetensors 2>/dev/null || echo "Fichier non trouvé"

echo ""
echo "=============================================="
echo "📁 Fichiers téléchargés dans: $BASE_DIR"
ls -lh *.safetensors 2>/dev/null
echo ""
echo "=============================================="
echo "⚠️  Étape suivante : Copier vers ComfyUI"
echo ""
echo "sudo cp *.safetensors $COMFYUI_DIR/checkpoints/"
echo "sudo cp ace_1.5_vae.safetensors $COMFYUI_DIR/vae/"
echo ""
echo "Puis redémarrer ComfyUI:"
echo "docker restart comfyui"
echo "=============================================="