fix(ui): fix canvas soft-lock if canceled before first generation

The canvas needs to be set to staging mode as soon as a canvas-destined batch is enqueued. If the batch is is fully canceled before an image is generated, we need to remove that batch from the canvas `batchIds` watchlist, else canvas gets stuck in staging mode with no way to exit. The changes here allow the batch status to be tracked, and if a batch has all its items completed, we can remove it from the `batchIds` watchlist. The `batchIds` watchlist now accurately represents *incomplete* canvas batches, fixing this cause of soft lock.
2024-08-30 20:32:17 +00:00 · 2023-10-09 11:15:19 +11:00 · 2023-10-09 11:15:19 +11:00 · ca95a3bd0d
commit ca95a3bd0d
parent 55b40a9425
8 changed files with 115 additions and 65 deletions
--- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketQueueItemStatusChanged.ts
+++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketQueueItemStatusChanged.ts
@ -11,44 +11,66 @@ export const addSocketQueueItemStatusChangedEventListener = () => {
    actionCreator: socketQueueItemStatusChanged,
    effect: async (action, { dispatch }) => {
      const log = logger('socketio');
-      const {
-        queue_item_id: item_id,
-        queue_batch_id,
-        status,
-      } = action.payload.data;
+
+      const { queue_item, batch_status, queue_status } = action.payload.data;
+
      log.debug(
        action.payload,
-        `Queue item ${item_id} status updated: ${status}`
+        `Queue item ${queue_item.item_id} status updated: ${queue_item.status}`
      );
      dispatch(appSocketQueueItemStatusChanged(action.payload));

      dispatch(
        queueApi.util.updateQueryData('listQueueItems', undefined, (draft) => {
          queueItemsAdapter.updateOne(draft, {
-            id: item_id,
-            changes: action.payload.data,
+            id: queue_item.item_id,
+            changes: queue_item,
          });
        })
      );

+      dispatch(
+        queueApi.util.updateQueryData('getQueueStatus', undefined, (draft) => {
+          Object.assign(draft.queue, queue_status);
+        })
+      );
+      dispatch(
+        queueApi.util.updateQueryData('getQueueStatus', undefined, (draft) => {
+          if (!draft) {
+            return;
+          }
+          Object.assign(draft.queue, queue_status);
+        })
+      );
+
+      dispatch(
+        queueApi.util.updateQueryData(
+          'getBatchStatus',
+          { batch_id: batch_status.batch_id },
+          () => batch_status
+        )
+      );
+
+      dispatch(
+        queueApi.util.updateQueryData(
+          'getQueueItem',
+          queue_item.item_id,
+          (draft) => {
+            if (!draft) {
+              return;
+            }
+            Object.assign(draft, queue_item);
+          }
+        )
+      );
+
      dispatch(
        queueApi.util.invalidateTags([
          'CurrentSessionQueueItem',
          'NextSessionQueueItem',
          'InvocationCacheStatus',
-          { type: 'SessionQueueItem', id: item_id },
-          { type: 'SessionQueueItemDTO', id: item_id },
-          { type: 'BatchStatus', id: queue_batch_id },
        ])
      );
-
-      const req = dispatch(
-        queueApi.endpoints.getQueueStatus.initiate(undefined, {
-          forceRefetch: true,
-        })
-      );
-      await req.unwrap();
-      req.unsubscribe();
    },
  });
 };
--- a/invokeai/frontend/web/src/features/canvas/store/canvasSlice.ts
+++ b/invokeai/frontend/web/src/features/canvas/store/canvasSlice.ts
@ -29,6 +29,7 @@ import {
  isCanvasBaseImage,
  isCanvasMaskLine,
 } from './canvasTypes';
