From 46a5bd64edece07f521409f0adace6cdd6f30a40 Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Mon, 5 Jun 2023 20:04:28 +0300
Subject: [PATCH] Restart: only do restart if running via the wrapper script

---
 modules/restart.py       | 23 +++++++++++++++++++++++
 modules/shared.py        |  9 ---------
 modules/ui_extensions.py | 11 ++++++++---
 webui.bat                |  2 +-
 webui.sh                 |  1 +
 5 files changed, 33 insertions(+), 13 deletions(-)
 create mode 100644 modules/restart.py

diff --git a/modules/restart.py b/modules/restart.py
new file mode 100644
index 00000000..18eacaf3
--- /dev/null
+++ b/modules/restart.py
@@ -0,0 +1,23 @@
+import os
+from pathlib import Path
+
+from modules.paths_internal import script_path
+
+
+def is_restartable() -> bool:
+    """
+    Return True if the webui is restartable (i.e. there is something watching to restart it with)
+    """
+    return bool(os.environ.get('SD_WEBUI_RESTART'))
+
+
+def restart_program() -> None:
+    """creates file tmp/restart and immediately stops the process, which webui.bat/webui.sh interpret as a command to start webui again"""
+
+    (Path(script_path) / "tmp" / "restart").touch()
+
+    stop_program()
+
+
+def stop_program() -> None:
+    os._exit(0)
diff --git a/modules/shared.py b/modules/shared.py
index 2bd7c6ec..c9ee2dd1 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -853,12 +853,3 @@ 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 5580dfaf..3d216912 100644
--- a/modules/ui_extensions.py
+++ b/modules/ui_extensions.py
@@ -11,7 +11,7 @@ import html
 import shutil
 import errno
 
-from modules import extensions, shared, paths, config_states, errors
+from modules import extensions, shared, paths, config_states, errors, restart
 from modules.paths_internal import config_states_dir
 from modules.call_queue import wrap_gradio_gpu_call
 
@@ -49,7 +49,11 @@ 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.restart_program()
+
+    if restart.is_restartable():
+        restart.restart_program()
+    else:
+        restart.stop_program()
 
 
 def save_config_state(name):
@@ -508,7 +512,8 @@ def create_ui():
             with gr.TabItem("Installed", id="installed"):
 
                 with gr.Row(elem_id="extensions_installed_top"):
-                    apply = gr.Button(value="Apply and restart UI", variant="primary")
+                    apply_label = ("Apply and restart UI" if restart.is_restartable() else "Apply and quit")
+                    apply = gr.Button(value=apply_label, variant="primary")
                     check = gr.Button(value="Check for updates")
                     extensions_disable_all = gr.Radio(label="Disable all extensions", choices=["none", "extra", "all"], value=shared.opts.disable_all_extensions, elem_id="extensions_disable_all")
                     extensions_disabled_list = gr.Text(elem_id="extensions_disabled_list", visible=False).style(container=False)
diff --git a/webui.bat b/webui.bat
index 961fc7d4..42e7d517 100644
--- a/webui.bat
+++ b/webui.bat
@@ -3,7 +3,7 @@
 if not defined PYTHON (set PYTHON=python)
 if not defined VENV_DIR (set "VENV_DIR=%~dp0%venv")
 
-
+set SD_WEBUI_RESTART=tmp/restart
 set ERROR_REPORTING=FALSE
 
 mkdir tmp 2>NUL
diff --git a/webui.sh b/webui.sh
index c407b3ef..6c48e969 100755
--- a/webui.sh
+++ b/webui.sh
@@ -204,6 +204,7 @@ prepare_tcmalloc() {
 }
 
 KEEP_GOING=1
+export SD_WEBUI_RESTART=tmp/restart
 while [[ "$KEEP_GOING" -eq "1" ]]; do
     if [[ ! -z "${ACCELERATE}" ]] && [ ${ACCELERATE}="True" ] && [ -x "$(command -v accelerate)" ]; then
         printf "\n%s\n" "${delimiter}"