Skip to content

Commit 47dc954

Browse files
rikublockpsychedelicious
authored andcommitted
feat(app): add cancel all except current queue item functionality
1 parent 8fc5d3d commit 47dc954

File tree

4 files changed

+59
-0
lines changed

4 files changed

+59
-0
lines changed

invokeai/app/api/routers/session_queue.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
QUEUE_ITEM_STATUS,
1111
Batch,
1212
BatchStatus,
13+
CancelAllExceptCurrentResult,
1314
CancelByBatchIDsResult,
1415
CancelByDestinationResult,
1516
ClearResult,
@@ -94,6 +95,18 @@ async def Pause(
9495
return ApiDependencies.invoker.services.session_processor.pause()
9596

9697

98+
@session_queue_router.put(
99+
"/{queue_id}/cancel_all_except_current",
100+
operation_id="cancel_all_except_current",
101+
responses={200: {"model": CancelAllExceptCurrentResult}},
102+
)
103+
async def cancel_all_except_current(
104+
queue_id: str = Path(description="The queue id to perform this operation on"),
105+
) -> CancelAllExceptCurrentResult:
106+
"""Immediately cancels all queue items except in-processing items"""
107+
return ApiDependencies.invoker.services.session_queue.cancel_all_except_current(queue_id=queue_id)
108+
109+
97110
@session_queue_router.put(
98111
"/{queue_id}/cancel_by_batch_ids",
99112
operation_id="cancel_by_batch_ids",

invokeai/app/services/session_queue/session_queue_base.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
QUEUE_ITEM_STATUS,
66
Batch,
77
BatchStatus,
8+
CancelAllExceptCurrentResult,
89
CancelByBatchIDsResult,
910
CancelByDestinationResult,
1011
CancelByQueueIDResult,
@@ -112,6 +113,11 @@ def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult:
112113
"""Cancels all queue items with matching queue ID"""
113114
pass
114115

116+
@abstractmethod
117+
def cancel_all_except_current(self, queue_id: str) -> CancelAllExceptCurrentResult:
118+
"""Cancels all queue items except in-progress items"""
119+
pass
120+
115121
@abstractmethod
116122
def list_queue_items(
117123
self,

invokeai/app/services/session_queue/session_queue_common.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -374,6 +374,12 @@ class CancelByQueueIDResult(CancelByBatchIDsResult):
374374
pass
375375

376376

377+
class CancelAllExceptCurrentResult(CancelByBatchIDsResult):
378+
"""Result of canceling all except current"""
379+
380+
pass
381+
382+
377383
class IsEmptyResult(BaseModel):
378384
"""Result of checking if the session queue is empty"""
379385

invokeai/app/services/session_queue/session_queue_sqlite.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
QUEUE_ITEM_STATUS,
1010
Batch,
1111
BatchStatus,
12+
CancelAllExceptCurrentResult,
1213
CancelByBatchIDsResult,
1314
CancelByDestinationResult,
1415
CancelByQueueIDResult,
@@ -510,6 +511,39 @@ def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult:
510511
self.__lock.release()
511512
return CancelByQueueIDResult(canceled=count)
512513

514+
def cancel_all_except_current(self, queue_id: str) -> CancelAllExceptCurrentResult:
515+
try:
516+
where = """--sql
517+
WHERE
518+
queue_id == ?
519+
AND status == 'pending'
520+
"""
521+
self.__lock.acquire()
522+
self.__cursor.execute(
523+
f"""--sql
524+
SELECT COUNT(*)
525+
FROM session_queue
526+
{where};
527+
""",
528+
(queue_id,),
529+
)
530+
count = self.__cursor.fetchone()[0]
531+
self.__cursor.execute(
532+
f"""--sql
533+
UPDATE session_queue
534+
SET status = 'canceled'
535+
{where};
536+
""",
537+
(queue_id,),
538+
)
539+
self.__conn.commit()
540+
except Exception:
541+
self.__conn.rollback()
542+
raise
543+
finally:
544+
self.__lock.release()
545+
return CancelAllExceptCurrentResult(canceled=count)
546+
513547
def get_queue_item(self, item_id: int) -> SessionQueueItem:
514548
try:
515549
self.__lock.acquire()

0 commit comments

Comments
 (0)