From 1dec4f027b102647317e930c66221c722d1739e2 Mon Sep 17 00:00:00 2001
From: erew123 <35898566+erew123@users.noreply.github.com>
Date: Tue, 12 Mar 2024 11:34:35 +0000
Subject: [PATCH] Added Windows DeepSpeed update for PyTorch 2.2.1

---
 atsetup.bat | 51 +++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 45 insertions(+), 6 deletions(-)

diff --git a/atsetup.bat b/atsetup.bat
index 9baf649..696c7e2 100644
--- a/atsetup.bat
+++ b/atsetup.bat
@@ -66,12 +66,18 @@ echo    2) Install the requirements for a %L_GREEN%AMD or MAC machine%RESET%.
 echo.
 echo    OPTIONAL
 echo    3) Install the Finetuning requirements.
-echo    4) Install DeepSpeed v11.2 for CUDA %L_GREEN%11.8%RESET% and Python-3.11.x.
-echo    5) Install DeepSpeed v11.2 for CUDA %L_GREEN%12.1%RESET% and Python-3.11.x.
-echo    6) Uninstall DeepSpeed.
+echo.
+echo    DEEPSPEED FOR %L_YELLOW=%PyTorch 2.1.x%RESET%
+echo    4) Install DeepSpeed v11.2 for CUDA %L_GREEN%11.8%RESET% and Python-3.11.x and %L_YELLOW%PyTorch 2.1.x%RESET%.
+echo    5) Install DeepSpeed v11.2 for CUDA %L_GREEN%12.1%RESET% and Python-3.11.x and %L_YELLOW%PyTorch 2.1.x%RESET%.
+echo.
+echo    DEEPSPEED FOR %L_YELLOW=%PyTorch 2.2.x%RESET% (March 2024 builds of Text-gen-webui and later)
+echo    6) Install DeepSpeed v14.0 for CUDA %L_GREEN%12.1%RESET% and Python-3.11.x and %L_YELLOW%PyTorch 2.2.x%RESET%.
+echo.
+echo    7) Uninstall DeepSpeed.
 echo.
 echo.   OTHER
-echo    7) Generate a diagnostics file.
+echo    8) Generate a diagnostics file.
 echo.
 echo    9) %L_RED%Exit/Quit%RESET%
 echo.
@@ -81,8 +87,9 @@ if "%WebUIOption%"=="2" goto InstallOtherTextGen
 if "%WebUIOption%"=="3" goto InstallFinetuneTextGen
 if "%WebUIOption%"=="4" goto InstallDeepSpeed118TextGen
 if "%WebUIOption%"=="5" goto InstallDeepSpeed121TextGen
-if "%WebUIOption%"=="6" goto UnInstallDeepSpeed
-if "%WebUIOption%"=="7" goto GenerateDiagsTextGen
+if "%WebUIOption%"=="6" goto InstallDeepSpeed121TextGenPytorch221
+if "%WebUIOption%"=="7" goto UnInstallDeepSpeed
+if "%WebUIOption%"=="8" goto GenerateDiagsTextGen
 if "%WebUIOption%"=="9" goto End
 goto WebUIMenu
 
@@ -227,6 +234,38 @@ del deepspeed-0.11.2+cuda121-cp311-cp311-win_amd64.whl
 pause
 goto WebUIMenu
 
+:InstallDeepSpeed121TextGenPytorch221
+echo Downloading DeepSpeed...
+curl -LO https://github.com/erew123/alltalk_tts/releases/download/DeepSpeed-14.0/deepspeed-0.14.0+ce78a63-cp311-cp311-win_amd64.whl
+if %ERRORLEVEL% neq 0 (
+    echo.
+    echo    Failed to download DeepSpeed wheel file.
+    echo    Please check your internet connection or try again later.
+    echo    Press any key to return to the menu.
+    echo.
+    pause
+    goto WebUIMenu
+)
+echo    DeepSpeed wheel file downloaded successfully.
+echo    Installing DeepSpeed...
+pip install deepspeed-0.14.0+ce78a63-cp311-cp311-win_amd64.whl
+if %ERRORLEVEL% neq 0 (
+    echo. 
+    echo    Failed to install DeepSpeed.
+    echo    Please check if the wheel file is compatible with your system.
+    echo    Press any key to return to the menu.
+    echo.
+    pause
+    goto WebUIMenu
+)
+Echo.
+echo.
+Echo    DeepSpeed installed successfully.
+Echo. 
+del deepspeed-0.14.0+ce78a63-cp311-cp311-win_amd64.whl
+pause
+goto WebUIMenu
+
 :UnInstallDeepSpeed
 pip uninstall deepspeed
 if %ERRORLEVEL% neq 0 (