From 73debbaafeaf876f1e06f6798b1b68f5caefaa52 Mon Sep 17 00:00:00 2001
From: deffcolony <61471128+deffcolony@users.noreply.github.com>
Date: Mon, 26 Aug 2024 15:35:07 +0200
Subject: [PATCH] launcher upgrades

+added ostris ai toolkit. Feature Request by: @BlueprintCoding
---
 Launcher.bat                                  | 61 +++++++++++++++++++
 .../install_ostris_aitoolkit.bat              | 56 +++++++++++++++++
 .../uninstall_ostris_aitoolkit.bat            | 36 +++++++++++
 3 files changed, 153 insertions(+)
 create mode 100644 bin/functions/Toolbox/App_Installer/Image_Generation/install_ostris_aitoolkit.bat
 create mode 100644 bin/functions/Toolbox/App_Uninstaller/Image_Generation/uninstall_ostris_aitoolkit.bat

diff --git a/Launcher.bat b/Launcher.bat
index c626ab4..1ce1093 100644
--- a/Launcher.bat
+++ b/Launcher.bat
@@ -126,6 +126,8 @@ set "sdwebuiforge_install_path=%image_generation_dir%\stable-diffusion-webui-for
 set "comfyui_install_path=%image_generation_dir%\ComfyUI"
 set "fooocus_install_path=%image_generation_dir%\Fooocus"
 set "invokeai_install_path=%image_generation_dir%\InvokeAI"
+set "ostrisaitoolkit_install_path=%image_generation_dir%\ai-toolkit"
+
 
 REM Define variables for install locations (Text Completion)
 set "text_completion_dir=%~dp0text-completion"
@@ -943,6 +945,7 @@ echo    2. Update Stable Diffusion web UI forge
 echo    3. Update ComfyUI
 echo    4. Update Fooocus
 echo    5. Update InvokeAI
+echo    6. Update Ostris AI Toolkit
 echo %cyan_fg_strong% ______________________________________________________________%reset%
 echo %cyan_fg_strong%^| Menu Options:                                                ^|%reset%
 echo    0. Back
