Merge pull request #11146 from AUTOMATIC1111/api-quit-restart

api quit restart

2 files changed, 19 insertions, 2 deletions

diff --git a/modules/api/api.py b/modules/api/api.py
index 2e49526e..80d45ca7 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -14,7 +14,7 @@ from fastapi.encoders import jsonable_encoder
 from secrets import compare_digest
 
 import modules.shared as shared
-from modules import sd_samplers, deepbooru, sd_hijack, images, scripts, ui, postprocessing, errors
+from modules import sd_samplers, deepbooru, sd_hijack, images, scripts, ui, postprocessing, errors, restart
 from modules.api import models
 from modules.shared import opts
 from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images
@@ -209,6 +209,11 @@ class Api:
         self.add_api_route("/sdapi/v1/scripts", self.get_scripts_list, methods=["GET"], response_model=models.ScriptsList)
         self.add_api_route("/sdapi/v1/script-info", self.get_script_info, methods=["GET"], response_model=List[models.ScriptInfo])
 
+        if shared.cmd_opts.add_stop_route:
+            self.add_api_route("/sdapi/v1/server-kill", self.kill_webui, methods=["POST"])
+            self.add_api_route("/sdapi/v1/server-restart", self.restart_webui, methods=["POST"])
+            self.add_api_route("/sdapi/v1/server-stop", self.stop_webui, methods=["POST"])
+
         self.default_script_arg_txt2img = []
         self.default_script_arg_img2img = []
 
@@ -715,3 +720,15 @@ class Api:
     def launch(self, server_name, port):
         self.app.include_router(self.router)
         uvicorn.run(self.app, host=server_name, port=port, timeout_keep_alive=0)
+
+    def kill_webui(self):
+        restart.stop_program()
+
+    def restart_webui(self):
+        if restart.is_restartable():
+            restart.restart_program()
+        return Response(status_code=501)
+
+    def stop_webui(request):
+        shared.state.server_command = "stop"
+        return Response("Stopping.")
diff --git a/modules/cmd_args.py b/modules/cmd_args.py
index de905caa..624dcb4f 100644
--- a/modules/cmd_args.py
+++ b/modules/cmd_args.py
@@ -106,4 +106,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')

+parser.add_argument('--add-stop-route', action='store_true', help='enable server stop/restart/kill via api')