From 93fad28a979727f9b1331dbdc447598824057cdc Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Tue, 24 Jan 2023 21:13:05 +0300
Subject: 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

---
 launch.py | 20 +++++++++++++++++---
 webui.bat |  8 +++++---
 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
 
-- 
cgit v1.2.1