@@ -967,6 +970,8 @@ if "%update_manager_img_gen_choice%"=="1" (
     goto :update_fooocus
 ) else if "%update_manager_img_gen_choice%"=="5" (
     goto :update_invokeai
+) else if "%update_manager_img_gen_choice%"=="6" (
+    goto :update_ostrisaitoolkit
 ) else if "%update_manager_img_gen_choice%"=="0" (
     goto :update_manager
 ) else (
@@ -1108,6 +1113,36 @@ pause
 goto :update_manager_image_generation
 
 
+:update_ostrisaitoolkit
+REM Check if ai-toolkit directory exists
+if not exist "%ostrisaitoolkit_install_path%" (
+    echo %yellow_bg%[%time%]%reset% %yellow_fg_strong%[WARN] ai-toolkit directory not found. Skipping update.%reset%
+    pause
+    goto :update_manager_image_generation
+)
+
+REM Update Ostris AI Toolkit
+set max_retries=3
+set retry_count=0
+
+:retry_update_ostrisaitoolkit
+echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% Updating Ostris AI Toolkit...
+cd /d "%ostrisaitoolkit_install_path%"
+call git pull
+git submodule update --init --recursive
+if %errorlevel% neq 0 (
+    set /A retry_count+=1
+    echo %yellow_bg%[%time%]%reset% %yellow_fg_strong%[WARN] Retry %retry_count% of %max_retries%%reset%
+    if %retry_count% lss %max_retries% goto :retry_update_ostrisaitoolkit
+    echo %red_bg%[%time%]%reset% %red_fg_strong%[ERROR] Failed to update Ostris AI Toolkit repository after %max_retries% retries.%reset%
+    pause
+    goto :update_manager_image_generation
+)
+echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% %green_fg_strong%Ostris AI Toolkit updated successfully.%reset%
+pause
+goto :update_manager_image_generation
+
+
 REM ############################################################
 REM ######## UPDATE MANAGER CORE UTILITIES - FRONTEND #########
 REM ############################################################
@@ -2470,6 +2505,7 @@ echo    2. Stable Diffusion web UI forge [Install options]
 echo    3. Install ComfyUI
 echo    4. Install Fooocus
 echo    5. Install InvokeAI
+echo    6. Install Ostris AI Toolkit
 echo %cyan_fg_strong% ______________________________________________________________%reset%
 echo %cyan_fg_strong%^| Menu Options:                                                ^|%reset%
 echo    0. Back
@@ -2521,6 +2557,17 @@ if "%app_installer_image_generation_choice%"=="1" (
         pause
         goto :app_installer_image_generation
     )
+) else if "%app_installer_image_generation_choice%"=="6" (
+    set "caller=app_installer_image_generation"
+    if exist "%app_installer_image_generation_dir%\install_ostris_aitoolkit.bat" (
+        call %app_installer_image_generation_dir%\install_ostris_aitoolkit.bat
+        goto :app_installer_image_generation
+    ) else (
+        echo [%DATE% %TIME%] ERROR: install_ostris_aitoolkit.bat not found in: %app_installer_image_generation_dir% >> %logs_stl_console_path%
+        echo %red_bg%[%time%]%reset% %red_fg_strong%[ERROR] install_ostris_aitoolkit.bat not found in: %app_installer_image_generation_dir%%reset%
+        pause
+        goto :app_installer_image_generation
+    )
 ) else if "%app_installer_image_generation_choice%"=="0" (
     goto :app_installer
 ) else (
@@ -3474,6 +3521,7 @@ echo    2. UNINSTALL Stable Diffusion web UI forge
 echo    3. UNINSTALL ComfyUI
 echo    4. UNINSTALL Fooocus
 echo    5. UNINSTALL InvokeAI
+echo    6. UNINSTALL Ostris AI Toolkit
 echo %cyan_fg_strong% ______________________________________________________________%reset%
 echo %cyan_fg_strong%^| Menu Options:                                                ^|%reset%
 echo    0. Back
@@ -3553,6 +3601,19 @@ if "%app_uninstaller_img_gen_choice%"=="1" (
         pause
         goto :app_uninstaller_image_generation
     )
+) else if "%app_uninstaller_img_gen_choice%"=="6" (
+    set "caller=app_uninstaller_image_generation"
+    if exist "%app_uninstaller_image_generation_dir%\uninstall_ostris_aitoolkit.bat" (
+        call %app_uninstaller_image_generation_dir%\uninstall_ostris_aitoolkit.bat
+        goto :app_uninstaller_image_generation
+    ) else (
+        echo [%DATE% %TIME%] ERROR: uninstall_ostris_aitoolkit.bat not found in: %app_uninstaller_image_generation_dir% >> %logs_stl_console_path%
+        echo %red_bg%[%time%]%reset% %red_fg_strong%[ERROR] uninstall_ostris_aitoolkit.bat not found in: %app_uninstaller_image_generation_dir%%reset%
+        echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% Running Automatic Repair...
+        git pull
+        pause
+        goto :app_uninstaller_image_generation
+    )
 ) else if "%app_uninstaller_img_gen_choice%"=="0" (
     goto :app_uninstaller
 ) else (
diff --git a/bin/functions/Toolbox/App_Installer/Image_Generation/install_ostris_aitoolkit.bat b/bin/functions/Toolbox/App_Installer/Image_Generation/install_ostris_aitoolkit.bat
new file mode 100644
index 0000000..bb8b84a
--- /dev/null
+++ b/bin/functions/Toolbox/App_Installer/Image_Generation/install_ostris_aitoolkit.bat
@@ -0,0 +1,56 @@
+@echo off
+
+:install_ostrisaitoolkit
+title STL [INSTALL OSTRIS AI TOOLKIT]
+cls
+echo %blue_fg_strong%/ Home / Toolbox / App Installer / Image Generation / Install Ostris AI Toolkit%reset%
+echo -------------------------------------------------------------
+echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% Installing Ostris AI Toolkit...
+
+REM Check if the folder exists
+if not exist "%image_generation_dir%" (
+    mkdir "%image_generation_dir%"
+    echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% Created folder: "image-generation"  
+) else (
+    echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO] "image-generation" folder already exists.%reset%
+)
+cd /d "%image_generation_dir%"
+
+
+set max_retries=3
+set retry_count=0
+:retry_install_ostrisaitoolkit
+echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% Cloning the ostris ai-toolkit repository...
+git clone https://github.com/ostris/ai-toolkit.git
+
+if %errorlevel% neq 0 (
+    set /A retry_count+=1
+    echo %yellow_bg%[%time%]%reset% %yellow_fg_strong%[WARN] Retry %retry_count% of %max_retries%%reset%
+    if %retry_count% lss %max_retries% goto :retry_install_ostrisaitoolkit
+    echo %red_bg%[%time%]%reset% %red_fg_strong%[ERROR] Failed to clone repository after %max_retries% retries.%reset%
+    pause
+    goto :app_installer_image_generation
+)
+cd /d "%ostrisaitoolkit_install_path%"
+git submodule update --init --recursive
+
+REM Run conda activate from the Miniconda installation
+echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% Activating Miniconda environment...
+call "%miniconda_path%\Scripts\activate.bat"
+
+REM Create a Conda environment named ostrisaitoolkit
+echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% Creating Conda environment: %cyan_fg_strong%ostrisaitoolkit%reset%
+call conda create -n ostrisaitoolkit python=3.11 -y
+
+REM Activate the ostrisaitoolkit environment
+echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% Activating Conda environment: %cyan_fg_strong%ostrisaitoolkit%reset%
+call conda activate ostrisaitoolkit
+
+REM Install pip requirements
+echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% Installing pip requirements
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
+pip install -r requirements.txt
+
+echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% %green_fg_strong%Ostris AI Toolkit installed Successfully.%reset%
+pause
+goto :install_ostrisaitoolkit_menu
\ No newline at end of file
diff --git a/bin/functions/Toolbox/App_Uninstaller/Image_Generation/uninstall_ostris_aitoolkit.bat b/bin/functions/Toolbox/App_Uninstaller/Image_Generation/uninstall_ostris_aitoolkit.bat
new file mode 100644
index 0000000..26df88e
--- /dev/null
+++ b/bin/functions/Toolbox/App_Uninstaller/Image_Generation/uninstall_ostris_aitoolkit.bat
@@ -0,0 +1,36 @@
+@echo off
+
+:uninstall_ostrisaitoolkit
+title STL [UNINSTALL OSTRIS AI TOOLKIT]
+setlocal enabledelayedexpansion
+chcp 65001 > nul
+
+REM Confirm with the user before proceeding
+echo.
+echo %red_bg%╔════ DANGER ZONE ══════════════════════════════════════════════════════════════════════════════╗%reset%
+echo %red_bg%║ WARNING: This will delete all data of Ostris AI Toolkit                                       ║%reset%
+echo %red_bg%║ If you want to keep any data, make sure to create a backup before proceeding.                 ║%reset%
+echo %red_bg%╚═══════════════════════════════════════════════════════════════════════════════════════════════╝%reset%
+echo.
+set /p "confirmation=Are you sure you want to proceed? [Y/N]: "
+if /i "%confirmation%"=="Y" (
+
+    REM Remove the Conda environment
+    echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% Removing the Conda enviroment: %cyan_fg_strong%ostrisaitoolkit%reset%
+    call conda deactivate
+    call conda remove --name ostrisaitoolkit --all -y
+    call conda clean -a -y
+
+    REM Remove the folder ai-toolkit
+    echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% Removing the ai-toolkit directory...
+    cd /d "%~dp0"
+    rmdir /s /q "%ostrisaitoolkit_install_path%"
+
+    echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% %green_fg_strong%Ostris AI Toolkit has been uninstalled successfully.%reset%
+    pause
+    goto :app_uninstaller_image_generation
+) else (
+    echo %blue_bg%[%time%]%reset% %blue_fg_strong%[INFO]%reset% Uninstall canceled.
+    pause
+    goto :app_uninstaller_image_generation
+)
\ No newline at end of file