#!/bin/bash
# Установка специализированной vision модели для Ollama

set -e

CONTAINER_OLLAMA="ollama"
MODEL_CHOICE="${1:-llava}"

echo "=== Установка Vision модели для Ollama ==="
echo ""

# Определяем команду docker
DOCKER_CMD="docker"
if ! docker ps >/dev/null 2>&1; then
    DOCKER_CMD="sudo docker"
fi

echo "Доступные vision модели:"
echo "  1. llava:latest (LLaVA 1.6) - 7B, лучший выбор для vision"
echo "  2. bakllava:latest (BakLLaVA) - 7B, Mistral + LLaVA"
echo "  3. llama3.2-vision:latest (Llama 3.2 Vision) - 11B, требует 8GB VRAM"
echo ""

if [ "$MODEL_CHOICE" = "llava" ]; then
    MODEL="llava:latest"
    echo "Выбрана модель: $MODEL (LLaVA 1.6)"
elif [ "$MODEL_CHOICE" = "bakllava" ]; then
    MODEL="bakllava:latest"
    echo "Выбрана модель: $MODEL (BakLLaVA)"
elif [ "$MODEL_CHOICE" = "llama3.2" ]; then
    MODEL="llama3.2-vision:latest"
    echo "Выбрана модель: $MODEL (Llama 3.2 Vision)"
else
    MODEL="llava:latest"
    echo "Используется модель по умолчанию: $MODEL"
fi

echo ""
echo "Загрузка модели $MODEL..."
echo "Это может занять несколько минут в зависимости от скорости интернета..."
echo ""

$DOCKER_CMD exec $CONTAINER_OLLAMA ollama pull "$MODEL" || {
    echo "✗ Ошибка при загрузке модели"
    exit 1
}

echo ""
echo "✓ Модель загружена!"
echo ""
echo "Проверка установленных моделей:"
$DOCKER_CMD exec $CONTAINER_OLLAMA ollama list

echo ""
echo "=== Готово! ==="
echo ""
echo "Теперь в Open WebUI выберите модель: $MODEL"
echo "И попробуйте загрузить изображение и задать вопрос о нем."