feat(ui): gallery optimistic updates for video

Author: psychedelicious

invokeai/frontend/web/src/services/api/endpoints/videos.ts

@@ -1,25 +1,25 @@
+import { getStore } from 'app/store/nanostores/store';
 import type { paths } from 'services/api/schema';
-import type {
-  GetVideoIdsArgs,
-  GetVideoIdsResult,
-  VideoDTO,
-} from 'services/api/types';
+import type { GetVideoIdsArgs, GetVideoIdsResult, VideoDTO } from 'services/api/types';
+import {
+  getTagsToInvalidateForBoardAffectingMutation,
+  getTagsToInvalidateForVideoMutation,
+} from 'services/api/util/tagInvalidation';
 import stableHash from 'stable-hash';
 import type { Param0 } from 'tsafe';
 
 import { api, buildV1Url, LIST_TAG } from '..';
-import { getTagsToInvalidateForBoardAffectingMutation, getTagsToInvalidateForImageMutation, getTagsToInvalidateForVideoMutation } from '../util/tagInvalidation';
 
 /**
  * Builds an endpoint URL for the videos router
  * @example
  * buildVideosUrl('some-path')
  * // '/api/v1/videos/some-path'
  */
-const buildVideosUrl  = (path: string = '', query?: Parameters<typeof buildV1Url>[1]) =>
+const buildVideosUrl = (path: string = '', query?: Parameters<typeof buildV1Url>[1]) =>
   buildV1Url(`videos/${path}`, query);
 
