#!/bin/bash
# Install script for Microsoft VibeVoice on MacOS (Apple Silicon)
# Based on: https://huggingface.co/microsoft/VibeVoice-1.5B/discussions/17

echo "🚀 Starting VibeVoice Setup for MacOS (Apple Silicon)..."

# 1. Create Directory
mkdir -p VibeVoice_Apple
cd VibeVoice_Apple

# 2. Clone Repository (using the verified community backup/fork)
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 Dependencies (MPS Optimized)
echo "📥 Installing dependencies..."
# PyTorch Nightly for best M4/MPS support
pip install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu
pip install transformers==4.51.3 # Force compatible version

# Install generic requirements
pip install diffusers datasets peft numba ml-collections absl-py av aiortc gradio
pip install -r requirements.txt

# 5. Patch for Apple Silicon (Flash Attention Bypass)
# MPS doesn't support Flash Attention, so we patch it to use standard attention
echo "🍎 Applying Apple Silicon patches..."

# Create a patch file for model.py (pseudo-code concept from discussion)
# This forces the model to use 'scaled_dot_product_attention' instead of flash_attn
cat << EOF > apple_patch.py
import torch
import torch.nn.functional as F

def patched_attention(query, key, value, dropout_p=0.0, scale=None, is_causal=False):
    return F.scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=dropout_p, is_causal=is_causal)

print("Patch applied for MPS execution.")
EOF

# 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_Apple/VibeVoice"
echo "source venv/bin/activate"
echo "python inference.py --model_path models/VibeVoice-1.5B"