Skip to content

Commit

Permalink
Switch to remote workflow after "Open Web UI" button is clicked
Browse files Browse the repository at this point in the history
  • Loading branch information
Acly committed Nov 12, 2024
1 parent 617c2ea commit 465f550
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 2 deletions.
25 changes: 24 additions & 1 deletion ai_diffusion/custom_workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion ai_diffusion/resources.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
version = "1.28.0"

comfy_url = "https://github.com/comfyanonymous/ComfyUI"
comfy_version = "5e29e7a488b3f48afc6c4a3cb8ed110976d0ebb8"
comfy_version = "2d28b0b4790e3f6c2287be49d9872419eadfe5bb"


class CustomNode(NamedTuple):
Expand Down
1 change: 1 addition & 0 deletions ai_diffusion/ui/custom_workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 465f550

Please sign in to comment.