#!/bin/bash
# Install script for Microsoft VibeVoice on Nvidia RTX (Windows/Linux via WSL/Bash)

echo "🚀 Starting VibeVoice Setup for Nvidia RTX..."

# 1. Create Directory
mkdir -p VibeVoice_RTX
cd VibeVoice_RTX

# 2. Clone Repository
echo "📦 Cloning VibeVoice repository..."
if [ ! -d "VibeVoice" ]; then
    git clone https://github.com/vibevoice-community/VibeVoice.git
fi
cd VibeVoice

# 3. Setup Python Environment
echo "🐍 Setting up Python environment..."
python3 -m venv venv
source venv/bin/activate

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

# 5. Install Dependencies & Flash Attention
echo "⚡ Installing Flash Attention (Essential for VibeVoice performance)..."
pip install packaging ninja
pip install flash-attn --no-build-isolation

echo "📥 Installing usage dependencies..."
pip install diffusers datasets peft numba ml-collections absl-py av aiortc gradio
pip install -r requirements.txt

# 6. Download Model
echo "💾 Downloading VibeVoice-1.5B Model..."
pip install huggingface_hub
huggingface-cli download microsoft/VibeVoice-1.5B --local-dir models/VibeVoice-1.5B --local-dir-use-symlinks False

# 7. Apply Fix for "custom_generate/generate.py not found" error
echo "🔧 Applying fix for missing generation config..."
mkdir -p models/VibeVoice-1.5B/custom_generate
touch models/VibeVoice-1.5B/custom_generate/__init__.py
echo "def generate(*args, **kwargs): pass" > models/VibeVoice-1.5B/custom_generate/generate.py

echo "✅ Setup Complete!"
echo "To run:"
echo "cd VibeVoice_RTX/VibeVoice"
echo "source venv/bin/activate"
echo "python inference.py --model_path models/VibeVoice-1.5B"