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

api quit restart

4 files changed, 21 insertions, 11 deletions

diff --git a/.github/workflows/run_tests.yaml b/.github/workflows/run_tests.yaml
index 226cf759..96546011 100644
--- a/.github/workflows/run_tests.yaml
+++ b/.github/workflows/run_tests.yaml
@@ -50,7 +50,7 @@ jobs:
           python -m pytest -vv --junitxml=test/results.xml --cov . --cov-report=xml --verify-base-url test
       - name: Kill test server
         if: always()
-        run: curl -vv -XPOST http://127.0.0.1:7860/_stop && sleep 10
+        run: curl -vv -XPOST http://127.0.0.1:7860/sdapi/v1/server-stop && sleep 10
       - name: Show coverage
         run: |
           python -m coverage combine .coverage*
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')

diff --git a/webui.py b/webui.py
index 136d036d..bad29f28 100644
--- a/webui.py
+++ b/webui.py
@@ -11,7 +11,7 @@ import json
 from threading import Thread

 from typing import Iterable

 

-from fastapi import FastAPI, Response

+from fastapi import FastAPI

 from fastapi.middleware.cors import CORSMiddleware

 from fastapi.middleware.gzip import GZipMiddleware

 from packaging import version

@@ -362,11 +362,6 @@ def api_only():
     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()

@@ -404,8 +399,6 @@ def webui():
                 "redoc_url": "/redoc",

             },

         )

-        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