From 1358c5eb7d5968177acc582a64018cfad5577335 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Tue, 4 Jul 2023 22:45:45 +1000
Subject: [PATCH] fix(ui): fix selector memoization

Every `GalleryImage` was rerendering any time the app rerendered bc the selector function itself was not memoized. This resulted in the memoization cache inside the selector constantly being reset.

Same for `BatchImage`.

Also updated memoization for a few other selectors.
---
 .../features/batch/components/BatchImage.tsx  | 40 +++++++++-------
 .../gallery/components/GalleryImage.tsx       | 47 ++++++++++---------
 .../lora/components/ParamLoraList.tsx         | 13 +++--
 .../Parameters/Core/ParamCFGScale.tsx         |  4 +-
 .../Parameters/Core/ParamHeight.tsx           |  4 +-
 .../Parameters/Core/ParamIterations.tsx       | 47 ++++++++++---------
 .../components/Parameters/Core/ParamSteps.tsx |  4 +-
 .../components/Parameters/Core/ParamWidth.tsx |  7 +--
 8 files changed, 92 insertions(+), 74 deletions(-)

diff --git a/invokeai/frontend/web/src/features/batch/components/BatchImage.tsx b/invokeai/frontend/web/src/features/batch/components/BatchImage.tsx
index 822b1cf183..3394946972 100644
--- a/invokeai/frontend/web/src/features/batch/components/BatchImage.tsx
+++ b/invokeai/frontend/web/src/features/batch/components/BatchImage.tsx
@@ -1,28 +1,29 @@
 import { Box, Icon, Skeleton } from '@chakra-ui/react';
+import { createSelector } from '@reduxjs/toolkit';
+import { TypesafeDraggableData } from 'app/components/ImageDnd/typesafeDnd';
+import { stateSelector } from 'app/store/store';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import { FaExclamationCircle } from 'react-icons/fa';
-import { useGetImageDTOQuery } from 'services/api/endpoints/images';
-import { MouseEvent, memo, useCallback, useMemo } from 'react';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
+import IAIDndImage from 'common/components/IAIDndImage';
 import {
   batchImageRangeEndSelected,
   batchImageSelected,
   batchImageSelectionToggled,
   imageRemovedFromBatch,
 } from 'features/batch/store/batchSlice';
-import IAIDndImage from 'common/components/IAIDndImage';
-import { createSelector } from '@reduxjs/toolkit';
-import { RootState, stateSelector } from 'app/store/store';
-import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
-import { TypesafeDraggableData } from 'app/components/ImageDnd/typesafeDnd';
+import { MouseEvent, memo, useCallback, useMemo } from 'react';
+import { FaExclamationCircle } from 'react-icons/fa';
+import { useGetImageDTOQuery } from 'services/api/endpoints/images';
 
