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 1ae516d7..4379a641 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): @@ -509,7 +513,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}"