Private "wrapper-script" to maintain and launch different installed AI tools

ai-start

#!/bin/bash
#
# Script to easier managing, launching, killing, restarting, and updating ALL of my precious AI tools
# Date: 13. October 2024
#
#            DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE
#                    Version 2, December 2004

# Copyright (C) 2024 Andreas Nilsen <[email protected]> - @adde88 - https://www.github.com/adde88

# Everyone is permitted to copy and distribute verbatim or modified
# copies of this license document, and changing it is allowed as long
# as the name is changed.

#            DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE
#   TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION
#
#  0. You just DO WHAT THE FUCK YOU WANT TO.
#

# XXX: Setup global variables for the script
A1111_SDF_HOME="$HOME/stable-diffusion-forge"
A1111_HOME="$HOME/stable-diffusion"
SDNEXT_HOME="$HOME/sd.next"
COMFYUI_HOME="$HOME/comfyui"
OMOST_HOME="$HOME/omost"
FOOOCUS_HOME="$HOME/fooocus"
TEXTUI_HOME="$HOME/text-generation-webui"

# XXX: Function: Update all AI tools to their latest Git commits
update_all_ai_tools() {
    echo -e "This will update every single AI tools you are using.\n"
    local dirs=(
        "$HOME/stable-diffusion-forge"
        "$HOME/stable-diffusion"
        "$HOME/sd.next"
        "$HOME/fooocus"
        "$HOME/omost"
        "$HOME/comfyui"
        "$HOME/text-generation-webui"
    )

    for dir in "${dirs[@]}"; do
        if [ -d "$dir" ]; then
            echo "Entering directory: $dir"
            cd "$dir"
            # Replace the following line with the command you want to run
            git pull && echo ""
        else
            echo "Directory: '$dir' does not exist."
        fi
    done
}

# XXX: Function: Update Stable Diffusion's core to it's latest Git commits
function update_sd_core(){
    cd "$A1111_HOME"
    echo -e "Checking for Updates for Stable Diffusion"
    git pull
    cd - > /dev/null 2>&1
}

# XXX: Function: Update Stable Diffusion Forge's core to it's latest Git commits
function update_sdf_core(){
    cd "$A1111_SDF_HOME"
    echo -e "Checking for Updates for Stable Diffusion Forge"
    git pull
    cd - > /dev/null 2>&1
}

# XXX: Function: Update SD.Next's core to it's latest Git commits
function update_sdnext_core(){
    cd "$SDNEXT_HOME"
    echo -e "Checking for Updates for SD.Next"
    git pull
    cd - > /dev/null 2>&1
}

# XXX: Function: Update Fooocus's core to it's latest Git commits
function update_fooocus_core(){
    cd "$FOOOCUS_HOME"
    echo -e "Checking for Updates for Fooocus"
    git pull
    cd - > /dev/null 2>&1
}

# XXX: Function: Update Omost's core to it's latest Git commits
function update_omost_core(){
    cd "$OMOST_HOME"
    echo -e "Checking for Updates for Omost"
    git pull
    cd - > /dev/null 2>&1
}

# XXX: Function: Update ComfyUI's core to it's latest Git commits
function update_comfyui_core(){
    cd "$COMFYUI_HOME"
    echo -e "Checking for Updates for ComfyUI"
    git pull
    cd - > /dev/null 2>&1
}

# XXX: Function: Update Text Generation WebUI's core to it's latest Git commits
function update_textgen_core(){
    cd "$TEXTUI_HOME"
    echo -e "Checking for Updates for Text Generation WebUI"
    git pull
    cd - > /dev/null 2>&1
}

# XXX: Function: Kill Stable Diffusion Forge processes, without any restart options
function sdf-kill(){
    # Get the PID of the process
    pid1=$(pgrep -f "bash ./webui.sh")
    pid2=$(pgrep -f "python3.10 -u launch.py")
    pid3=$(pgrep -f "stable-diffusion-forge/extensions/sd-civitai-browser-plus/aria2/lin/aria2")
    pid4=$(pgrep -f "stable-diffusion-forge/venv/lib/python3.10/site-packages/gradio/frpc_linux_amd64_v0.2")
    pid5=$(pgrep -f "accelerate launch --num_cpu_threads_per_process")
    # If the process is running, kill it
    if [[ -z "$pid1" && -z "$pid2" && -z "$pid3" && -z "$pid4" && -z "$pid5" ]]; then
      echo -e "No Stable Diffusion Forge processes was found running!"
      return 1
    else
      echo -e "Killing processes!\nPID 1: $pid1\nPID 2: $pid2\nPID 3: $pid3\nPID 4: $pid4\nPID 5: $pid5"
      for pids in "$pid1" "$pid2" "$pid3" "$pid4" "$pid5"; do
        kill -9 "$pids"  > /dev/null 2>&1
      done
      echo -e "Stable Diffusion Forge  processes was shut down properly..."
    fi
}

# XXX: Function: Kill Stable Diffusion processes, without any restart options
function sd-kill(){
    echo -e "This function is not finished, yet! Please be patient."
}

