#!/bin/bash
# MuseTalk V1.5 Installation Script
# Tested on: RTX 4090, CUDA 12.x, Ubuntu
# Author: Dumont Talker Team
# Date: 2024-12-24

set -e

echo "=============================================="
echo "  MuseTalk V1.5 Installation Script"
echo "=============================================="

# Variables
INSTALL_DIR="${INSTALL_DIR:-/workspace}"
MINICONDA_DIR="${MINICONDA_DIR:-/opt/miniconda}"
ENV_NAME="musetalk"

# Check if running as root or with sudo
if [ "$EUID" -ne 0 ]; then
    echo "Please run as root or with sudo"
    exit 1
fi

# Install system dependencies
echo ""
echo "=== Installing system dependencies ==="
apt-get update -qq
apt-get install -y -qq wget git ffmpeg espeak-ng libgl1-mesa-glx libglib2.0-0 > /dev/null

# Install Miniconda if not present
if [ ! -d "$MINICONDA_DIR" ]; then
    echo ""
    echo "=== Installing Miniconda ==="
    wget -q -O /tmp/miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
    bash /tmp/miniconda.sh -b -p "$MINICONDA_DIR"
    rm /tmp/miniconda.sh
fi

# Source conda
source "$MINICONDA_DIR/bin/activate"

# Create conda environment
echo ""
echo "=== Creating conda environment: $ENV_NAME ==="
if conda env list | grep -q "^$ENV_NAME "; then
    echo "Environment $ENV_NAME already exists, skipping creation"
else
    conda create -n "$ENV_NAME" python=3.10 -y -q
fi

# Activate environment
conda activate "$ENV_NAME"

# Clone MuseTalk repository
echo ""
echo "=== Cloning MuseTalk repository ==="
cd "$INSTALL_DIR"
if [ -d "MuseTalk" ]; then
    echo "MuseTalk directory already exists, pulling latest changes"
    cd MuseTalk && git pull && cd ..
else
    git clone --depth 1 https://github.com/TMElyralab/MuseTalk.git
fi
cd MuseTalk

# Install PyTorch with CUDA support
echo ""
echo "=== Installing PyTorch 2.1.0 with CUDA 12.1 ==="
pip install -q torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121

# Install OpenMMLab dependencies
echo ""
echo "=== Installing OpenMMLab dependencies ==="
pip install -q mmengine==0.10.4
pip install -q mmcv==2.1.0 -f https://download.openmmlab.com/mmcv/dist/cu121/torch2.1/index.html
pip install -q mmdet==3.3.0
pip install -q mmpose==1.3.1 --no-deps

# Install compatible versions of key dependencies
echo ""
echo "=== Installing compatible dependencies ==="
pip install -q 'numpy<2' 'numpy==1.26.4'
pip install -q 'transformers==4.42.0'
pip install -q 'diffusers==0.28.0'
pip install -q 'huggingface_hub==0.23.2'
pip install -q 'opencv-python==4.8.0.74'

# Install other requirements
echo ""
echo "=== Installing other requirements ==="
pip install -q omegaconf tqdm einops av librosa soundfile scikit-image
pip install -q ffmpeg-python accelerate safetensors

# Download models
echo ""
echo "=== Downloading models ==="

# Create model directories
mkdir -p models/dwpose models/sd-vae-ft-mse models/whisper

# Create symlink for sd-vae
ln -sf sd-vae-ft-mse models/sd-vae 2>/dev/null || true

# Download DWPose model
echo "Downloading DWPose model..."
if [ ! -f "models/dwpose/dw-ll_ucoco_384.pth" ]; then
    wget -q --show-progress -O models/dwpose/dw-ll_ucoco_384.pth \
        'https://huggingface.co/yzd-v/DWPose/resolve/main/dw-ll_ucoco_384.pth'
fi

# Download SD-VAE model
echo "Downloading SD-VAE model..."
if [ ! -f "models/sd-vae-ft-mse/config.json" ]; then
    wget -q -O models/sd-vae-ft-mse/config.json \
        'https://huggingface.co/stabilityai/sd-vae-ft-mse/resolve/main/config.json'
fi
if [ ! -f "models/sd-vae-ft-mse/diffusion_pytorch_model.bin" ]; then
    wget -q --show-progress -O models/sd-vae-ft-mse/diffusion_pytorch_model.bin \
        'https://huggingface.co/stabilityai/sd-vae-ft-mse/resolve/main/diffusion_pytorch_model.bin'
fi

# Download Whisper model from HuggingFace
echo "Downloading Whisper model..."
python3 -c "
from huggingface_hub import snapshot_download
snapshot_download(repo_id='openai/whisper-tiny', local_dir='models/whisper', local_dir_use_symlinks=False)
" 2>/dev/null

# Create config.json in musetalk folder
echo "Fixing config files..."
if [ -f "models/musetalk/musetalk.json" ] && [ ! -f "models/musetalk/config.json" ]; then
    cp models/musetalk/musetalk.json models/musetalk/config.json
fi

# Download MuseTalk V1.5 weights if not present (using git lfs)
echo ""
echo "=== Downloading MuseTalk weights ==="
echo "Note: MuseTalk weights should be downloaded from HuggingFace:"
echo "  https://huggingface.co/TMElyralab/MuseTalk"
echo ""
echo "If weights are not present, run:"
echo "  cd $INSTALL_DIR/MuseTalk"
echo "  git lfs install"
echo "  git lfs pull"

# Verify installation
echo ""
echo "=== Verifying installation ==="
python3 -c "
import torch
print(f'PyTorch version: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'CUDA version: {torch.version.cuda}')
    print(f'GPU: {torch.cuda.get_device_name(0)}')
"

echo ""
echo "=============================================="
echo "  Installation Complete!"
echo "=============================================="
echo ""
echo "To run MuseTalk V1.5:"
echo "  1. Activate the environment:"
echo "     source $MINICONDA_DIR/bin/activate $ENV_NAME"
echo ""
echo "  2. Go to MuseTalk directory:"
echo "     cd $INSTALL_DIR/MuseTalk"
echo ""
echo "  3. Create inference config (configs/inference/test.yaml):"
echo "     task_0:"
echo "       video_path: \"data/video/your_avatar.mp4\""
echo "       audio_path: \"path/to/your_audio.wav\""
echo ""
echo "  4. Run inference:"
echo "     PYTHONPATH=. python3 scripts/inference.py --version v15 --inference_config configs/inference/test.yaml"
echo ""
echo "Output will be saved in ./results/v15/"
echo "=============================================="