-const isSelectedSelector = createSelector(
-  [stateSelector, (state: RootState, imageName: string) => imageName],
-  (state, imageName) => ({
-    selection: state.batch.selection,
-    isSelected: state.batch.selection.includes(imageName),
-  }),
-  defaultSelectorOptions
-);
+const makeSelector = (image_name: string) =>
+  createSelector(
+    [stateSelector],
+    (state) => ({
+      selection: state.batch.selection,
+      isSelected: state.batch.selection.includes(image_name),
+    }),
+    defaultSelectorOptions
+  );
 
 type BatchImageProps = {
   imageName: string;
@@ -37,10 +38,13 @@ const BatchImage = (props: BatchImageProps) => {
   } = useGetImageDTOQuery(props.imageName);
   const dispatch = useAppDispatch();
 
-  const { isSelected, selection } = useAppSelector((state) =>
-    isSelectedSelector(state, props.imageName)
+  const selector = useMemo(
+    () => makeSelector(props.imageName),
+    [props.imageName]
   );
 
+  const { isSelected, selection } = useAppSelector(selector);
+
   const handleClickRemove = useCallback(() => {
     dispatch(imageRemovedFromBatch(props.imageName));
   }, [dispatch, props.imageName]);
diff --git a/invokeai/frontend/web/src/features/gallery/components/GalleryImage.tsx b/invokeai/frontend/web/src/features/gallery/components/GalleryImage.tsx
index 30e1c5abf3..7b2e27ddbe 100644
--- a/invokeai/frontend/web/src/features/gallery/components/GalleryImage.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/GalleryImage.tsx
@@ -1,34 +1,35 @@
 import { Box } from '@chakra-ui/react';
-import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import { MouseEvent, memo, useCallback, useMemo } from 'react';
-import { FaTrash } from 'react-icons/fa';
-import { useTranslation } from 'react-i18next';
 import { createSelector } from '@reduxjs/toolkit';
-import { ImageDTO } from 'services/api/types';
 import { TypesafeDraggableData } from 'app/components/ImageDnd/typesafeDnd';
 import { stateSelector } from 'app/store/store';
-import ImageContextMenu from './ImageContextMenu';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import IAIDndImage from 'common/components/IAIDndImage';
+import { imageToDeleteSelected } from 'features/imageDeletion/store/imageDeletionSlice';
+import { MouseEvent, memo, useCallback, useMemo } from 'react';
+import { useTranslation } from 'react-i18next';
+import { FaTrash } from 'react-icons/fa';
+import { ImageDTO } from 'services/api/types';
 import {
   imageRangeEndSelected,
   imageSelected,
   imageSelectionToggled,
 } from '../store/gallerySlice';
-import { imageToDeleteSelected } from 'features/imageDeletion/store/imageDeletionSlice';
+import ImageContextMenu from './ImageContextMenu';
 
-export const selector = createSelector(
-  [stateSelector, (state, { image_name }: ImageDTO) => image_name],
-  ({ gallery }, image_name) => {
-    const isSelected = gallery.selection.includes(image_name);
-    const selection = gallery.selection;
-    return {
-      isSelected,
-      selection,
-    };
-  },
-  defaultSelectorOptions
-);
+export const makeSelector = (image_name: string) =>
+  createSelector(
+    [stateSelector],
+    ({ gallery }) => {
+      const isSelected = gallery.selection.includes(image_name);
+      const selection = gallery.selection;
+      return {
+        isSelected,
+        selection,
+      };
+    },
+    defaultSelectorOptions
+  );
 
 interface HoverableImageProps {
   imageDTO: ImageDTO;
@@ -38,13 +39,13 @@ interface HoverableImageProps {
  * Gallery image component with delete/use all/use seed buttons on hover.
  */
 const GalleryImage = (props: HoverableImageProps) => {
-  const { isSelected, selection } = useAppSelector((state) =>
-    selector(state, props.imageDTO)
-  );
-
   const { imageDTO } = props;
   const { image_url, thumbnail_url, image_name } = imageDTO;
 
+  const localSelector = useMemo(() => makeSelector(image_name), [image_name]);
+
+  const { isSelected, selection } = useAppSelector(localSelector);
+
   const dispatch = useAppDispatch();
 
   const { t } = useTranslation();
diff --git a/invokeai/frontend/web/src/features/lora/components/ParamLoraList.tsx b/invokeai/frontend/web/src/features/lora/components/ParamLoraList.tsx
index 8d6ff98498..89432ac862 100644
--- a/invokeai/frontend/web/src/features/lora/components/ParamLoraList.tsx
+++ b/invokeai/frontend/web/src/features/lora/components/ParamLoraList.tsx
@@ -1,14 +1,19 @@
 import { createSelector } from '@reduxjs/toolkit';
 import { stateSelector } from 'app/store/store';
 import { useAppSelector } from 'app/store/storeHooks';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import { map } from 'lodash-es';
 import ParamLora from './ParamLora';
 
-const selector = createSelector(stateSelector, ({ lora }) => {
-  const { loras } = lora;
+const selector = createSelector(
+  stateSelector,
+  ({ lora }) => {
+    const { loras } = lora;
 
-  return { loras };
-});
+    return { loras };
+  },
+  defaultSelectorOptions
+);
 
 const ParamLoraList = () => {
   const { loras } = useAppSelector(selector);
diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamCFGScale.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamCFGScale.tsx
index 111e3d3ae8..d32ff960d5 100644
--- a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamCFGScale.tsx
+++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamCFGScale.tsx
@@ -1,5 +1,6 @@
 import { createSelector } from '@reduxjs/toolkit';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import IAINumberInput from 'common/components/IAINumberInput';
 import IAISlider from 'common/components/IAISlider';
 import { generationSelector } from 'features/parameters/store/generationSelectors';
@@ -27,7 +28,8 @@ const selector = createSelector(
       shouldUseSliders,
       shift,
     };
-  }
+  },
+  defaultSelectorOptions
 );
 
 const ParamCFGScale = () => {
diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamHeight.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamHeight.tsx
index 9501c8b475..6939ede424 100644
--- a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamHeight.tsx
+++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamHeight.tsx
@@ -1,5 +1,6 @@
 import { createSelector } from '@reduxjs/toolkit';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import IAISlider, { IAIFullSliderProps } from 'common/components/IAISlider';
 import { generationSelector } from 'features/parameters/store/generationSelectors';
 import { setHeight } from 'features/parameters/store/generationSlice';
@@ -25,7 +26,8 @@ const selector = createSelector(
       inputMax,
       step,
     };
-  }
+  },
+  defaultSelectorOptions
 );
 
 type ParamHeightProps = Omit<
diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamIterations.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamIterations.tsx
index a8cdabc8c9..1e203a1e45 100644
--- a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamIterations.tsx
+++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamIterations.tsx
@@ -1,37 +1,38 @@
 import { createSelector } from '@reduxjs/toolkit';
 import { stateSelector } from 'app/store/store';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import IAINumberInput from 'common/components/IAINumberInput';
 import IAISlider from 'common/components/IAISlider';
-import { generationSelector } from 'features/parameters/store/generationSelectors';
 import { setIterations } from 'features/parameters/store/generationSlice';
-import { configSelector } from 'features/system/store/configSelectors';
-import { hotkeysSelector } from 'features/ui/store/hotkeysSlice';
-import { uiSelector } from 'features/ui/store/uiSelectors';
 import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 
-const selector = createSelector([stateSelector], (state) => {
-  const { initial, min, sliderMax, inputMax, fineStep, coarseStep } =
-    state.config.sd.iterations;
-  const { iterations } = state.generation;
-  const { shouldUseSliders } = state.ui;
-  const isDisabled =
-    state.dynamicPrompts.isEnabled && state.dynamicPrompts.combinatorial;
+const selector = createSelector(
+  [stateSelector],
+  (state) => {
+    const { initial, min, sliderMax, inputMax, fineStep, coarseStep } =
+      state.config.sd.iterations;
+    const { iterations } = state.generation;
+    const { shouldUseSliders } = state.ui;
+    const isDisabled =
+      state.dynamicPrompts.isEnabled && state.dynamicPrompts.combinatorial;
 
-  const step = state.hotkeys.shift ? fineStep : coarseStep;
+    const step = state.hotkeys.shift ? fineStep : coarseStep;
 
-  return {
-    iterations,
-    initial,
-    min,
-    sliderMax,
-    inputMax,
-    step,
-    shouldUseSliders,
-    isDisabled,
-  };
-});
+    return {
+      iterations,
+      initial,
+      min,
+      sliderMax,
+      inputMax,
+      step,
+      shouldUseSliders,
+      isDisabled,
+    };
+  },
+  defaultSelectorOptions
+);
 
 const ParamIterations = () => {
   const {
diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamSteps.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamSteps.tsx
index f43cdd425b..d939113c7c 100644
--- a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamSteps.tsx
+++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamSteps.tsx
@@ -1,5 +1,6 @@
 import { createSelector } from '@reduxjs/toolkit';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import IAINumberInput from 'common/components/IAINumberInput';
 
 import IAISlider from 'common/components/IAISlider';
@@ -33,7 +34,8 @@ const selector = createSelector(
       step,
       shouldUseSliders,
     };
-  }
+  },
+  defaultSelectorOptions
 );
 
 const ParamSteps = () => {
diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamWidth.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamWidth.tsx
index b7d63038d1..b4121184b5 100644
--- a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamWidth.tsx
+++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamWidth.tsx
@@ -1,7 +1,7 @@
 import { createSelector } from '@reduxjs/toolkit';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import IAISlider from 'common/components/IAISlider';
-import { IAIFullSliderProps } from 'common/components/IAISlider';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
+import IAISlider, { IAIFullSliderProps } from 'common/components/IAISlider';
 import { generationSelector } from 'features/parameters/store/generationSelectors';
 import { setWidth } from 'features/parameters/store/generationSlice';
 import { configSelector } from 'features/system/store/configSelectors';
@@ -26,7 +26,8 @@ const selector = createSelector(
       inputMax,
       step,
     };
-  }
+  },
+  defaultSelectorOptions
 );
 
 type ParamWidthProps = Omit<IAIFullSliderProps, 'label' | 'value' | 'onChange'>;