From 85b4f89926f7c3aaa7846dcbb47df3fd3b483b6b Mon Sep 17 00:00:00 2001 From: Aarni Koskela Date: Thu, 11 May 2023 23:46:45 +0300 Subject: [PATCH 1/2] Replace state.need_restart with state.server_command + replace poll loop with signal --- modules/shared.py | 42 +++++++++++++++++++++++++++++++++++++++- modules/ui.py | 6 +----- modules/ui_extensions.py | 7 ++----- webui.py | 39 +++++++++++++++++++++++-------------- 4 files changed, 68 insertions(+), 26 deletions(-) diff --git a/modules/shared.py b/modules/shared.py index 3abf71c0..648a2a19 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -2,6 +2,7 @@ import datetime import json import os import sys +import threading import time import gradio as gr @@ -110,8 +111,47 @@ class State: id_live_preview = 0 textinfo = None time_start = None - need_restart = False server_start = None + _server_command_signal = threading.Event() + _server_command: str | None = None + + @property + def need_restart(self) -> bool: + # Compatibility getter for need_restart. + return self.server_command == "restart" + + @need_restart.setter + def need_restart(self, value: bool) -> None: + # Compatibility setter for need_restart. + if value: + self.server_command = "restart" + + @property + def server_command(self): + return self._server_command + + @server_command.setter + def server_command(self, value: str | None) -> None: + """ + Set the server command to `value` and signal that it's been set. + """ + self._server_command = value + self._server_command_signal.set() + + def wait_for_server_command(self, timeout: float | None = None) -> str | None: + """ + Wait for server command to get set; return and clear the value and signal. + """ + if self._server_command_signal.wait(timeout): + self._server_command_signal.clear() + req = self._server_command + self._server_command = None + return req + return None + + def request_restart(self) -> None: + self.interrupt() + self.server_command = True def skip(self): self.skipped = True diff --git a/modules/ui.py b/modules/ui.py index 8e51e782..bed8464e 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -1609,12 +1609,8 @@ def create_ui(): outputs=[] ) - def request_restart(): - shared.state.interrupt() - shared.state.need_restart = True - restart_gradio.click( - fn=request_restart, + fn=shared.state.request_restart, _js='restart_reload', inputs=[], outputs=[], diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py index d7a0f685..4ba3bdd7 100644 --- a/modules/ui_extensions.py +++ b/modules/ui_extensions.py @@ -52,9 +52,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.interrupt() - shared.state.need_restart = True + shared.state.request_restart() def save_config_state(name): @@ -92,8 +90,7 @@ def restore_config_state(confirmed, config_state_name, restore_type): if restore_type == "webui" or restore_type == "both": config_states.restore_webui_config(config_state) - shared.state.interrupt() - shared.state.need_restart = True + shared.state.request_restart() return "" diff --git a/webui.py b/webui.py index 293a16cc..39dec3ca 100644 --- a/webui.py +++ b/webui.py @@ -234,7 +234,10 @@ def initialize(): print(f'Interrupted with signal {sig} in {frame}') os._exit(0) - signal.signal(signal.SIGINT, sigint_handler) + if not os.environ.get("COVERAGE_RUN"): + # Don't install the immediate-quit handler when running under coverage, + # as then the coverage report won't be generated. + signal.signal(signal.SIGINT, sigint_handler) def setup_middleware(app): @@ -255,19 +258,6 @@ def create_api(app): return api -def wait_on_server(demo=None): - while 1: - time.sleep(0.5) - if shared.state.need_restart: - shared.state.need_restart = False - time.sleep(0.5) - demo.close() - time.sleep(0.5) - - modules.script_callbacks.app_reload_callback() - break - - def api_only(): initialize() @@ -328,6 +318,7 @@ def webui(): inbrowser=cmd_opts.autolaunch, prevent_thread_lock=True ) + # after initial launch, disable --autolaunch for subsequent restarts cmd_opts.autolaunch = False @@ -359,8 +350,26 @@ def webui(): redirector.get("/") gradio.mount_gradio_app(redirector, shared.demo, path=f"/{cmd_opts.subpath}") - wait_on_server(shared.demo) + try: + while True: + server_command = shared.state.wait_for_server_command(timeout=5) + if server_command: + if server_command in ("stop", "restart"): + break + else: + print(f"Unknown server command: {server_command}") + except KeyboardInterrupt: + server_command = "stop" + + if server_command == "stop": + # If we catch a keyboard interrupt, we want to stop the server and exit. + print('Caught KeyboardInterrupt, stopping...') + shared.demo.close() + break print('Restarting UI...') + shared.demo.close() + time.sleep(0.5) + modules.script_callbacks.app_reload_callback() startup_timer.reset() From 875990a23213c63c19b8fdd3c87345f7a8ea2ceb Mon Sep 17 00:00:00 2001 From: Aarni Koskela Date: Tue, 16 May 2023 20:58:35 +0300 Subject: [PATCH 2/2] Add option for /_stop route (for graceful shutdown) --- modules/cmd_args.py | 1 + webui.py | 13 +++++++++++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/modules/cmd_args.py b/modules/cmd_args.py index f4a4ab36..6144db5c 100644 --- a/modules/cmd_args.py +++ b/modules/cmd_args.py @@ -103,3 +103,4 @@ parser.add_argument("--skip-version-check", action='store_true', help="Do not ch parser.add_argument("--no-hashing", action='store_true', help="disable sha256 hashing of checkpoints to help loading performance", default=False) parser.add_argument("--no-download-sd-model", action='store_true', help="don't download SD1.5 model even if no model is found in --ckpt-dir", default=False) parser.add_argument('--subpath', type=str, help='customize the subpath for gradio, use with reverse proxy') +parser.add_argument('--add-stop-route', action='store_true', help='add /_stop route to stop server') diff --git a/webui.py b/webui.py index 39dec3ca..5172f049 100644 --- a/webui.py +++ b/webui.py @@ -8,7 +8,7 @@ import warnings import json from threading import Thread -from fastapi import FastAPI +from fastapi import FastAPI, Response from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.gzip import GZipMiddleware from packaging import version @@ -270,6 +270,12 @@ def api_only(): print(f"Startup time: {startup_timer.summary()}.") api.launch(server_name="0.0.0.0" if cmd_opts.listen else "127.0.0.1", port=cmd_opts.port if cmd_opts.port else 7861) + +def stop_route(request): + shared.state.server_command = "stop" + return Response("Stopping.") + + def webui(): launch_api = cmd_opts.api initialize() @@ -318,6 +324,8 @@ def webui(): inbrowser=cmd_opts.autolaunch, prevent_thread_lock=True ) + if cmd_opts.add_stop_route: + app.add_route("/_stop", stop_route, methods=["POST"]) # after initial launch, disable --autolaunch for subsequent restarts cmd_opts.autolaunch = False @@ -359,11 +367,12 @@ def webui(): else: print(f"Unknown server command: {server_command}") except KeyboardInterrupt: + print('Caught KeyboardInterrupt, stopping...') server_command = "stop" if server_command == "stop": + print("Stopping server...") # If we catch a keyboard interrupt, we want to stop the server and exit. - print('Caught KeyboardInterrupt, stopping...') shared.demo.close() break print('Restarting UI...')