From 333e63c0911c148ea306d7b72580d5c6d2f2c41a Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sat, 3 Jun 2023 09:59:56 +0300
Subject: [PATCH] a yet another method to restart webui

---
 modules/launch_utils.py  |  6 ++++++
 modules/shared.py        |  9 +++++++++
 modules/ui_extensions.py |  2 +-
 webui.bat                |  2 ++
 webui.sh                 | 34 ++++++++++++++++++++--------------
 5 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/modules/launch_utils.py b/modules/launch_utils.py
index 0c8c4db0b..af8d8b378 100644
--- a/modules/launch_utils.py
+++ b/modules/launch_utils.py
@@ -244,6 +244,12 @@ def prepare_environment():
     codeformer_commit_hash = os.environ.get('CODEFORMER_COMMIT_HASH', "c5b4593074ba6214284d6acd5f1719b6c5d739af")
     blip_commit_hash = os.environ.get('BLIP_COMMIT_HASH', "48211a1594f1321b00f14c9f7a5b4813144b2fb9")
 
+    try:
+        # the existance of this file is a signal to webui.sh/bat that webui needs to be restarted when it stops execution
+        os.remove(os.path.join(script_path, "tmp", "restart"))
+    except OSError:
+        pass
+
     if not args.skip_python_version_check:
         check_python_version()
 
diff --git a/modules/shared.py b/modules/shared.py
index 7025a7547..c4c719ad1 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -841,3 +841,12 @@ def walk_files(path, allowed_extensions=None):
                 continue
 
             yield os.path.join(root, filename)
+
+
+def restart_program():
+    """creates file tmp/restart and immediately stops the process, which webui.bat/webui.sh interpret as a command to start webui again"""
+
+    with open(os.path.join(script_path, "tmp", "restart"), "w"):
+        pass
+
+    os._exit(0)
diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py
index 3140ed649..5580dfafe 100644
--- a/modules/ui_extensions.py
+++ b/modules/ui_extensions.py
@@ -49,7 +49,7 @@ def apply_and_restart(disable_list, update_list, disable_all):
     shared.opts.disabled_extensions = disabled
     shared.opts.disable_all_extensions = disable_all
     shared.opts.save(shared.config_filename)
-    shared.state.request_restart()
+    shared.restart_program()
 
 
 def save_config_state(name):
diff --git a/webui.bat b/webui.bat
index 209d972bd..961fc7d4c 100644
--- a/webui.bat
+++ b/webui.bat
@@ -51,12 +51,14 @@ if EXIST %ACCELERATE% goto :accelerate_launch
 
 :launch
 %PYTHON% launch.py %*
+if EXIST tmp/restart goto :skip_venv
 pause
 exit /b
 
 :accelerate_launch
 echo Accelerating
 %ACCELERATE% launch --num_cpu_threads_per_process=6 launch.py
+if EXIST tmp/restart goto :skip_venv
 pause
 exit /b
 
diff --git a/webui.sh b/webui.sh
index 1e7288136..c407b3efe 100755
--- a/webui.sh
+++ b/webui.sh
@@ -203,17 +203,23 @@ prepare_tcmalloc() {
     fi
 }
 
-if [[ ! -z "${ACCELERATE}" ]] && [ ${ACCELERATE}="True" ] && [ -x "$(command -v accelerate)" ]
-then
-    printf "\n%s\n" "${delimiter}"
-    printf "Accelerating launch.py..."
-    printf "\n%s\n" "${delimiter}"
-    prepare_tcmalloc
-    exec accelerate launch --num_cpu_threads_per_process=6 "${LAUNCH_SCRIPT}" "$@"
-else
-    printf "\n%s\n" "${delimiter}"
-    printf "Launching launch.py..."
-    printf "\n%s\n" "${delimiter}"
-    prepare_tcmalloc
-    exec "${python_cmd}" "${LAUNCH_SCRIPT}" "$@"
-fi
+KEEP_GOING=1
+while [[ "$KEEP_GOING" -eq "1" ]]; do
+    if [[ ! -z "${ACCELERATE}" ]] && [ ${ACCELERATE}="True" ] && [ -x "$(command -v accelerate)" ]; then
+        printf "\n%s\n" "${delimiter}"
+        printf "Accelerating launch.py..."
+        printf "\n%s\n" "${delimiter}"
+        prepare_tcmalloc
+        accelerate launch --num_cpu_threads_per_process=6 "${LAUNCH_SCRIPT}" "$@"
+    else
+        printf "\n%s\n" "${delimiter}"
+        printf "Launching launch.py..."
+        printf "\n%s\n" "${delimiter}"
+        prepare_tcmalloc
+        "${python_cmd}" "${LAUNCH_SCRIPT}" "$@"
+    fi
+
+    if [[ ! -f tmp/restart ]]; then
+        KEEP_GOING=0
+    fi
+done