#!/bin/bash
# Setup CUDA venv for Soprano on GTX 1660

set -e

echo "Setting up CUDA environment for Soprano..."

cd /home/koko210Serve/docker/miku-discord/soprano_to_rvc

# Remove old venv if exists
if [ -d ".venv-cuda" ]; then
    echo "Removing old CUDA venv..."
    rm -rf .venv-cuda
fi

# Create fresh venv with Python 3.11
echo "Creating Python 3.11 venv..."
python3.11 -m venv .venv-cuda

# Activate CUDA venv
source .venv-cuda/bin/activate

echo "Upgrading pip..."
pip install --upgrade pip

echo "Step 1: Installing latest CUDA PyTorch..."
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

echo "Step 2: Installing Soprano from source with lmdeploy..."
cd soprano
pip install -e '.[lmdeploy]'
cd ..

echo "Step 3: Installing IPC libraries..."
pip install redis pyzmq numpy

echo "Step 4: Verifying CUDA setup..."
python -c "
import torch
print(f'PyTorch version: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
print(f'CUDA version: {torch.version.cuda}')
print(f'Device count: {torch.cuda.device_count()}')
if torch.cuda.is_available():
    print(f'Device 0: {torch.cuda.get_device_name(0)}')
"

echo "Step 5: Testing Soprano..."
python -c "
from soprano import SopranoTTS
print('✓ Soprano imported successfully')
model = SopranoTTS(backend='lmdeploy', device='cuda', cache_size_mb=500, decoder_batch_size=2)
print(f'✓ Soprano loaded on {model.device} using {model.backend} backend')
"

echo ""
echo "========================================"
echo "✓ CUDA environment setup complete!"
echo "========================================"
echo ""
echo "GPU: NVIDIA GTX 1660 (6GB VRAM)"
echo "Backend: lmdeploy"
echo "Device: cuda:0"
echo ""
echo "NOTE: When running scripts, avoid running from"
echo "the soprano_to_rvc directory due to the soprano/"
echo "subdirectory shadowing the installed package."
echo "Run from /tmp or remove current dir from sys.path"
echo ""