From fde8fc7575e93080ec8ddf21a58a0166abf77b24 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Wed, 26 Jun 2024 22:11:50 +1000
Subject: [PATCH] perf(ui): optimistic updates for getImageDTO query cache

---
 .../web/src/services/api/endpoints/images.ts        | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/invokeai/frontend/web/src/services/api/endpoints/images.ts b/invokeai/frontend/web/src/services/api/endpoints/images.ts
index 1e18d86aec..f6e71d3433 100644
--- a/invokeai/frontend/web/src/services/api/endpoints/images.ts
+++ b/invokeai/frontend/web/src/services/api/endpoints/images.ts
@@ -52,6 +52,16 @@ export const imagesApi = api.injectEndpoints({
           'FetchOnReconnect',
         ];
       },
+      async onQueryStarted(_, { dispatch, queryFulfilled }) {
+        // Populate the getImageDTO cache with these images. This makes image selection smoother, because it doesn't
+        // need to re-fetch image data when the user selects an image. The getImageDTO cache keeps data for the default
+        // of 60s, so this data won't stick around too long.
+        queryFulfilled.then(({ data }) => {
+          for (const imageDTO of data.items) {
+            dispatch(imagesApi.util.upsertQueryData('getImageDTO', imageDTO.image_name, imageDTO));
+          }
+        });
+      },
     }),
     getIntermediatesCount: build.query<number, void>({
       query: () => ({ url: buildImagesUrl('intermediates') }),
@@ -64,17 +74,14 @@ export const imagesApi = api.injectEndpoints({
     getImageDTO: build.query<ImageDTO, string>({
       query: (image_name) => ({ url: buildImagesUrl(`i/${image_name}`) }),
       providesTags: (result, error, image_name) => [{ type: 'Image', id: image_name }],
-      keepUnusedDataFor: 86400, // 24 hours
     }),
     getImageMetadata: build.query<JSONObject | undefined, string>({
       query: (image_name) => ({ url: buildImagesUrl(`i/${image_name}/metadata`) }),
       providesTags: (result, error, image_name) => [{ type: 'ImageMetadata', id: image_name }],
-      keepUnusedDataFor: 86400, // 24 hours
     }),
     getImageWorkflow: build.query<GraphAndWorkflowResponse, string>({
       query: (image_name) => ({ url: buildImagesUrl(`i/${image_name}/workflow`) }),
       providesTags: (result, error, image_name) => [{ type: 'ImageWorkflow', id: image_name }],
-      keepUnusedDataFor: 86400, // 24 hours
     }),
     deleteImage: build.mutation<void, ImageDTO>({
       query: ({ image_name }) => ({