feat: add button to reconnect UI without having to reload the page (#…

…2727) * feat: add button to reconnect UI without having to reload the page * qa: add missing semicolon
lllyasviel · Apr 17, 2024 · dbf49d3 · dbf49d3
1 parent e641303
commit dbf49d3
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 1 deletion.
diff --git a/javascript/script.js b/javascript/script.js
@@ -122,6 +122,43 @@ document.addEventListener("DOMContentLoaded", function() {
     initStylePreviewOverlay();
 });
 
+var onAppend = function(elem, f) {
+    var observer = new MutationObserver(function(mutations) {
+        mutations.forEach(function(m) {
+            if (m.addedNodes.length) {
+                f(m.addedNodes);
+            }
+        });
+    });
+    observer.observe(elem, {childList: true});
+}
+
+function addObserverIfDesiredNodeAvailable(querySelector, callback) {
+    var elem = document.querySelector(querySelector);
+    if (!elem) {
+        window.setTimeout(() => addObserverIfDesiredNodeAvailable(querySelector, callback), 1000);
+        return;
+    }
+
+    onAppend(elem, callback);
+}
+
+/**
+ * Show reset button on toast "Connection errored out."
+ */
+addObserverIfDesiredNodeAvailable(".toast-wrap", function(added) {
+    added.forEach(function(element) {
+         if (element.innerText.includes("Connection errored out.")) {
+             window.setTimeout(function() {
+                document.getElementById("reset_button").classList.remove("hidden");
+                document.getElementById("generate_button").classList.add("hidden");
+                document.getElementById("skip_button").classList.add("hidden");
+                document.getElementById("stop_button").classList.add("hidden");
+            });
+         }
+    });
+});
+
 /**
  * Add a ctrl+enter as a shortcut to start a generation
  */

diff --git a/language/en.json b/language/en.json
@@ -4,6 +4,7 @@
     "Generate": "Generate",
     "Skip": "Skip",
     "Stop": "Stop",
+    "Reconnect and Reset UI": "Reconnect and Reset UI",
     "Input Image": "Input Image",
     "Advanced": "Advanced",
     "Upscale or Variation": "Upscale or Variation",

diff --git a/webui.py b/webui.py
@@ -123,8 +123,9 @@ def generate_clicked(task: worker.AsyncTask):
 
                 with gr.Column(scale=3, min_width=0):
                     generate_button = gr.Button(label="Generate", value="Generate", elem_classes='type_row', elem_id='generate_button', visible=True)
+                    reset_button = gr.Button(label="Reconnect and Reset UI", value="Reconnect and Reset UI", elem_classes='type_row', elem_id='reset_button', visible=False)
                     load_parameter_button = gr.Button(label="Load Parameters", value="Load Parameters", elem_classes='type_row', elem_id='load_parameter_button', visible=False)
-                    skip_button = gr.Button(label="Skip", value="Skip", elem_classes='type_row_half', visible=False)
+                    skip_button = gr.Button(label="Skip", value="Skip", elem_classes='type_row_half', elem_id='skip_button', visible=False)
                     stop_button = gr.Button(label="Stop", value="Stop", elem_classes='type_row_half', elem_id='stop_button', visible=False)
 
                     def stop_clicked(currentTask):
@@ -688,6 +689,14 @@ def trigger_metadata_import(filepath, state_is_generating):
             .then(fn=update_history_link, outputs=history_link) \
             .then(fn=lambda: None, _js='playNotification').then(fn=lambda: None, _js='refresh_grid_delayed')
 
+        reset_button.click(lambda: [worker.AsyncTask(args=[]), False, gr.update(visible=True, interactive=True)] +
+                                   [gr.update(visible=False)] * 6 +
+                                   [gr.update(visible=True, value=[])],
+                           outputs=[currentTask, state_is_generating, generate_button,
+                                    reset_button, stop_button, skip_button,
+                                    progress_html, progress_window, progress_gallery, gallery],
+                           queue=False)
+
         for notification_file in ['notification.ogg', 'notification.mp3']:
             if os.path.exists(notification_file):
                 gr.Audio(interactive=False, value=notification_file, elem_id='audio_notification', visible=False)