Upgrade to PyTorch 2.0; Doesn't use a special repo url for pytorch on Linux

scripts/check_modules.py

@@ -1,13 +1,23 @@
-'''
+"""
 This script checks if the given modules exist
-'''
+
+E.g. python check_modules.py sdkit==1.0.3 sdkit.models ldm transformers numpy antlr4 gfpgan realesrgan
+"""
 
 import sys
 import pkgutil
+from importlib.metadata import version
 
 modules = sys.argv[1:]
 missing_modules = []
 for m in modules:
-    if pkgutil.find_loader(m) is None:
+    m = m.split("==")
-        print('module', m, 'not found')
+    module_name = m[0]
+    module_version = m[1] if len(m) > 1 else None
+    is_installed = pkgutil.find_loader(module_name) is not None
+    if not is_installed:
+        print("module", module_name, "not found")
+        exit(1)
+    elif module_version and version(module_name) != module_version:
+        print("module version is different! expected: ", module_version, ", actual: ", version(module_name))
         exit(1)

scripts/on_sd_start.bat

@@ -7,6 +7,11 @@
 @copy sd-ui-files\scripts\bootstrap.bat scripts\ /Y
 @copy sd-ui-files\scripts\check_modules.py scripts\ /Y
 
+set TORCH_VERSION=2.0.0+cu117
+set TORCHVISION_VERSION=0.15.1+cu117
+set SDKIT_VERSION=1.0.71
+set SD_VERSION=2.1.4
+
 if exist "%cd%\profile" (
     set USERPROFILE=%cd%\profile
 )
@@ -65,7 +70,7 @@ if not exist "%INSTALL_ENV_DIR%\DLLs\libssl-1_1-x64.dll"    copy "%INSTALL_ENV_D
 if not exist "%INSTALL_ENV_DIR%\DLLs\libcrypto-1_1-x64.dll" copy "%INSTALL_ENV_DIR%\Library\bin\libcrypto-1_1-x64.dll" "%INSTALL_ENV_DIR%\DLLs\"
 
 @rem install torch and torchvision
-call python ..\scripts\check_modules.py torch torchvision
+call python ..\scripts\check_modules.py torch==%TORCH_VERSION% torchvision==%TORCHVISION_VERSION%
 if "%ERRORLEVEL%" EQU "0" (
     echo "torch and torchvision have already been installed."
 ) else (
@@ -75,13 +80,15 @@ if "%ERRORLEVEL%" EQU "0" (
     set PYTHONNOUSERSITE=1
     set PYTHONPATH=%INSTALL_ENV_DIR%\lib\site-packages
 
-    call python -m pip install --upgrade torch==1.13.1+cu116 torchvision==0.14.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116 || (
+    call python -m pip install --upgrade torch==%TORCH_VERSION% torchvision==%TORCHVISION_VERSION% --index-url https://download.pytorch.org/whl/cu117 || (
         echo "Error installing torch. Sorry about that, please try to:" & echo "  1. Run this installer again." & echo "  2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo "  3. If those steps don't help, please copy *all* the error messages in this window, and ask the community at https://discord.com/invite/u9yhsFmEkB" & echo "  4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!"
         pause
         exit /b
     )
 )
 
+call python -c "from importlib.metadata import version; print('torch version:', version('torch'))"
+
 set PATH=C:\Windows\System32;%PATH%
 
 @rem install/upgrade sdkit
@@ -95,7 +102,7 @@ if "%ERRORLEVEL%" EQU "0" (
         set PYTHONNOUSERSITE=1
         set PYTHONPATH=%INSTALL_ENV_DIR%\lib\site-packages
 
-        call python -m pip install --upgrade sdkit==1.0.70 -q || (
+        call python -m pip install --upgrade sdkit==%SDKIT_VERSION% -q || (
             echo "Error updating sdkit"
         )
     )
@@ -106,7 +113,7 @@ if "%ERRORLEVEL%" EQU "0" (
     set PYTHONNOUSERSITE=1
     set PYTHONPATH=%INSTALL_ENV_DIR%\lib\site-packages
 
-    call python -m pip install sdkit==1.0.70 || (
+    call python -m pip install sdkit==%SDKIT_VERSION% || (
         echo "Error installing sdkit. Sorry about that, please try to:" & echo "  1. Run this installer again." & echo "  2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo "  3. If those steps don't help, please copy *all* the error messages in this window, and ask the community at https://discord.com/invite/u9yhsFmEkB" & echo "  4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!"
         pause
         exit /b
@@ -116,7 +123,7 @@ if "%ERRORLEVEL%" EQU "0" (
 call python -c "from importlib.metadata import version; print('sdkit version:', version('sdkit'))"
 
 @rem upgrade stable-diffusion-sdkit
-call python -m pip install --upgrade stable-diffusion-sdkit==2.1.4 -q || (
+call python -m pip install --upgrade stable-diffusion-sdkit==%SD_VERSION% -q || (
     echo "Error updating stable-diffusion-sdkit"
 )
 call python -c "from importlib.metadata import version; print('stable-diffusion version:', version('stable-diffusion-sdkit'))"

scripts/on_sd_start.sh

@@ -7,6 +7,11 @@ cp sd-ui-files/scripts/check_modules.py scripts/
 
 source ./scripts/functions.sh
 
+TORCH_VERSION="2.0.0"
+TORCHVISION_VERSION="0.15.1"
+SDKIT_VERSION="1.0.71"
+SD_VERSION="2.1.4"
+
 # activate the installer env
 CONDA_BASEPATH=$(conda info --base)
 source "$CONDA_BASEPATH/etc/profile.d/conda.sh" # avoids the 'shell not initialized' error
@@ -64,14 +69,7 @@ case "${OS_NAME}" in
 esac
 
 # install torch and torchvision
-if python ../scripts/check_modules.py torch torchvision; then
+if python ../scripts/check_modules.py torch==$TORCH_VERSION torchvision==$TORCHVISION_VERSION; then
-    # temp fix for installations that installed torch 2.0 by mistake
-    if [ "$OS_NAME" == "linux" ]; then
-        python -m pip install --upgrade torch==1.13.1+cu116 torchvision==0.14.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116 -q
-    elif [ "$OS_NAME" == "macos" ]; then
-        python -m pip install --upgrade torch==1.13.1 torchvision==0.14.1 -q
-    fi
-
     echo "torch and torchvision have already been installed."
 else
     echo "Installing torch and torchvision.."
@@ -79,20 +77,14 @@ else
     export PYTHONNOUSERSITE=1
     export PYTHONPATH="$INSTALL_ENV_DIR/lib/python3.8/site-packages"
 
-    if [ "$OS_NAME" == "linux" ]; then
+    if python -m pip install --upgrade torch==$TORCH_VERSION torchvision==$TORCHVISION_VERSION ; then
-        if python -m pip install --upgrade torch==1.13.1+cu116 torchvision==0.14.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116 ; then
-            echo "Installed."
-        else
-            fail "torch install failed"
-        fi
-    elif [ "$OS_NAME" == "macos" ]; then
-        if python -m pip install --upgrade torch==1.13.1 torchvision==0.14.1 ; then
         echo "Installed."
     else
         fail "torch install failed"
     fi
 fi
-fi
+
+python -c "from importlib.metadata import version; print('torch version:', version('torch'))"
 
 # install/upgrade sdkit
 if python ../scripts/check_modules.py sdkit sdkit.models ldm transformers numpy antlr4 gfpgan realesrgan ; then
@@ -103,7 +95,7 @@ if python ../scripts/check_modules.py sdkit sdkit.models ldm transformers numpy
         export PYTHONNOUSERSITE=1
         export PYTHONPATH="$INSTALL_ENV_DIR/lib/python3.8/site-packages"
 
-        python -m pip install --upgrade sdkit==1.0.70 -q
+        python -m pip install --upgrade sdkit==$SDKIT_VERSION -q
     fi
 else
     echo "Installing sdkit: https://pypi.org/project/sdkit/"
@@ -111,7 +103,7 @@ else
     export PYTHONNOUSERSITE=1
     export PYTHONPATH="$INSTALL_ENV_DIR/lib/python3.8/site-packages"
 
-    if python -m pip install sdkit==1.0.70 ; then
+    if python -m pip install sdkit==$SDKIT_VERSION ; then
         echo "Installed."
     else
         fail "sdkit install failed"
@@ -121,7 +113,7 @@ fi
 python -c "from importlib.metadata import version; print('sdkit version:', version('sdkit'))"
 
 # upgrade stable-diffusion-sdkit
-python -m pip install --upgrade stable-diffusion-sdkit==2.1.4 -q
+python -m pip install --upgrade stable-diffusion-sdkit==$SD_VERSION -q
 python -c "from importlib.metadata import version; print('stable-diffusion version:', version('stable-diffusion-sdkit'))"
 
 # install rich