From 689195cc93559810bd6c1acf2bb1553bf46e34f1 Mon Sep 17 00:00:00 2001
From: bluelovers <codelovers@users.sourceforge.net>
Date: Tue, 4 Jun 2024 17:40:01 +0800
Subject: [PATCH] feat: support restart webui

https://github.com/ArtVentureX/sd-webui-agent-scheduler/issues/243
---
 agent_scheduler/task_runner.py | 21 +++++++++++++++++++++
 scripts/task_scheduler.py      |  2 +-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/agent_scheduler/task_runner.py b/agent_scheduler/task_runner.py
index 12c65dc..36a8da4 100644
--- a/agent_scheduler/task_runner.py
+++ b/agent_scheduler/task_runner.py
@@ -620,6 +620,27 @@ def __on_completed(self):
         elif action == "Stop webui":
             log.info("[AgentScheduler] Stopping webui...")
             _exit(0)
+        elif action == "Restart webui":
+            log.info("[AgentScheduler] Restarting webui...")
+            with gr.Blocks() as demo:
+                gr.HTML("""
+    <script type="text/javascript">
+    console.info(`Restart webui...after 5sec`);
+    
+    var requestPing = function() {
+        requestGet("./internal/ping", {}, function(data) {
+            location.reload();
+        }, function() {
+            console.info(`Waiting webui start...`);
+            setTimeout(requestPing, 500);
+        });
+    };
+    
+    setTimeout(requestPing, 5000);
+    </script>
+                """)
+            shared.state.request_restart()
+            demo.launch()
 
         if command:
             subprocess.Popen(command)
diff --git a/scripts/task_scheduler.py b/scripts/task_scheduler.py
index 3198d35..42d5793 100644
--- a/scripts/task_scheduler.py
+++ b/scripts/task_scheduler.py
@@ -41,7 +41,7 @@
 placement_under_generate = "Under Generate button"
 placement_between_prompt_and_generate = "Between Prompt and Generate button"
 
-completion_action_choices = ["Do nothing", "Shut down", "Restart", "Sleep", "Hibernate", "Stop webui"]
+completion_action_choices = ["Do nothing", "Shut down", "Restart", "Sleep", "Hibernate", "Stop webui", "Restart webui"]
 
 task_filter_choices = ["All", "Bookmarked", "Done", "Failed", "Interrupted"]