From 8dc1ef33c83e474a7db297b829830d89aecb84f8 Mon Sep 17 00:00:00 2001
From: ljleb <lebel.louisjacob@gmail.com>
Date: Sun, 29 Oct 2023 16:11:49 -0400
Subject: [PATCH] commit

---
 javascript/compute-height.js           | 14 +++++++++++++-
 lib_comfyui/comfyui/iframe_requests.py | 10 +++-------
 lib_comfyui/webui/tab.py               |  7 +++++++
 scripts/comfyui.py                     |  8 ++++----
 4 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/javascript/compute-height.js b/javascript/compute-height.js
index 391d268..8025718 100644
--- a/javascript/compute-height.js
+++ b/javascript/compute-height.js
@@ -22,7 +22,8 @@ function onComfyuiTabLoaded(callback) {
     if (getClearEnabledDisplayNamesButtons().some(e => e === null) ||
         getWorkflowTypeIds() === null ||
         getComfyuiContainer() === null ||
-        getTabNav() === null
+        getTabNav() === null ||
+        getWebuiClientIdTextArea() === null
     ) {
         // webui not yet ready, try again in a bit
         setTimeout(() => { onComfyuiTabLoaded(callback); }, POLLING_TIMEOUT);
@@ -39,6 +40,7 @@ function clearEnabledDisplayNames() {
 }
 
 function setupComfyuiTabEvents() {
+    setupWebuiClientId();
     setupResizeTabEvent();
     setupToggleFooterEvent();
 
@@ -54,6 +56,12 @@ function reloadComfyuiIFrames() {
     });
 }
 
+function setupWebuiClientId() {
+    const textArea = getWebuiClientIdTextArea();
+    textArea.value = WEBUI_CLIENT_ID;
+    textArea.dispatchEvent(new Event('input'));
+}
+
 function setupResizeTabEvent() {
     window.addEventListener("resize", updateComfyuiTabHeight);
 }
@@ -108,6 +116,10 @@ function getComfyuiContainer() {
     return document.getElementById("comfyui_webui_container") ?? null;
 }
 
+function getWebuiClientIdTextArea() {
+    return document.querySelector("#comfyui_webui_client_id textarea") ?? null;
+}
+
 function getFooter() {
     return document.querySelector('#footer') ?? null;
 }
diff --git a/lib_comfyui/comfyui/iframe_requests.py b/lib_comfyui/comfyui/iframe_requests.py
index e3eed61..eea7c85 100644
--- a/lib_comfyui/comfyui/iframe_requests.py
+++ b/lib_comfyui/comfyui/iframe_requests.py
@@ -8,7 +8,6 @@
 
 class ComfyuiIFrameRequests:
     finished_comfyui_queue = multiprocessing.Queue()
-    focused_webui_client_id = None
     server_instance = None
     sid_map = {}
 
@@ -19,12 +18,12 @@ def send(request, workflow_type, data=None):
             data = {}
 
         cls = ComfyuiIFrameRequests
-        if cls.focused_webui_client_id is None:
+        if global_state.webui_client_id is None:
             raise RuntimeError('No active webui connection')
 
-        ws_client_ids = cls.sid_map[cls.focused_webui_client_id]
+        ws_client_ids = cls.sid_map[global_state.webui_client_id]
         if workflow_type not in ws_client_ids:
-            raise RuntimeError(f"The workflow type {workflow_type} has not been registered by the active webui client {cls.focused_webui_client_id}")
+            raise RuntimeError(f"The workflow type {workflow_type} has not been registered by the active webui client {global_state.webui_client_id}")
 
         clear_queue(cls.finished_comfyui_queue)
         cls.server_instance.send_sync(request, data, ws_client_ids[workflow_type])
@@ -76,9 +75,6 @@ def register_client(request):
         webui_client_id = request['webuiClientId']
         sid = request['sid']
 
-        # TODO: generalize this
-        ComfyuiIFrameRequests.focused_webui_client_id = webui_client_id
-
         if webui_client_id not in ComfyuiIFrameRequests.sid_map:
             ComfyuiIFrameRequests.sid_map[webui_client_id] = {}
 
diff --git a/lib_comfyui/webui/tab.py b/lib_comfyui/webui/tab.py
index 54abc47..7e7d9ac 100644
--- a/lib_comfyui/webui/tab.py
+++ b/lib_comfyui/webui/tab.py
@@ -8,6 +8,12 @@
 from lib_comfyui.default_workflow_types import sandbox_tab_workflow_type
 
 
+webui_client_id = gr.Text(
+    elem_id='comfyui_webui_client_id',
+    visible=False,
+)
+
+
 def create_tab():
     install_location = settings.get_install_location()
     with gr.Blocks() as tab:
@@ -33,6 +39,7 @@ def create_tab():
             key='workflow_type_ids',
             value=external_code.get_workflow_type_ids(),
         )
+        webui_client_id.render()
     return [(tab, sandbox_tab_workflow_type.display_name, 'comfyui_webui_root')]
 
 
diff --git a/scripts/comfyui.py b/scripts/comfyui.py
index 44aafa3..d483915 100644
--- a/scripts/comfyui.py
+++ b/scripts/comfyui.py
@@ -2,7 +2,7 @@
 
 from modules import scripts
 from lib_comfyui import global_state, platform_utils, external_code, default_workflow_types, comfyui_process
-from lib_comfyui.webui import callbacks, settings, patches, gradio_utils, accordion
+from lib_comfyui.webui import callbacks, settings, patches, gradio_utils, accordion, tab
 from lib_comfyui.comfyui import iframe_requests, type_conversion
 
 
@@ -39,17 +39,17 @@ def ui(self, is_img2img):
         self.accordion.arrange_components()
         self.accordion.connect_events()
         self.accordion.setup_infotext_fields(self)
-        return self.accordion.get_script_ui_components()
+        return (tab.webui_client_id,) + self.accordion.get_script_ui_components()
 
-    def process(self, p, queue_front, enabled_workflow_type_ids, **kwargs):
+    def process(self, p, webui_client_id, queue_front, enabled_workflow_type_ids, *args, **kwargs):
         if not getattr(global_state, 'enabled', True):
             return
 
         if not hasattr(global_state, 'enabled_workflow_type_ids'):
             global_state.enabled_workflow_type_ids = {}
 
+        global_state.webui_client_id = webui_client_id
         global_state.enabled_workflow_type_ids.update(enabled_workflow_type_ids)
-
         global_state.queue_front = queue_front
         patches.patch_processing(p)