From 465f5509a0de92b02f09462711c79d29bdae190a Mon Sep 17 00:00:00 2001
From: Acly <aclysia@gmail.com>
Date: Tue, 12 Nov 2024 12:18:31 +0100
Subject: [PATCH] Switch to remote workflow after "Open Web UI" button is
 clicked

---
 ai_diffusion/custom_workflow.py    | 25 ++++++++++++++++++++++++-
 ai_diffusion/resources.py          |  2 +-
 ai_diffusion/ui/custom_workflow.py |  1 +
 3 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/ai_diffusion/custom_workflow.py b/ai_diffusion/custom_workflow.py
index ed681b662..81487c1b5 100644
--- a/ai_diffusion/custom_workflow.py
+++ b/ai_diffusion/custom_workflow.py
@@ -8,7 +8,7 @@
 from typing import Any, Awaitable, Callable, NamedTuple, Literal, TYPE_CHECKING
 from pathlib import Path
 from PyQt5.QtCore import Qt, QObject, QUuid, QAbstractListModel, QSortFilterProxyModel, QModelIndex
-from PyQt5.QtCore import pyqtSignal
+from PyQt5.QtCore import QMetaObject, QTimer, pyqtSignal
 
 from .api import WorkflowInput
 from .client import TextOutput, ClientOutput
@@ -385,6 +385,8 @@ def __init__(self, workflows: WorkflowCollection, generator: ImageGenerator, job
         self._last_result: Image | None = None
         self._last_job: JobParams | None = None
         self._new_outputs: list[str] = []
+        self._switch_workflow_bind: QMetaObject.Connection | None = None
+        self._switch_workflow_timer: QTimer | None = None
 
         jobs.job_finished.connect(self._handle_job_finished)
         workflows.dataChanged.connect(self._update_workflow)
@@ -503,6 +505,27 @@ def collect_parameters(self, layers: "LayerManager", bounds: Bounds):
 
         return params
 
+    def switch_to_web_workflow(self):
+        self._switch_workflow_bind = self._workflows.rowsInserted.connect(self._set_workflow_index)
+        self._switch_workflow_timer = QTimer()
+        self._switch_workflow_timer.timeout.connect(self._clear_switch_workflow)
+        self._switch_workflow_timer.start(5 * 60 * 1000)
+
+    def _set_workflow_index(self, parent: QModelIndex, first: int, last: int):
+        for i in range(first, last + 1):
+            if self._workflows[i].source is WorkflowSource.remote:
+                self._clear_switch_workflow()
+                self.workflow_id = self._workflows[i].id
+                return
+
+    def _clear_switch_workflow(self):
+        if self._switch_workflow_timer is not None:
+            self._switch_workflow_timer.stop()
+            self._switch_workflow_timer = None
+        if self._switch_workflow_bind is not None:
+            self._workflows.rowsInserted.disconnect(self._switch_workflow_bind)
+            self._switch_workflow_bind = None
+
     def show_output(self, output: ClientOutput | None):
         if isinstance(output, TextOutput):
             self._new_outputs.append(output.key)
diff --git a/ai_diffusion/resources.py b/ai_diffusion/resources.py
index d86253053..7d379ee2c 100644
--- a/ai_diffusion/resources.py
+++ b/ai_diffusion/resources.py
@@ -9,7 +9,7 @@
 version = "1.28.0"
 
 comfy_url = "https://github.com/comfyanonymous/ComfyUI"
-comfy_version = "5e29e7a488b3f48afc6c4a3cb8ed110976d0ebb8"
+comfy_version = "2d28b0b4790e3f6c2287be49d9872419eadfe5bb"
 
 
 class CustomNode(NamedTuple):
diff --git a/ai_diffusion/ui/custom_workflow.py b/ai_diffusion/ui/custom_workflow.py
index 100d80cad..278f28dfc 100644
--- a/ai_diffusion/ui/custom_workflow.py
+++ b/ai_diffusion/ui/custom_workflow.py
@@ -807,6 +807,7 @@ def _delete_workflow(self):
     def _open_webui(self):
         if client := root.connection.client_if_connected:
             QDesktopServices.openUrl(QUrl(client.url))
+            self.model.custom.switch_to_web_workflow()
 
     @property
     def is_edit_mode(self):