install_musubi.sh

#!/bin/bash
set -e # Sai imediatamente se um comando falhar

echo "==========================================================="
echo "        MusubiPrime: Iniciando Instalação Otimizada        "
echo "==========================================================="

# --- 1. Verificação do Ambiente ---
echo "INFO: Verificando a versão do Python..."
python3 --version
# O ambiente RunPod já possui uma versão adequada do Python (3.11).

# --- 2. Clonando o Repositório Musubi Trainer ---
echo "INFO: Clonando o repositório Musubi Trainer..."
if [ -d "musubi-tuner" ]; then
    echo "WARN: O diretório 'musubi-tuner' já existe. Pulando a clonagem."
else
    git clone https://github.com/kohya-ss/musubi-tuner.git
fi
cd musubi-tuner

# --- 3. Configurando o Ambiente Virtual ---
echo "INFO: Configurando ambiente virtual Python em ./venv..."
if [ -d "venv" ]; then
    echo "WARN: O ambiente virtual 'venv' já existe. Ativando..."
else
    python3 -m venv venv
fi
source venv/bin/activate
echo "INFO: Ambiente virtual ativado."

# --- 4. Instalando Dependências Essenciais ---
echo "INFO: Atualizando pip..."
pip install --upgrade pip

echo "INFO: Atualizando PyTorch para uma versão estável compatível com CUDA 12.x..."
# Conforme README.md e Issue #273, uma versão mais recente do PyTorch é necessária.
# Usamos o índice cu121, que é o correto para o ambiente CUDA 12.4.1 da imagem.
pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

echo "INFO: Instalando dependências do requirements.txt..."
pip install -r requirements.txt

# --- 5. Aplicando Correções de Conflitos Conhecidas ---
echo "INFO: Aplicando correção para o problema de dependência do bitsandbytes (Issue #149)..."
pip install bitsandbytes --upgrade --pre

# --- 6. Instalando Dependências Opcionais para Monitoramento ---
echo "INFO: Instalando dependências opcionais (tensorboard, matplotlib)..."
pip install tensorboard matplotlib

# --- 7. Instalação Opcional de Aceleradores de INFERÊNCIA ---
# Estes são para INFERÊNCIA e NÃO são necessários para o TREINAMENTO.
# Eles são conhecidos por causar problemas se usados incorretamente durante o treinamento (Issue #92).
# Descomente as linhas abaixo apenas se você precisar de inferência acelerada.

# echo "INFO: Comandos para instalação opcional de aceleradores de inferência:"
# echo "      (Descomente no script se necessário)"
#
# # FlashAttention (para GPUs Ampere e mais recentes)
# pip install flash-attn --no-build-isolation
#
# # SageAttention (requer Triton)
# # pip install -U "triton-pre-release" -f https://storage.googleapis.com/triton-nightly/index.html
# # git clone https://github.com/thu-ml/SageAttention.git
# # cd SageAttention && pip install . && cd ..

echo "==========================================================="
echo "      Instalação do Musubi Trainer Concluída com Sucesso!    "
echo "==========================================================="
echo "Para ativar o ambiente em sessões futuras, execute:"
echo "cd /workspace/musubi-tuner && source venv/bin/activate"
echo "-----------------------------------------------------------"