#!/bin/bash

# Remove existing virtual environment if it exists
if [ -d "venv" ]; then
    echo "Removing existing virtual environment..."
    rm -rf venv
fi

# Create and activate a new virtual environment
echo "Creating and activating a new virtual environment..."
python3 -m venv venv
source venv/bin/activate

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

# Install PyTorch with CUDA support (for NVIDIA RTX 3060 Ti, likely CUDA 11.8)
echo "Installing PyTorch with CUDA support..."
pip install torch==2.1.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118

# Install Audiocraft, Gradio, pydub, and other dependencies
echo "Installing Audiocraft, Gradio, pydub, and other dependencies..."
pip install audiocraft==1.3.0 gradio==4.44.1 numpy==1.23.5 pydub==0.25.1 psutil==6.1.0

# Install ffmpeg for MP3 export
echo "Installing ffmpeg..."
sudo apt-get update
sudo apt-get install -y ffmpeg

# Create the models directory if it doesn't exist
echo "Creating models directory for musicgen-medium..."
mkdir -p models/musicgen-medium

# Download the musicgen-medium model using Audiocraft
# This will download the model to ~/.cache/audiocraft by default
echo "Downloading musicgen-medium model..."
python3 -c "from audiocraft.models import MusicGen; MusicGen.get_pretrained('facebook/musicgen-medium')"

# Move the model to the specified directory
# Audiocraft typically downloads models to ~/.cache/audiocraft
# We'll move it to /home/ubuntu/ghostai_music_generator/models/musicgen-medium
echo "Moving model to /home/ubuntu/ghostai_music_generator/models/musicgen-medium..."
mv ~/.cache/audiocraft/models--facebook--musicgen-medium/* models/musicgen-medium/

# Deactivate the virtual environment
echo "Deactivating virtual environment..."
deactivate

echo "Setup complete! You can now activate the virtual environment and run your application:"
echo "source venv/bin/activate"
echo "python3 app.py"