StableDiffusion/WebUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add option for /_stop route (for graceful shutdown)

Browse Source
This commit is contained in:
Aarni Koskela 2023-05-16 20:58:35 +03:00
parent 85b4f89926
commit 875990a232
2 changed files with 12 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
modules/cmd_args.py
View File

@ -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-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("--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('--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')

13
webui.py
View File

@ -8,7 +8,7 @@ import warnings
import json import json
from threading import Thread from threading import Thread
from fastapi import FastAPI from fastapi import FastAPI, Response
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.gzip import GZipMiddleware from fastapi.middleware.gzip import GZipMiddleware
from packaging import version from packaging import version
@ -270,6 +270,12 @@ def api_only():
print(f"Startup time: {startup_timer.summary()}.") 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) 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(): def webui():
launch_api = cmd_opts.api launch_api = cmd_opts.api
initialize() initialize()
@ -318,6 +324,8 @@ def webui():
inbrowser=cmd_opts.autolaunch, inbrowser=cmd_opts.autolaunch,
prevent_thread_lock=True 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 # after initial launch, disable --autolaunch for subsequent restarts
cmd_opts.autolaunch = False cmd_opts.autolaunch = False
@ -359,11 +367,12 @@ def webui():
else: else:
print(f"Unknown server command: {server_command}") print(f"Unknown server command: {server_command}")
except KeyboardInterrupt: except KeyboardInterrupt:
print('Caught KeyboardInterrupt, stopping...')
server_command = "stop" server_command = "stop"
if server_command == "stop": if server_command == "stop":
print("Stopping server...")
# If we catch a keyboard interrupt, we want to stop the server and exit. # If we catch a keyboard interrupt, we want to stop the server and exit.
print('Caught KeyboardInterrupt, stopping...')
shared.demo.close() shared.demo.close()
break break
print('Restarting UI...') print('Restarting UI...')