-const buildBoardVideosUrl = (path: string = '') => buildV1Url(`board_videos/${path}`);  
+const buildBoardVideosUrl = (path: string = '') => buildV1Url(`board_videos/${path}`);
 
 export const videosApi = api.injectEndpoints({
   endpoints: (build) => ({
@@ -31,7 +31,6 @@ export const videosApi = api.injectEndpoints({
       query: (video_id) => ({ url: buildVideosUrl(`i/${video_id}`) }),
       providesTags: (result, error, video_id) => [{ type: 'Video', id: video_id }],
     }),
-    
 
     /**
      * Get ordered list of image names for selection operations
@@ -204,4 +203,24 @@ export const {
   useRemoveVideosFromBoardMutation,
 } = videosApi;
 
-
+/**
+ * Imperative RTKQ helper to fetch an VideoDTO.
+ * @param id The id of the video to fetch
+ * @param options The options for the query. By default, the query will not subscribe to the store.
+ * @returns The ImageDTO if found, otherwise null
+ */
+export const getVideoDTOSafe = async (
+  id: string,
+  options?: Parameters<typeof videosApi.endpoints.getVideoDTOsByNames.initiate>[1]
+): Promise<VideoDTO | null> => {
+  const _options = {
+    subscribe: false,
+    ...options,
+  };
+  const req = getStore().dispatch(videosApi.endpoints.getVideoDTOsByNames.initiate({ video_ids: [id] }, _options));
+  try {
+    return (await req.unwrap())[0] ?? null;
+  } catch {
+    return null;
+  }
+};

invokeai/frontend/web/src/services/api/util/optimisticUpdates.ts

@@ -1,5 +1,5 @@
 import type { OrderDir } from 'features/gallery/store/types';
-import type { GetImageNamesResult, ImageDTO } from 'services/api/types';
+import type { GetImageNamesResult, GetVideoIdsResult, ImageDTO, VideoDTO } from 'services/api/types';
 
 /**
  * Calculates the optimal insertion position for a new image in the names list.
@@ -57,3 +57,60 @@ export function insertImageIntoNamesResult(
     total_count: currentResult.total_count + 1,
   };
 }
+
+/**
+ * Calculates the optimal insertion position for a new image in the names list.
+ * For starred_first=true: starred images go to position 0, unstarred go after all starred images
+ * For starred_first=false: all new images go to position 0 (newest first)
+ */
+function calculateVideoInsertionPosition(
+  videoDTO: VideoDTO,
+  starredFirst: boolean,
+  starredCount: number,
+  orderDir: OrderDir = 'DESC'
+): number {
+  if (!starredFirst) {
+    // When starred_first is false, insertion depends on order direction
+    return orderDir === 'DESC' ? 0 : Number.MAX_SAFE_INTEGER;
+  }
+
+  // When starred_first is true
+  if (videoDTO.starred) {
+    // Starred images: beginning for desc, after existing starred for asc
+    return orderDir === 'DESC' ? 0 : starredCount;
+  }
+
+  // Unstarred images go after all starred images
+  return orderDir === 'DESC' ? starredCount : Number.MAX_SAFE_INTEGER;
+}
+
+/**
+ * Optimistically inserts a new image into the ImageNamesResult at the correct position
+ */
+export function insertVideoIntoGetVideoIdsResult(
+  currentResult: GetVideoIdsResult,
+  videoDTO: VideoDTO,
+  starredFirst: boolean,
+  orderDir: OrderDir = 'DESC'
+): GetVideoIdsResult {
+  // Don't insert if the image is already in the list
+  if (currentResult.video_ids.includes(videoDTO.video_id)) {
+    return currentResult;
+  }
+
+  const insertPosition = calculateVideoInsertionPosition(videoDTO, starredFirst, currentResult.starred_count, orderDir);
+
+  const newVideoIds = [...currentResult.video_ids];
+  // Handle MAX_SAFE_INTEGER by pushing to end
+  if (insertPosition >= newVideoIds.length) {
+    newVideoIds.push(videoDTO.video_id);
+  } else {
+    newVideoIds.splice(insertPosition, 0, videoDTO.video_id);
+  }
+
+  return {
+    video_ids: newVideoIds,
+    starred_count: starredFirst && videoDTO.starred ? currentResult.starred_count + 1 : currentResult.starred_count,
+    total_count: currentResult.total_count + 1,
+  };
+}

invokeai/frontend/web/src/services/events/onInvocationComplete.tsx

@@ -5,6 +5,7 @@ import {
   selectAutoSwitch,
   selectGalleryView,
   selectGetImageNamesQueryArgs,
+  selectGetVideoIdsQueryArgs,
   selectListBoardsQueryArgs,
   selectSelectedBoardId,
 } from 'features/gallery/store/gallerySelectors';
@@ -17,9 +18,10 @@ import { generatedVideoChanged } from 'features/parameters/store/videoSlice';
 import type { LRUCache } from 'lru-cache';
 import { boardsApi } from 'services/api/endpoints/boards';
 import { getImageDTOSafe, imagesApi } from 'services/api/endpoints/images';
-import type { ImageDTO, S } from 'services/api/types';
+import { getVideoDTOSafe, videosApi } from 'services/api/endpoints/videos';
+import type { ImageDTO, S, VideoDTO } from 'services/api/types';
 import { getCategories } from 'services/api/util';
-import { insertImageIntoNamesResult } from 'services/api/util/optimisticUpdates';
+import { insertImageIntoNamesResult, insertVideoIntoGetVideoIdsResult } from 'services/api/util/optimisticUpdates';
 import { $lastProgressEvent } from 'services/events/stores';
 import stableHash from 'stable-hash';
 import type { Param0 } from 'tsafe';
@@ -185,6 +187,154 @@ export const buildOnInvocationComplete = (
     }
   };
 
+  const addVideosToGallery = async (data: S['InvocationCompleteEvent']) => {
+    if (nodeTypeDenylist.includes(data.invocation.type)) {
+      log.trace(`Skipping denylisted node type (${data.invocation.type})`);
+      return;
+    }
+
+    const videoDTOs = await getResultVideoDTOs(data);
+    if (videoDTOs.length === 0) {
+      return;
+    }
+
+    // For efficiency's sake, we want to minimize the number of dispatches and invalidations we do.
+    // We'll keep track of each change we need to make and do them all at once.
+    const boardTotalAdditions: Record<string, number> = {};
+    const getVideoIdsArg = selectGetVideoIdsQueryArgs(getState());
+
+    for (const videoDTO of videoDTOs) {
+      if (videoDTO.is_intermediate) {
+        return;
+      }
+
+      const board_id = videoDTO.board_id ?? 'none';
+
+      boardTotalAdditions[board_id] = (boardTotalAdditions[board_id] || 0) + 1;
+    }
+
+    // Update all the board image totals at once
+    const entries: Param0<typeof boardsApi.util.upsertQueryEntries> = [];
+    for (const [boardId, amountToAdd] of objectEntries(boardTotalAdditions)) {
+      // upsertQueryEntries doesn't provide a "recipe" function for the update - we must provide the new value
+      // directly. So we need to select the board totals first.
+      const total = boardsApi.endpoints.getBoardImagesTotal.select(boardId)(getState()).data?.total;
+      if (total === undefined) {
+        // No cache exists for this board, so we can't update it.
+        continue;
+      }
+      entries.push({
+        endpointName: 'getBoardImagesTotal',
+        arg: boardId,
+        value: { total: total + amountToAdd },
+      });
+    }
+    dispatch(boardsApi.util.upsertQueryEntries(entries));
+
+    dispatch(
+      boardsApi.util.updateQueryData('listAllBoards', selectListBoardsQueryArgs(getState()), (draft) => {
+        for (const board of draft) {
+          board.image_count = board.image_count + (boardTotalAdditions[board.board_id] ?? 0);
+        }
+      })
+    );
+
+    /**
+     * Optimistic update and cache invalidation for image names queries that match this image's board and categories.
+     * - Optimistic update for the cache that does not have a search term (we cannot derive the correct insertion
+     *   position when a search term is present).
+     * - Cache invalidation for the query that has a search term, so it will be refetched.
+     *
+     * Note: The image DTO objects are already implicitly cached by the getResultImageDTOs function. We do not need
+     * to explicitly cache them again here.
+     */
+    for (const videoDTO of videoDTOs) {
+      // Override board_id and categories for this specific image to build the "expected" args for the query.
+      const videoSpecificArgs = {
+        board_id: videoDTO.board_id ?? 'none',
+      };
+
+      const expectedQueryArgs = {
+        ...getVideoIdsArg,
+        ...videoSpecificArgs,
+        search_term: '',
+      };
+
+      // If the cache for the query args provided here does not exist, RTK Query will ignore the update.
+      dispatch(
+        videosApi.util.updateQueryData(
+          'getVideoIds',
+          {
+            ...getVideoIdsArg,
+            ...videoSpecificArgs,
+            search_term: '',
+          },
+          (draft) => {
+            const updatedResult = insertVideoIntoGetVideoIdsResult(
+              draft,
+              videoDTO,
+              expectedQueryArgs.starred_first ?? true,
+              expectedQueryArgs.order_dir
+            );
+
+            draft.video_ids = updatedResult.video_ids;
+            draft.starred_count = updatedResult.starred_count;
+            draft.total_count = updatedResult.total_count;
+          }
+        )
+      );
+
+      // If there is a search term present, we need to invalidate that query to ensure the search results are updated.
+      if (getVideoIdsArg.search_term) {
+        const expectedQueryArgs = {
+          ...getVideoIdsArg,
+          ...videoSpecificArgs,
+        };
+        dispatch(imagesApi.util.invalidateTags([{ type: 'ImageNameList', id: stableHash(expectedQueryArgs) }]));
+      }
+    }
+
+    // No need to invalidate tags since we're doing optimistic updates
+    // Board totals are already updated above via upsertQueryEntries
+
+    const autoSwitch = selectAutoSwitch(getState());
+
+    if (!autoSwitch) {
+      return;
+    }
+
+    // Finally, we may need to autoswitch to the new image. We'll only do it for the last image in the list.
+    const lastVideoDTO = videoDTOs.at(-1);
+
+    if (!lastVideoDTO) {
+      return;
+    }
+
+    const { video_id } = lastVideoDTO;
+    const board_id = lastVideoDTO.board_id ?? 'none';
+
+    // With optimistic updates, we can immediately switch to the new image
+    const selectedBoardId = selectSelectedBoardId(getState());
+
+    // If the image is from a different board, switch to that board & select the image - otherwise just select the
+    // image. This implicitly changes the view to 'images' if it was not already.
+    if (board_id !== selectedBoardId) {
+      dispatch(
+        boardIdSelected({
+          boardId: board_id,
+          selectedImageName: video_id,
+        })
+      );
+    } else {
+      // Ensure we are on the 'images' gallery view - that's where this image will be displayed
+      const galleryView = selectGalleryView(getState());
+      if (galleryView !== 'videos') {
+        dispatch(galleryViewChanged('videos'));
+      }
+      // Select the image immediately since we've optimistically updated the cache
+      dispatch(imageSelected(lastVideoDTO.video_id));
+    }
+  };
   const getResultImageDTOs = async (data: S['InvocationCompleteEvent']): Promise<ImageDTO[]> => {
     const { result } = data;
     const imageDTOs: ImageDTO[] = [];
@@ -206,17 +356,20 @@ export const buildOnInvocationComplete = (
     return imageDTOs;
   };
 
-  const getResultVideoFields = (data: S['InvocationCompleteEvent']): VideoField[] => {
+  const getResultVideoDTOs = async (data: S['InvocationCompleteEvent']): Promise<VideoDTO[]> => {
     const { result } = data;
-    const videoFields: VideoField[] = [];
+    const videoDTOs: VideoDTO[] = [];
 
     for (const [_name, value] of objectEntries(result)) {
       if (isVideoField(value)) {
-        videoFields.push(value);
+        const videoDTO = await getVideoDTOSafe(value.video_id);
+        if (videoDTO) {
+          videoDTOs.push(videoDTO);
+        }
       }
     }
 
-    return videoFields;
+    return videoDTOs;
   };
 
   return async (data: S['InvocationCompleteEvent']) => {
@@ -239,11 +392,7 @@ export const buildOnInvocationComplete = (
     }
 
     await addImagesToGallery(data);
-
-    const videoField = getResultVideoFields(data)[0];
-    if (videoField) {
-      dispatch(generatedVideoChanged({ videoField }));
-    }
+    await addVideosToGallery(data);
 
     $lastProgressEvent.set(null);
   };