Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(events): add batch status and queue status to queue item status changed events

Browse Source 
The UI will always re-fetch queue and batch status on receiving this event, so we may as well jsut include that data in the event and save the extra network roundtrips.
This commit is contained in:
psychedelicious 2023-10-09 11:04:03 +11:00
parent 90083cc88d
commit 55b40a9425
2 changed files with 48 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
invokeai/app/services/events.py
View File

@ -4,7 +4,12 @@ from typing import Any, Optional
from invokeai.app.models.image import ProgressImage from invokeai.app.models.image import ProgressImage
from invokeai.app.services.model_manager_service import BaseModelType, ModelInfo, ModelType, SubModelType from invokeai.app.services.model_manager_service import BaseModelType, ModelInfo, ModelType, SubModelType
from invokeai.app.services.session_queue.session_queue_common import EnqueueBatchResult, SessionQueueItem from invokeai.app.services.session_queue.session_queue_common import (
BatchStatus,
EnqueueBatchResult,
SessionQueueItem,
SessionQueueStatus,
)
from invokeai.app.util.misc import get_timestamp from invokeai.app.util.misc import get_timestamp
@ -262,13 +267,20 @@ class EventServiceBase:
), ),
) )
def emit_queue_item_status_changed(self, session_queue_item: SessionQueueItem) -> None: def emit_queue_item_status_changed(
self,
session_queue_item: SessionQueueItem,
batch_status: BatchStatus,
queue_status: SessionQueueStatus,
) -> None:
"""Emitted when a queue item's status changes""" """Emitted when a queue item's status changes"""
self.__emit_queue_event( self.__emit_queue_event(
event_name="queue_item_status_changed", event_name="queue_item_status_changed",
payload=dict( payload=dict(
queue_id=queue_status.queue_id,
queue_item=dict(
queue_id=session_queue_item.queue_id, queue_id=session_queue_item.queue_id,
queue_item_id=session_queue_item.item_id, item_id=session_queue_item.item_id,
status=session_queue_item.status, status=session_queue_item.status,
batch_id=session_queue_item.batch_id, batch_id=session_queue_item.batch_id,
session_id=session_queue_item.session_id, session_id=session_queue_item.session_id,
@ -278,6 +290,9 @@ class EventServiceBase:
started_at=str(session_queue_item.started_at) if session_queue_item.started_at else None, started_at=str(session_queue_item.started_at) if session_queue_item.started_at else None,
completed_at=str(session_queue_item.completed_at) if session_queue_item.completed_at else None, completed_at=str(session_queue_item.completed_at) if session_queue_item.completed_at else None,
), ),
batch_status=batch_status.dict(),
queue_status=queue_status.dict(),
),
) )
def emit_batch_enqueued(self, enqueue_result: EnqueueBatchResult) -> None: def emit_batch_enqueued(self, enqueue_result: EnqueueBatchResult) -> None:

24
invokeai/app/services/session_queue/session_queue_sqlite.py
View File

@ -427,7 +427,13 @@ class SqliteSessionQueue(SessionQueueBase):
finally: finally:
self.__lock.release() self.__lock.release()
queue_item = self.get_queue_item(item_id) queue_item = self.get_queue_item(item_id)
self.__invoker.services.events.emit_queue_item_status_changed(queue_item) batch_status = self.get_batch_status(queue_id=queue_item.queue_id, batch_id=queue_item.batch_id)
queue_status = self.get_queue_status(queue_id=queue_item.queue_id)
self.__invoker.services.events.emit_queue_item_status_changed(
session_queue_item=queue_item,
batch_status=batch_status,
queue_status=queue_status,
)
return queue_item return queue_item
def is_empty(self, queue_id: str) -> IsEmptyResult: def is_empty(self, queue_id: str) -> IsEmptyResult:
@ -609,7 +615,13 @@ class SqliteSessionQueue(SessionQueueBase):
queue_batch_id=current_queue_item.batch_id, queue_batch_id=current_queue_item.batch_id,
graph_execution_state_id=current_queue_item.session_id, graph_execution_state_id=current_queue_item.session_id,
) )
self.__invoker.services.events.emit_queue_item_status_changed(current_queue_item) batch_status = self.get_batch_status(queue_id=queue_id, batch_id=current_queue_item.batch_id)
queue_status = self.get_queue_status(queue_id=queue_id)
self.__invoker.services.events.emit_queue_item_status_changed(
session_queue_item=current_queue_item,
batch_status=batch_status,
queue_status=queue_status,
)
except Exception: except Exception:
self.__conn.rollback() self.__conn.rollback()
raise raise
@ -655,7 +667,13 @@ class SqliteSessionQueue(SessionQueueBase):
queue_batch_id=current_queue_item.batch_id, queue_batch_id=current_queue_item.batch_id,
graph_execution_state_id=current_queue_item.session_id, graph_execution_state_id=current_queue_item.session_id,
) )
self.__invoker.services.events.emit_queue_item_status_changed(current_queue_item) batch_status = self.get_batch_status(queue_id=queue_id, batch_id=current_queue_item.batch_id)
queue_status = self.get_queue_status(queue_id=queue_id)
self.__invoker.services.events.emit_queue_item_status_changed(
session_queue_item=current_queue_item,
batch_status=batch_status,
queue_status=queue_status,
)
except Exception: except Exception:
self.__conn.rollback() self.__conn.rollback()
raise raise