Fix issue 1137: Query ComfyUI queue state and use it to display count…

… of jobs in the queue before ours.

ai_diffusion/client.py

@@ -25,6 +25,7 @@ class ClientEvent(Enum):
     disconnected = 5
     queued = 6
     upload = 7
+    foreign_jobs = 8
 
 
 class ClientMessage(NamedTuple):
@@ -34,6 +35,7 @@ class ClientMessage(NamedTuple):
     images: ImageCollection | None = None
     result: dict | None = None
     error: str | None = None
+    foreign_jobs: int | None = None
 
 
 class User(QObject, ObservableProperties):

ai_diffusion/comfy_client.py

@@ -7,6 +7,7 @@
 from enum import Enum
 from collections import deque
 from itertools import chain, product
+from operator import itemgetter
 from typing import NamedTuple, Optional, Sequence
 
 from .api import WorkflowInput
@@ -264,6 +265,7 @@ async def _listen_websocket(self, websocket: websockets_client.WebSocketClientPr
 
                 if msg["type"] == "status":
                     await self._report(ClientEvent.connected, "")
+                    await self._poll_server_queue()
 
                 if msg["type"] == "execution_start":
                     id = msg["data"]["prompt_id"]
@@ -345,6 +347,29 @@ async def clear_queue(self):
 
         self._jobs.clear()
 
+    async def _poll_server_queue(self):
+        queue = await self._get("api/queue")
+        server_jobs = queue["queue_running"] + queue["queue_pending"]
+        # why are they unsorted to start with...?
+        server_jobs = sorted(server_jobs, key=itemgetter(0))
+        server_jobs = [entry[1] for entry in server_jobs]
+        if not (self._jobs or self._active):
+            await self._report(ClientEvent.foreign_jobs, "", foreign_jobs=len(server_jobs))
+            return
+
+        if self._active:
+            first_remote_id = await self._active.get_remote_id()
+        else:
+            first_remote_id = await self._jobs[0].get_remote_id()
+        # if we got it from _jobs or _active, this field must have been set (in _run_job).
+        first_remote_id = util.ensure(first_remote_id)
+        try:
+            offset = server_jobs.index(first_remote_id)
+        except ValueError:
+            # probably just haven't gotten the notification yet
+            return
+        await self._report(ClientEvent.foreign_jobs, "", foreign_jobs=offset)
+
     async def disconnect(self):
         if self._is_connected:
             self._is_connected = False
@@ -445,7 +470,7 @@ def _get_active_job(self, remote_id: str) -> Optional[JobInfo]:
             return self._active
         elif self._active:
             log.warning(f"Received message for job {remote_id}, but job {self._active} is active")
-        if len(self._jobs) == 0:
+        if not self._jobs:
             log.warning(f"Received unknown job {remote_id}")
             return None
         active = next((j for j in self._jobs if j.remote_id == remote_id), None)
@@ -456,7 +481,7 @@ def _get_active_job(self, remote_id: str) -> Optional[JobInfo]:
     async def _start_job(self, remote_id: str):
         if self._active is not None:
             log.warning(f"Started job {remote_id}, but {self._active} was never finished")
-        if len(self._jobs) == 0:
+        if not self._jobs:
             log.warning(f"Received unknown job {remote_id}")
             return None
 

ai_diffusion/connection.py

@@ -31,11 +31,13 @@ class Connection(QObject, ObservableProperties):
     state = Property(ConnectionState.disconnected)
     error = Property("")
     missing_resource: MissingResource | None = None
+    foreign_jobs: int = 0
 
     state_changed = pyqtSignal(ConnectionState)
     error_changed = pyqtSignal(str)
     models_changed = pyqtSignal()
     message_received = pyqtSignal(ClientMessage)
+    foreign_jobs_changed = pyqtSignal(int)
 
     _client: Client | None = None
     _task: asyncio.Task | None = None
@@ -169,6 +171,9 @@ async def _handle_messages(self):
                             if temporary_disconnect:
                                 temporary_disconnect = False
                                 self.error = ""
+                        elif msg.event is ClientEvent.foreign_jobs:
+                            self.foreign_jobs = util.ensure(msg.foreign_jobs)
+                            self.foreign_jobs_changed.emit(self.foreign_jobs)
                         else:
                             self.message_received.emit(msg)
                     except asyncio.CancelledError:

ai_diffusion/ui/widget.py

@@ -209,6 +209,7 @@ def _connect_model(self):
         self._connections = [
             self._model.jobs.count_changed.connect(self._update),
             self._model.progress_kind_changed.connect(self._update),
+            root.connection.foreign_jobs_changed.connect(self._update),
         ]
 
     def _update(self):
@@ -220,6 +221,10 @@ def _update(self):
             self.setIcon(theme.icon("queue-upload"))
             self.setToolTip(_("Uploading models.") + f" {queued_msg} {cancel_msg}")
             count += 1
+        elif root.connection.foreign_jobs > 0:
+            self.setIcon(theme.icon("queue-inactive"))
+            self.setToolTip(_("Server is busy."))
+            count = f"+{root.connection.foreign_jobs}"
         elif self._model.jobs.any_executing():
             self.setIcon(theme.icon("queue-active"))
             if count > 0: