diff --git a/invokeai/app/api/routers/session_queue.py b/invokeai/app/api/routers/session_queue.py index 660d09728c..f19f4bc47c 100644 --- a/invokeai/app/api/routers/session_queue.py +++ b/invokeai/app/api/routers/session_queue.py @@ -14,6 +14,7 @@ from invokeai.app.services.session_queue.session_queue_common import ( CancelByBatchIDsResult, CancelByDestinationResult, ClearResult, + DeleteByDestinationResult, EnqueueBatchResult, FieldIdentifier, PruneResult, @@ -293,6 +294,18 @@ async def get_queue_item( return ApiDependencies.invoker.services.session_queue.get_queue_item(item_id) +@session_queue_router.delete( + "/{queue_id}/i/{item_id}", + operation_id="delete_queue_item", +) +async def delete_queue_item( + queue_id: str = Path(description="The queue id to perform this operation on"), + item_id: int = Path(description="The queue item to delete"), +) -> None: + """Deletes a queue item""" + ApiDependencies.invoker.services.session_queue.delete_queue_item(item_id) + + @session_queue_router.put( "/{queue_id}/i/{item_id}/cancel", operation_id="cancel_queue_item", @@ -322,3 +335,18 @@ async def counts_by_destination( return ApiDependencies.invoker.services.session_queue.get_counts_by_destination( queue_id=queue_id, destination=destination ) + + +@session_queue_router.delete( + "/{queue_id}/d/{destination}", + operation_id="delete_by_destination", + responses={200: {"model": DeleteByDestinationResult}}, +) +async def delete_by_destination( + queue_id: str = Path(description="The queue id to query"), + destination: str = Path(description="The destination to query"), +) -> DeleteByDestinationResult: + """Deletes all items with the given destination""" + return ApiDependencies.invoker.services.session_queue.delete_by_destination( + queue_id=queue_id, destination=destination + ) diff --git a/invokeai/app/services/session_queue/session_queue_base.py b/invokeai/app/services/session_queue/session_queue_base.py index a2aa844742..17c31e77a7 100644 --- a/invokeai/app/services/session_queue/session_queue_base.py +++ b/invokeai/app/services/session_queue/session_queue_base.py @@ -10,6 +10,7 @@ from invokeai.app.services.session_queue.session_queue_common import ( CancelByDestinationResult, CancelByQueueIDResult, ClearResult, + DeleteByDestinationResult, EnqueueBatchResult, IsEmptyResult, IsFullResult, @@ -91,6 +92,11 @@ class SessionQueueBase(ABC): """Cancels a session queue item""" pass + @abstractmethod + def delete_queue_item(self, item_id: int) -> None: + """Deletes a session queue item""" + pass + @abstractmethod def fail_queue_item( self, item_id: int, error_type: str, error_message: str, error_traceback: str @@ -108,6 +114,11 @@ class SessionQueueBase(ABC): """Cancels all queue items with the given batch destination""" pass + @abstractmethod + def delete_by_destination(self, queue_id: str, destination: str) -> DeleteByDestinationResult: + """Deletes all queue items with the given batch destination""" + pass + @abstractmethod def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult: """Cancels all queue items with matching queue ID""" diff --git a/invokeai/app/services/session_queue/session_queue_common.py b/invokeai/app/services/session_queue/session_queue_common.py index 1861110c97..d41fb44533 100644 --- a/invokeai/app/services/session_queue/session_queue_common.py +++ b/invokeai/app/services/session_queue/session_queue_common.py @@ -363,6 +363,12 @@ class CancelByDestinationResult(CancelByBatchIDsResult): pass +class DeleteByDestinationResult(BaseModel): + """Result of deleting by a destination""" + + deleted: int = Field(..., description="Number of queue items deleted") + + class CancelByQueueIDResult(CancelByBatchIDsResult): """Result of canceling by queue id""" diff --git a/invokeai/app/services/session_queue/session_queue_sqlite.py b/invokeai/app/services/session_queue/session_queue_sqlite.py index 0a0ef72456..c31226581a 100644 --- a/invokeai/app/services/session_queue/session_queue_sqlite.py +++ b/invokeai/app/services/session_queue/session_queue_sqlite.py @@ -17,6 +17,7 @@ from invokeai.app.services.session_queue.session_queue_common import ( CancelByDestinationResult, CancelByQueueIDResult, ClearResult, + DeleteByDestinationResult, EnqueueBatchResult, IsEmptyResult, IsFullResult, @@ -212,6 +213,19 @@ class SqliteSessionQueue(SessionQueueBase): ) -> SessionQueueItem: try: cursor = self._conn.cursor() + cursor.execute( + """--sql + SELECT status FROM session_queue WHERE item_id = ? + """, + (item_id,), + ) + row = cursor.fetchone() + if row is None: + raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}") + current_status = row[0] + # Only update if not already finished (completed, failed or canceled) + if current_status in ("completed", "failed", "canceled"): + return self.get_queue_item(item_id) cursor.execute( """--sql UPDATE session_queue @@ -323,6 +337,27 @@ class SqliteSessionQueue(SessionQueueBase): queue_item = self._set_queue_item_status(item_id=item_id, status="canceled") return queue_item + def delete_queue_item(self, item_id: int) -> None: + """Deletes a session queue item""" + try: + self.cancel_queue_item(item_id) + except SessionQueueItemNotFoundError: + pass + try: + cursor = self._conn.cursor() + cursor.execute( + """--sql + DELETE + FROM session_queue + WHERE item_id = ? + """, + (item_id,), + ) + self._conn.commit() + except Exception: + self._conn.rollback() + raise + def complete_queue_item(self, item_id: int) -> SessionQueueItem: queue_item = self._set_queue_item_status(item_id=item_id, status="completed") return queue_item @@ -420,6 +455,40 @@ class SqliteSessionQueue(SessionQueueBase): raise return CancelByDestinationResult(canceled=count) + def delete_by_destination(self, queue_id: str, destination: str) -> DeleteByDestinationResult: + try: + cursor = self._conn.cursor() + current_queue_item = self.get_current(queue_id) + if current_queue_item is not None and current_queue_item.destination == destination: + self.cancel_queue_item(current_queue_item.item_id) + params = (queue_id, destination) + cursor.execute( + """--sql + SELECT COUNT(*) + FROM session_queue + WHERE + queue_id = ? + AND destination = ?; + """, + params, + ) + count = cursor.fetchone()[0] + cursor.execute( + """--sql + DELETE + FROM session_queue + WHERE + queue_id = ? + AND destination = ?; + """, + params, + ) + self._conn.commit() + except Exception: + self._conn.rollback() + raise + return DeleteByDestinationResult(deleted=count) + def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult: try: cursor = self._conn.cursor()