ychappyboy/open-webui

1#!/bin/bash2 3# Define color and formatting codes4BOLD='\033[1m'5GREEN='\033[1;32m'6WHITE='\033[1;37m'7RED='\033[0;31m'8NC='\033[0m' # No Color9# Unicode character for tick mark10TICK='\u2713'11 12# Detect GPU driver13get_gpu_driver() {14    # Detect NVIDIA GPUs using lspci or nvidia-smi15    if lspci | grep -i nvidia >/dev/null || nvidia-smi >/dev/null 2>&1; then16        echo "nvidia"17        return18    fi19 20    # Detect AMD GPUs (including GCN architecture check for amdgpu vs radeon)21    if lspci | grep -i amd >/dev/null; then22        # List of known GCN and later architecture cards23        # This is a simplified list, and in a real-world scenario, you'd want a more comprehensive one24        local gcn_and_later=("Radeon HD 7000" "Radeon HD 8000" "Radeon R5" "Radeon R7" "Radeon R9" "Radeon RX")25 26        # Get GPU information27        local gpu_info=$(lspci | grep -i 'vga.*amd')28 29        for model in "${gcn_and_later[@]}"; do30            if echo "$gpu_info" | grep -iq "$model"; then31                echo "amdgpu"32                return33            fi34        done35 36        # Default to radeon if no GCN or later architecture is detected37        echo "radeon"38        return39    fi40 41    # Detect Intel GPUs42    if lspci | grep -i intel >/dev/null; then43        echo "i915"44        return45    fi46 47    # If no known GPU is detected48    echo "Unknown or unsupported GPU driver"49    exit 150}51 52# Function for rolling animation53show_loading() {54    local spin='-\|/'55    local i=056 57    printf " "58 59    while kill -0 $1 2>/dev/null; do60        i=$(( (i+1) %4 ))61        printf "\b${spin:$i:1}"62        sleep .163    done64 65    # Replace the spinner with a tick66    printf "\b${GREEN}${TICK}${NC}"67}68 69# Usage information70usage() {71    echo "Usage: $0 [OPTIONS]"72    echo "Options:"73    echo "  --enable-gpu[count=COUNT]  Enable GPU support with the specified count."74    echo "  --enable-api[port=PORT]    Enable API and expose it on the specified port."75    echo "  --webui[port=PORT]         Set the port for the web user interface."76    echo "  --data[folder=PATH]        Bind mount for ollama data folder (by default will create the 'ollama' volume)."77    echo "  --build                    Build the docker image before running the compose project."78    echo "  --drop                     Drop the compose project."79    echo "  -q, --quiet                Run script in headless mode."80    echo "  -h, --help                 Show this help message."81    echo ""82    echo "Examples:"83    echo "  $0 --drop"84    echo "  $0 --enable-gpu[count=1]"85    echo "  $0 --enable-gpu[count=all]"86    echo "  $0 --enable-api[port=11435]"87    echo "  $0 --enable-gpu[count=1] --enable-api[port=12345] --webui[port=3000]"88    echo "  $0 --enable-gpu[count=1] --enable-api[port=12345] --webui[port=3000] --data[folder=./ollama-data]"89    echo "  $0 --enable-gpu[count=1] --enable-api[port=12345] --webui[port=3000] --data[folder=./ollama-data] --build"90    echo ""91    echo "This script configures and runs a docker-compose setup with optional GPU support, API exposure, and web UI configuration."92    echo "About the gpu to use, the script automatically detects it using the "lspci" command."93    echo "In this case the gpu detected is: $(get_gpu_driver)"94}95 96# Default values97gpu_count=198api_port=1143599webui_port=3000100headless=false101build_image=false102kill_compose=false103 104# Function to extract value from the parameter105extract_value() {106    echo "$1" | sed -E 's/.*\[.*=(.*)\].*/\1/; t; s/.*//'107}108 109# Parse arguments110while [[ $# -gt 0 ]]; do111    key="$1"112 113    case $key in114        --enable-gpu*)115            enable_gpu=true116            value=$(extract_value "$key")117            gpu_count=${value:-1}118            ;;119        --enable-api*)120            enable_api=true121            value=$(extract_value "$key")122            api_port=${value:-11435}123            ;;124        --webui*)125            value=$(extract_value "$key")126            webui_port=${value:-3000}127            ;;128        --data*)129            value=$(extract_value "$key")130            data_dir=${value:-"./ollama-data"}131            ;;132        --drop)133            kill_compose=true134            ;;135        --build)136            build_image=true137            ;;138        -q|--quiet)139            headless=true140            ;;141        -h|--help)142            usage143            exit144            ;;145        *)146            # Unknown option147            echo "Unknown option: $key"148            usage149            exit 1150            ;;151    esac152    shift # past argument or value153done154 155if [[ $kill_compose == true ]]; then156    docker compose down --remove-orphans157    echo -e "${GREEN}${BOLD}Compose project dropped successfully.${NC}"158    exit159else160    DEFAULT_COMPOSE_COMMAND="docker compose -f docker-compose.yaml"161    if [[ $enable_gpu == true ]]; then162        # Validate and process command-line arguments163        if [[ -n $gpu_count ]]; then164            if ! [[ $gpu_count =~ ^([0-9]+|all)$ ]]; then165                echo "Invalid GPU count: $gpu_count"166                exit 1167            fi168            echo "Enabling GPU with $gpu_count GPUs"169            # Add your GPU allocation logic here170            export OLLAMA_GPU_DRIVER=$(get_gpu_driver)171            export OLLAMA_GPU_COUNT=$gpu_count # Set OLLAMA_GPU_COUNT environment variable172        fi173        DEFAULT_COMPOSE_COMMAND+=" -f docker-compose.gpu.yaml"174    fi175    if [[ $enable_api == true ]]; then176        DEFAULT_COMPOSE_COMMAND+=" -f docker-compose.api.yaml"177        if [[ -n $api_port ]]; then178            export OLLAMA_WEBAPI_PORT=$api_port # Set OLLAMA_WEBAPI_PORT environment variable179        fi180    fi181    if [[ -n $data_dir ]]; then182        DEFAULT_COMPOSE_COMMAND+=" -f docker-compose.data.yaml"183        export OLLAMA_DATA_DIR=$data_dir # Set OLLAMA_DATA_DIR environment variable184    fi185    if [[ -n $webui_port ]]; then186        export OPEN_WEBUI_PORT=$webui_port # Set OPEN_WEBUI_PORT environment variable187    fi188    DEFAULT_COMPOSE_COMMAND+=" up -d"189    DEFAULT_COMPOSE_COMMAND+=" --remove-orphans"190    DEFAULT_COMPOSE_COMMAND+=" --force-recreate"191    if [[ $build_image == true ]]; then192        DEFAULT_COMPOSE_COMMAND+=" --build"193    fi194fi195 196# Recap of environment variables197echo198echo -e "${WHITE}${BOLD}Current Setup:${NC}"199echo -e "   ${GREEN}${BOLD}GPU Driver:${NC} ${OLLAMA_GPU_DRIVER:-Not Enabled}"200echo -e "   ${GREEN}${BOLD}GPU Count:${NC} ${OLLAMA_GPU_COUNT:-Not Enabled}"201echo -e "   ${GREEN}${BOLD}WebAPI Port:${NC} ${OLLAMA_WEBAPI_PORT:-Not Enabled}"202echo -e "   ${GREEN}${BOLD}Data Folder:${NC} ${data_dir:-Using ollama volume}"203echo -e "   ${GREEN}${BOLD}WebUI Port:${NC} $webui_port"204echo205 206if [[ $headless == true ]]; then207    echo -ne "${WHITE}${BOLD}Running in headless mode... ${NC}"208    choice="y"209else210    # Ask for user acceptance211    echo -ne "${WHITE}${BOLD}Do you want to proceed with current setup? (Y/n): ${NC}"212    read -n1 -s choice213fi214 215echo216 217if [[ $choice == "" || $choice == "y" ]]; then218    # Execute the command with the current user219    eval "$DEFAULT_COMPOSE_COMMAND" &220 221    # Capture the background process PID222    PID=$!223 224    # Display the loading animation225    #show_loading $PID226 227    # Wait for the command to finish228    wait $PID229 230    echo231    # Check exit status232    if [ $? -eq 0 ]; then233        echo -e "${GREEN}${BOLD}Compose project started successfully.${NC}"234    else235        echo -e "${RED}${BOLD}There was an error starting the compose project.${NC}"236    fi237else238    echo "Aborted."239fi240 241echo242