#!/bin/bash
# RunPod Quick Setup Script

echo "🚀 Setting up VR180 Matting on RunPod..."
echo "GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader)"
echo ""

# Update system
echo "📦 Installing system dependencies..."
apt-get update && apt-get install -y ffmpeg git wget nano

# Clone your repository (replace with your actual repo URL)
cd /workspace
if [ ! -d "sam2e" ]; then
    echo "📥 Cloning repository..."
    git clone https://git.10n.app/scott/test2.git
fi
cd sam2e

# Install Python dependencies
echo "🐍 Installing Python dependencies..."
pip install --upgrade pip
pip install -r requirements.txt

# Install SAM2 separately (not on PyPI)
echo "🎯 Installing SAM2..."
pip install git+https://github.com/facebookresearch/segment-anything-2.git

# Install project
echo "📦 Installing VR180 matting package..."
pip install -e .

# Download models
echo "📥 Downloading models..."
mkdir -p models

# Download YOLOv8 models
python -c "from ultralytics import YOLO; YOLO('yolov8n.pt'); YOLO('yolov8m.pt')"

# Download SAM2 checkpoint
cd models
if [ ! -f "sam2_hiera_large.pt" ]; then
    echo "Downloading SAM2 model weights..."
    wget -q --show-progress https://dl.fbaipublicfiles.com/segment_anything_2/sam2_hiera_large.pt
fi
cd ..

# Create working directories
mkdir -p /workspace/data /workspace/output

# Test installation
echo ""
echo "🧪 Testing installation..."
python test_installation.py

echo ""
echo "✅ Setup complete!"
echo ""
echo "📝 Quick start:"
echo "1. Upload your VR180 video to /workspace/data/"
echo "   wget -O /workspace/data/video.mp4 'your-video-url'"
echo ""
echo "2. Use the RunPod optimized config:"
echo "   cp config_runpod.yaml config.yaml"
echo "   nano config.yaml  # Update video path"
echo ""
echo "3. Run the matting:"
echo "   vr180-matting config.yaml"
echo ""
echo "💡 For A40 GPU, you can use higher quality settings:"
echo "   vr180-matting config.yaml --scale 0.75"