#!/bin/bash

# ASCII Art
echo "================================="
echo " █████╗  ██╗ ██╗ ██╗ ██╗    ██╗███╗   ██╗███████╗████████╗ █████╗ ██╗     ██╗     ███████╗██████╗ 
██╔══██╗███║███║███║███║    ██║████╗  ██║██╔════╝╚══██╔══╝██╔══██╗██║     ██║     ██╔════╝██╔══██╗
███████║╚██║╚██║╚██║╚██║    ██║██╔██╗ ██║███████╗   ██║   ███████║██║     ██║     █████╗  ██████╔╝
██╔══██║ ██║ ██║ ██║ ██║    ██║██║╚██╗██║╚════██║   ██║   ██╔══██║██║     ██║     ██╔══╝  ██╔══██╗
██║  ██║ ██║ ██║ ██║ ██║    ██║██║ ╚████║███████║   ██║   ██║  ██║███████╗███████╗███████╗██║  ██║
╚═╝  ╚═╝ ╚═╝ ╚═╝ ╚═╝ ╚═╝    ╚═╝╚═╝  ╚═══╝╚══════╝   ╚═╝   ╚═╝  ╚═╝╚══════╝╚══════╝╚══════╝╚═╝  ╚═╝"
echo "================================="

# Function to install Nvidia CUDA toolkit
install_cuda() {
    echo "Step 1: Installing Nvidia CUDA Toolkit..."
    sudo apt update
    sudo apt install -y nvidia-driver-550 cuda-toolkit
    echo "NVIDIA CUDA Toolkit installation completed."
}

# Function to install dependencies for AUTOMATIC1111 Stable Diffusion
install_dependencies() {
    echo "Step 2: Installing required dependencies for AUTOMATIC1111 Stable Diffusion..."
    sudo apt update
    sudo apt install -y git libgl1-mesa-dev libfreetype6-dev pkg-config build-essential zlib1g-dev libncurses5-dev libgdbm-dev libnss3-dev libssl-dev libreadline-dev libffi-dev curl wget
    # Install Python 3.11 from source
    echo "Downloading and installing Python 3.11..."
    wget https://www.python.org/ftp/python/3.11.0/Python-3.11.0.tgz
    tar -xf Python-3.11.0.tgz
    cd Python-3.11.0
    ./configure --enable-optimizations
    make -j $(nproc)
    sudo make altinstall
    echo "Python 3.11 installation completed."
    
    # Install pip and dependencies for stable-diffusion-webui
    python3.11 -m pip install --upgrade pip
    echo "Installing Python packages..."
    pip3.11 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
    pip3.11 install gradio transformers diffusers
    echo "Dependencies installation completed."
}

# Function to clone and setup AUTOMATIC1111 Stable Diffusion web ui
install_stable_diffusion() {
    echo "Step 3: Cloning AUTOMATIC1111 Stable Diffusion web UI..."
    git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
    cd stable-diffusion-webui
    echo "Setting up web UI..."
    python3.11 launch.py --share
}

# Main menu
while true; do
    echo "=========================="
    echo "AUTOMATIC1111 Stable Diffusion Installer"
    echo "=========================="
    echo "Please select an option:"
    echo "1. Install Nvidia CUDA Toolkit (NVIDIA Driver 550)"
    echo "2. Install required dependencies for AUTOMATIC1111 Stable Diffusion (Python 3.11)"
    echo "3. Install AUTOMATIC1111 stable diffusion web ui itself"
    echo "4. Exit"

    read -p "Enter your choice [1-4]: " choice

    case $choice in
        1)
            install_cuda
            ;;
        2)
            install_dependencies
            ;;
        3)
            install_stable_diffusion
            ;;
        4)
            echo "Exiting..."
            exit 0
            ;;
        *)
            echo "Invalid option. Please try again."
            ;;
    esac
done