option to enable proxy

lib_comfyui/webui/accordion.py

@@ -66,7 +66,7 @@ def __init__(self, get_elem_id, tab):
         self._rendered = False
 
     def get_iframes_html(self) -> str:
-        comfyui_client_url = settings.get_comfyui_client_url()
+        comfyui_client_url = settings.get_comfyui_iframe_url()
         first_workflow_type_id = self.workflow_type_ids[self.first_workflow_type.display_name]
 
         iframes = []

lib_comfyui/webui/reverse_proxy.py

@@ -2,6 +2,9 @@
 
 
 def register_comfyui(fast_api):
+    if not settings.is_reverse_proxy_enabled():
+        return
+
     from starlette.requests import Request
     from starlette.responses import StreamingResponse
     from starlette.background import BackgroundTask
@@ -13,11 +16,11 @@ def register_comfyui(fast_api):
     from websockets.exceptions import ConnectionClosedOK
 
     # Constants & Configuration
-    comfyui_target_url = f"http://localhost:{settings.get_port()}/"
+    comfyui_target_url = settings.get_comfyui_server_url()
     client = httpx.AsyncClient(base_url=comfyui_target_url)
     ws_client_url = http_to_ws(comfyui_target_url)
-    proxy_path = "/sd-webui-comfyui/comfyui-reverse-proxy"
-    proxy_path_bytes = bytes("/sd-webui-comfyui/comfyui-reverse-proxy", "utf-8")
+    proxy_path = settings.get_comfyui_reverse_proxy_route()
+    proxy_path_bytes = bytes(proxy_path, "utf-8")
 
     async def async_iter_raw_patched(response):
         async for chunk in response.aiter_raw():

lib_comfyui/webui/settings.py

@@ -32,6 +32,9 @@ def create_section():
     shared.opts.onchange('comfyui_graceful_termination_timeout', update_comfyui_graceful_termination_timeout)
     update_comfyui_graceful_termination_timeout()
 
+    shared.opts.add_option("comfyui_reverse_proxy", shared.OptionInfo(
+        False, "Load ComfyUI iframes through a reverse proxy (requires reload UI. Needs --api. Check this if you are on colab or if all comfyui iframes fail to reload manually)", section=section))
+
 
 @ipc.restrict_to_process('webui')
 def update_enabled():
@@ -91,26 +94,65 @@ def get_setting_value(setting_key):
 
 
 @ipc.restrict_to_process('webui')
-def get_port():
+def get_comfyui_iframe_url():
+    if is_reverse_proxy_enabled():
+        return get_comfyui_reverse_proxy_url()
+    else:
+        return get_comfyui_client_url()
+
+
+@ipc.restrict_to_process('webui')
+def get_comfyui_reverse_proxy_url():
+    """
+    comfyui reverse proxy url, as seen from the browser
+    """
     from modules import shared
-    return get_setting_value('--port') or getattr(shared.cmd_opts, 'comfyui_port', 8188)
+    webui_port = shared.cmd_opts.port if shared.cmd_opts.port is not None else 7860
+    return f"http://localhost:{webui_port}{get_comfyui_reverse_proxy_route()}/"
+
+
+def get_comfyui_reverse_proxy_route():
+    return "/sd-webui-comfyui/comfyui"
 
 
 @ipc.restrict_to_process('webui')
 def get_comfyui_client_url():
+    """
+    comfyui server direct url, as seen from the browser
+    """
     from modules import shared
     loopback_address = '127.0.0.1'
-    server_url = get_setting_value('--listen') or getattr(shared.cmd_opts, 'comfyui_listen', loopback_address)
+    server_url = "http://" + (get_setting_value('--listen') or getattr(shared.cmd_opts, 'comfyui_listen', loopback_address)) + ":" + str(get_port()) + "/"
     client_url = shared.opts.data.get('comfyui_client_address', None) or getattr(shared.cmd_opts, 'webui_comfyui_client_address', None) or server_url
-    if client_url == '0.0.0.0':
+    if client_url.startswith(('http://0.0.0.0', 'https://0.0.0.0')):
         print(textwrap.dedent(f"""
-            [ComfyUI extension] changing the ComfyUI client address from {client_url} to {loopback_address}
+            [sd-webui-comfyui] changing the ComfyUI client address from {client_url} to http://{loopback_address}
             This does not change the --listen address passed to ComfyUI, but instead the address used by the extension to load the iframe
             To override this behavior, navigate to the extension settings or use the --webui-comfyui-client-address <address> cli argument
         """), sys.stderr)
-        client_url = loopback_address
+        client_url = client_url.replace("0.0.0.0", "127.0.0.1", 1)
+
+    return client_url
 
-    return f'http://{client_url}:{get_port()}/'
+
+@ipc.restrict_to_process('webui')
+def get_comfyui_server_url():
+    """
+    comfyui server url, as seen from the webui server
+    """
+    return f"http://localhost:{get_port()}/"
+
+
+@ipc.restrict_to_process('webui')
+def get_port():
+    from modules import shared
+    return get_setting_value('--port') or getattr(shared.cmd_opts, 'comfyui_port', 8188)
+
+
+@ipc.restrict_to_process('webui')
+def is_reverse_proxy_enabled():
+    from modules import shared
+    return getattr(shared.cmd_opts, "api", False) and shared.opts.data.get('comfyui_reverse_proxy', False)
 
 
 class WebuiOptions:

lib_comfyui/webui/tab.py

@@ -72,7 +72,7 @@ def get_comfyui_app_html():
     return textwrap.dedent(f'''
         <div id="comfyui_webui_container">
             <iframe
-                base_src="{settings.get_comfyui_client_url()}"
+                base_src="{settings.get_comfyui_iframe_url()}"
                 workflow_type_id="{sandbox_tab_workflow_type.get_ids()[0]}">
             </iframe>
         </div>