# XXX: Function: Kill ComfyUI processes, without any restart options
function comfyui-kill(){
    echo -e "This function is not finished, yet! Please be patient."
}

# XXX: Function: Kill SD.Next processes, without any restart options
function sdnext-kill(){
    echo -e "This function is not finished, yet! Please be patient."
}

# XXX: Function: Kill Omost processes, without any restart options
function omost-kill(){
    echo -e "This function is not finished, yet! Please be patient."
}

# XXX: Function: Kill Fooocus processes, without any restart options
function fooocus-kill(){
    echo -e "This function is not finished, yet! Please be patient."
}

# XXX: Function: Kill Text Generation WebUI processes, without any restart options
function textui-kill(){
    echo -e "This function is not finished, yet! Please be patient."
}

# XXX: Function: Kill and ask for restart of Stable Diffusion
function sd-restart(){
    sd-kill
    # Ask the user if they want to restart Stable Diffusion
    read -p "Would you like to re-start Stable Diffusion? Y/N: " user_input
    case $user_input in
        [Yy]* ) main_sd;;
        [Nn]* ) ;;
        * ) echo "Invalid input";;
    esac
}

# XXX: Function: Kill and ask for restart of Stable Diffusion Forge
function sdf-restart(){
    sdf-kill
    # Ask the user if they want to restart Stable Diffusion Forge
    read -p "Would you like to re-start Stable Diffusion Forge? Y/N: " user_input
    case $user_input in
        [Yy]* ) main_sdf;;
        [Nn]* ) ;;
        * ) echo "Invalid input";;
    esac
}

# XXX: Function: Main function that launches Stable Diffusion Forge
function main_sdf(){
    cd "$A1111_SDF_HOME"
    echo -e "Stable Diffusion Forge, and it's extensions should now be updated to the latest commits, if there was any updates to be found."
    echo -e "Please manually scroll up and read the text, to manually check for any errors, or other info if you're curious."
    ./webui-user.sh ; ./webui.sh &
    echo -e "\nLaunched Stable Diffusion Forge v1.10 on GPU + CPU"
    echo -e "http://localhost:7860"
}

# XXX: Function: Main function that launches Stable Diffusion
function main_sd(){
    cd "$A1111_HOME"
    echo -e "Stable Diffusion, and it's extensions should now be updated to the latest commits, if there was any updates to be found."
    echo -e "Please manually scroll up and read the text, to manually check for any errors, or other info if you're curious."
    ./webui-user.sh ; ./webui.sh &
    echo -e "\nLaunched Stable Diffusion v1.10 on GPU + CPU"
    echo -e "http://localhost:7860"
}

# XXX: Function: Main function that launches ComfyUI
function main_comfyui(){
    cd "$COMFYUI_HOME"
    source ./venv/bin/activate
    echo -e "ComfyUI should now be updated to the latest commits, if there was any updates to be found."
    echo -e "Please manually scroll up and read the text, to manually check for any errors, or other info if you're curious."
    python main.py
    echo -e "\nLaunched ComfyUI v0.2.3 on GPU + CPU"
    echo -e "Remember we have Gligen-GUI at: '$HOME/gligen-gui' available"
    echo -e "http://localhost:8188"
}

function help_msg(){
    echo -e "To chose which AI tool, or option you want to launch, please follow these instructions:"
    echo -e "Usage: ai-start [OPTIONS]\n"
    echo -e "\nLaunch Options:"
    echo -e "-s, --sd                       This will launch Stable Diffusion"
    echo -e "-f, --sdf                      This will launch Stable Diffusion Forge"
    echo -e "-c, --comfyui                  This will launch ComfyUI"
    echo -e "\n-U, --update-all-tools         This will update all AI tools that we have"
}

# XXX: Check if --update-all is part of any arguments, if so it will pull the latest git commits from the public repositories
update_all=false
for arg in "$@"; do
    if [[ "$arg" == "--update-all-tools" || "$arg" == "-U" ]]; then
        update_all=true
        break
    fi
done

if $update_all; then
    update_all_ai_tools
fi

# XXX: Check if --sdf, --sd, or -comfyui is part of any arguments, and launch the corresponding AI tool
if [ $# -eq 0 ]; then
    help_msg
    exit 0
fi
    
for arg in "$@"; do
    if [[ "$arg" == "--sdf" || "$arg" == "-s" ]]; then
        sdf-kill && sleep 1 && sdf-kill
        update_sdf_core
        main_sdf
        break
    elif [[ "$arg" == "--sd" || "$arg" == "-f" ]]; then
        sd-kill && sleep 1 && sd-kill
        update_sd_core
        main_sd
        break
    elif [[ "$arg" == "--comfyui" || "$arg" == "-c" ]]; then
        comfyui-kill && sleep 1 && comfyui-kill
        update_comfyui_core
        main_comfyui
        break
    elif [[ "$arg" == "--help" ]]; then
        help_msg
    else
        help_msg
    fi
done