From 85b4f89926f7c3aaa7846dcbb47df3fd3b483b6b Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Thu, 11 May 2023 23:46:45 +0300
Subject: Replace state.need_restart with state.server_command + replace poll
 loop with signal

---
 webui.py | 39 ++++++++++++++++++++++++---------------
 1 file changed, 24 insertions(+), 15 deletions(-)

(limited to 'webui.py')

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()
 
-- 
cgit v1.2.3


From 875990a23213c63c19b8fdd3c87345f7a8ea2ceb Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Tue, 16 May 2023 20:58:35 +0300
Subject: Add option for /_stop route (for graceful shutdown)

---
 webui.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

(limited to 'webui.py')

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...')
-- 
cgit v1.2.3