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

---
 modules/shared.py        | 42 +++++++++++++++++++++++++++++++++++++++++-
 modules/ui.py            |  6 +-----
 modules/ui_extensions.py |  7 ++-----
 3 files changed, 44 insertions(+), 11 deletions(-)

(limited to 'modules')

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