#!/bin/bash
set -e

echo "=== AI Model Installation ==="

# Check Python version
python3 --version

# Install PyTorch (CPU version for staging)
echo "Installing PyTorch..."
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu

# Install Transformers (for CLIP)
echo "Installing Transformers..."
pip install transformers

# Install Ultralytics (for YOLO)
echo "Installing Ultralytics..."
pip install ultralytics

# Download YOLOv8 model
echo "Downloading YOLOv8 model..."
python3 -c "from ultralytics import YOLO; YOLO('yolov8n.pt')"

# Download CLIP model
echo "Downloading CLIP model..."
python3 -c "from transformers import CLIPProcessor, CLIPModel; CLIPModel.from_pretrained('openai/clip-vit-base-patch32'); CLIPProcessor.from_pretrained('openai/clip-vit-base-patch32')"

echo "=== AI Models Installed ==="
echo "Models downloaded to ~/.cache/"
echo ""
echo "For GPU support, install CUDA version:"
echo "pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118"