+import { appSocketQueueItemStatusChanged } from 'services/events/actions';

 export const initialLayerState: CanvasLayerState = {
  objects: [],
@ -786,6 +787,18 @@ export const canvasSlice = createSlice({
    },
  },
  extraReducers: (builder) => {
+    builder.addCase(appSocketQueueItemStatusChanged, (state, action) => {
+      const batch_status = action.payload.data.batch_status;
+      if (!state.batchIds.includes(batch_status.batch_id)) {
+        return;
+      }
+
+      if (batch_status.in_progress === 0 && batch_status.pending === 0) {
+        state.batchIds = state.batchIds.filter(
+          (id) => id !== batch_status.batch_id
+        );
+      }
+    });
    builder.addCase(setAspectRatio, (state, action) => {
      const ratio = action.payload;
      if (ratio) {
--- a/invokeai/frontend/web/src/features/nodes/store/nodesSlice.ts
+++ b/invokeai/frontend/web/src/features/nodes/store/nodesSlice.ts
@ -984,7 +984,7 @@ const nodesSlice = createSlice({
      }
    });
    builder.addCase(appSocketQueueItemStatusChanged, (state, action) => {
-      if (['in_progress'].includes(action.payload.data.status)) {
+      if (['in_progress'].includes(action.payload.data.queue_item.status)) {
        forEach(state.nodeExecutionStates, (nes) => {
          nes.status = NodeStatus.PENDING;
          nes.error = null;
--- a/invokeai/frontend/web/src/features/system/store/systemSlice.ts
+++ b/invokeai/frontend/web/src/features/system/store/systemSlice.ts
@ -164,7 +164,9 @@ export const systemSlice = createSlice({

    builder.addCase(appSocketQueueItemStatusChanged, (state, action) => {
      if (
-        ['completed', 'canceled', 'failed'].includes(action.payload.data.status)
+        ['completed', 'canceled', 'failed'].includes(
+          action.payload.data.queue_item.status
+        )
      ) {
        state.status = 'CONNECTED';
        state.denoiseProgress = null;
--- a/invokeai/frontend/web/src/services/api/endpoints/queue.ts
+++ b/invokeai/frontend/web/src/services/api/endpoints/queue.ts
@ -135,12 +135,7 @@ export const queueApi = api.injectEndpoints({
        url: `queue/${$queueId.get()}/prune`,
        method: 'PUT',
      }),
-      invalidatesTags: [
-        'SessionQueueStatus',
-        'BatchStatus',
-        'SessionQueueItem',
-        'SessionQueueItemDTO',
-      ],
+      invalidatesTags: ['SessionQueueStatus', 'BatchStatus'],
      onQueryStarted: async (arg, api) => {
        const { dispatch, queryFulfilled } = api;
        try {
@ -165,8 +160,6 @@ export const queueApi = api.injectEndpoints({
        'BatchStatus',
        'CurrentSessionQueueItem',
        'NextSessionQueueItem',
-        'SessionQueueItem',
-        'SessionQueueItemDTO',
      ],
      onQueryStarted: async (arg, api) => {
        const { dispatch, queryFulfilled } = api;
@ -218,7 +211,6 @@ export const queueApi = api.injectEndpoints({
        url: `queue/${$queueId.get()}/status`,
        method: 'GET',
      }),
-
      providesTags: ['SessionQueueStatus'],
    }),
    getBatchStatus: build.query<
@ -269,7 +261,11 @@ export const queueApi = api.injectEndpoints({
              (draft) => {
                queueItemsAdapter.updateOne(draft, {
                  id: item_id,
-                  changes: { status: data.status },
+                  changes: {
+                    status: data.status,
+                    completed_at: data.completed_at,
+                    updated_at: data.updated_at,
+                  },
                });
              }
            )
@ -284,7 +280,6 @@ export const queueApi = api.injectEndpoints({
        }
        return [
          { type: 'SessionQueueItem', id: result.item_id },
-          { type: 'SessionQueueItemDTO', id: result.item_id },
          { type: 'BatchStatus', id: result.batch_id },
        ];
      },
@ -307,11 +302,7 @@ export const queueApi = api.injectEndpoints({
          // no-op
        }
      },
-      invalidatesTags: [
-        'SessionQueueItem',
-        'SessionQueueItemDTO',
-        'BatchStatus',
-      ],
+      invalidatesTags: ['SessionQueueStatus', 'BatchStatus'],
    }),
    listQueueItems: build.query<
      EntityState<components['schemas']['SessionQueueItemDTO']> & {
--- a/invokeai/frontend/web/src/services/api/index.ts
+++ b/invokeai/frontend/web/src/services/api/index.ts
@ -21,8 +21,6 @@ export const tagTypes = [
  'ImageMetadataFromFile',
  'IntermediatesCount',
  'SessionQueueItem',
-  'SessionQueueItemDTO',
-  'SessionQueueItemDTOList',
  'SessionQueueStatus',
  'SessionProcessorStatus',
  'CurrentSessionQueueItem',
--- a/invokeai/frontend/web/src/services/api/schema.d.ts
+++ b/invokeai/frontend/web/src/services/api/schema.d.ts
@ -9701,11 +9701,23 @@ export type components = {
      ui_order?: number;
    };
    /**
-     * StableDiffusionOnnxModelFormat
+     * T2IAdapterModelFormat
     * @description An enumeration.
     * @enum {string}
     */
-    StableDiffusionOnnxModelFormat: "olive" | "onnx";
+    T2IAdapterModelFormat: "diffusers";
+    /**
+     * ControlNetModelFormat
+     * @description An enumeration.
+     * @enum {string}
+     */
+    ControlNetModelFormat: "checkpoint" | "diffusers";
+    /**
+     * StableDiffusion2ModelFormat
+     * @description An enumeration.
+     * @enum {string}
+     */
+    StableDiffusion2ModelFormat: "checkpoint" | "diffusers";
    /**
     * StableDiffusionXLModelFormat
     * @description An enumeration.
@ -9713,11 +9725,11 @@ export type components = {
     */
    StableDiffusionXLModelFormat: "checkpoint" | "diffusers";
    /**
-     * StableDiffusion2ModelFormat
+     * StableDiffusionOnnxModelFormat
     * @description An enumeration.
     * @enum {string}
     */
-    StableDiffusion2ModelFormat: "checkpoint" | "diffusers";
+    StableDiffusionOnnxModelFormat: "olive" | "onnx";
    /**
     * CLIPVisionModelFormat
     * @description An enumeration.
@ -9736,18 +9748,6 @@ export type components = {
     * @enum {string}
     */
    StableDiffusion1ModelFormat: "checkpoint" | "diffusers";
-    /**
-     * T2IAdapterModelFormat
-     * @description An enumeration.
-     * @enum {string}
-     */
-    T2IAdapterModelFormat: "diffusers";
-    /**
-     * ControlNetModelFormat
-     * @description An enumeration.
-     * @enum {string}
-     */
-    ControlNetModelFormat: "checkpoint" | "diffusers";
  };
  responses: never;
  parameters: never;
--- a/invokeai/frontend/web/src/services/events/types.ts
+++ b/invokeai/frontend/web/src/services/events/types.ts
@ -170,16 +170,40 @@ export type InvocationRetrievalErrorEvent = {
 */
 export type QueueItemStatusChangedEvent = {
  queue_id: string;
-  queue_item_id: number;
-  queue_batch_id: string;
-  session_id: string;
-  graph_execution_state_id: string;
-  status: components['schemas']['SessionQueueItemDTO']['status'];
-  error: string | undefined;
-  created_at: string;
-  updated_at: string;
-  started_at: string | undefined;
-  completed_at: string | undefined;
+  queue_item: {
+    queue_id: string;
+    item_id: number;
+    batch_id: string;
+    session_id: string;
+    status: components['schemas']['SessionQueueItemDTO']['status'];
+    error: string | undefined;
+    created_at: string;
+    updated_at: string;
+    started_at: string | undefined;
+    completed_at: string | undefined;
+  };
+  batch_status: {
+    queue_id: string;
+    batch_id: string;
+    pending: number;
+    in_progress: number;
+    completed: number;
+    failed: number;
+    canceled: number;
+    total: number;
+  };
+  queue_status: {
+    queue_id: string;
+    item_id?: number;
+    batch_id?: string;
+    session_id?: string;
+    pending: number;
+    in_progress: number;
+    completed: number;
+    failed: number;
+    canceled: number;
+    total: number;
+  };
 };

 export type ClientEmitSubscribeQueue = {