print progress when installing torch

add PIP_INSTALLER_LOCATION env var to install pip if it's not installed
remove accidental call to accelerate when venv is disabled
add another env var to skip venv - SKIP_VENV

2 files changed, 22 insertions, 6 deletions

diff --git a/launch.py b/launch.py
index 6d523a34..f578c1c7 100644
--- a/launch.py
+++ b/launch.py
@@ -48,10 +48,19 @@ def extract_opt(args, name):
     return args, is_present, opt

 

 

-def run(command, desc=None, errdesc=None, custom_env=None):

+def run(command, desc=None, errdesc=None, custom_env=None, live=False):

     if desc is not None:

         print(desc)

 

+    if live:

+        result = subprocess.run(command, shell=True, env=os.environ if custom_env is None else custom_env)

+        if result.returncode != 0:

+            raise RuntimeError(f"""{errdesc or 'Error running command'}.

+Command: {command}

+Error code: {result.returncode}""")

+

+        return ""

+

     result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, env=os.environ if custom_env is None else custom_env)

 

     if result.returncode != 0:

@@ -179,6 +188,8 @@ def run_extensions_installers(settings_file):
 def prepare_environment():

     global skip_install

 

+    pip_installer_location = os.environ.get('PIP_INSTALLER_LOCATION', None)

+

     torch_command = os.environ.get('TORCH_COMMAND', "pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117")

     requirements_file = os.environ.get('REQS_FILE', "requirements_versions.txt")

     commandline_args = os.environ.get('COMMANDLINE_ARGS', "")

@@ -219,9 +230,12 @@ def prepare_environment():
 

     print(f"Python {sys.version}")

     print(f"Commit hash: {commit}")

-    

+

+    if pip_installer_location is not None and not is_installed("pip"):

+        run(f'"{python}" "{pip_installer_location}"', "Installing pip", "Couldn't install pip")

+

     if reinstall_torch or not is_installed("torch") or not is_installed("torchvision"):

-        run(f'"{python}" -m {torch_command}', "Installing torch and torchvision", "Couldn't install torch")

+        run(f'"{python}" -m {torch_command}', "Installing torch and torchvision", "Couldn't install torch", live=True)

 

     if not skip_torch_cuda_test:

         run_python("import torch; assert torch.cuda.is_available(), 'Torch is not able to use GPU; add --skip-torch-cuda-test to COMMANDLINE_ARGS variable to disable this check'")

diff --git a/webui.bat b/webui.bat
index 3165b94d..0d6865c9 100644
--- a/webui.bat
+++ b/webui.bat
@@ -3,6 +3,7 @@
 if not defined PYTHON (set PYTHON=python)

 if not defined VENV_DIR (set "VENV_DIR=%~dp0%venv")

 

+

 set ERROR_REPORTING=FALSE

 

 mkdir tmp 2>NUL

@@ -14,6 +15,7 @@ goto :show_stdout_stderr
 

 :start_venv

 if ["%VENV_DIR%"] == ["-"] goto :skip_venv

+if ["%SKIP_VENV%"] == ["1"] goto :skip_venv

 

 dir "%VENV_DIR%\Scripts\Python.exe" >tmp/stdout.txt 2>tmp/stderr.txt

 if %ERRORLEVEL% == 0 goto :activate_venv

@@ -28,13 +30,13 @@ goto :show_stdout_stderr
 :activate_venv

 set PYTHON="%VENV_DIR%\Scripts\Python.exe"

 echo venv %PYTHON%

-if [%ACCELERATE%] == ["True"] goto :accelerate

-goto :launch

 

 :skip_venv

+if [%ACCELERATE%] == ["True"] goto :accelerate

+goto :launch

 

 :accelerate

-echo "Checking for accelerate"

+echo Checking for accelerate

 set ACCELERATE="%VENV_DIR%\Scripts\accelerate.exe"

 if EXIST %ACCELERATE% goto :accelerate_launch