From 2e4e9434c199484ea5a21e67f6493a1c41980792 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Sat, 29 Apr 2023 19:11:20 +1000
Subject: [PATCH] fix(ui): fix initial image for uploads

---
 .../components/CurrentImageButtons.tsx        |  2 +-
 .../ImageToImage/InitialImagePreview.tsx      | 13 +++++--------
 .../parameters/store/generationSelectors.ts   | 19 +++++++++----------
 .../parameters/store/generationSlice.ts       | 10 ++++++++--
 4 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/invokeai/frontend/web/src/features/gallery/components/CurrentImageButtons.tsx b/invokeai/frontend/web/src/features/gallery/components/CurrentImageButtons.tsx
index 3cfe932b75..481e3feb40 100644
--- a/invokeai/frontend/web/src/features/gallery/components/CurrentImageButtons.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/CurrentImageButtons.tsx
@@ -166,7 +166,7 @@ const CurrentImageButtons = (props: CurrentImageButtonsProps) => {
   const handleClickUseAsInitialImage = useCallback(() => {
     if (!image) return;
     if (isLightboxOpen) dispatch(setIsLightboxOpen(false));
-    dispatch(initialImageSelected(image.name));
+    dispatch(initialImageSelected({ name: image.name, type: image.type }));
     // dispatch(setInitialImage(currentImage));
 
     // dispatch(setActiveTab('img2img'));
diff --git a/invokeai/frontend/web/src/features/parameters/components/AdvancedParameters/ImageToImage/InitialImagePreview.tsx b/invokeai/frontend/web/src/features/parameters/components/AdvancedParameters/ImageToImage/InitialImagePreview.tsx
index 2082c2a015..8a91602ada 100644
--- a/invokeai/frontend/web/src/features/parameters/components/AdvancedParameters/ImageToImage/InitialImagePreview.tsx
+++ b/invokeai/frontend/web/src/features/parameters/components/AdvancedParameters/ImageToImage/InitialImagePreview.tsx
@@ -5,7 +5,6 @@ import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import SelectImagePlaceholder from 'common/components/SelectImagePlaceholder';
 import { useGetUrl } from 'common/util/getUrl';
 import useGetImageByNameAndType from 'features/gallery/hooks/useGetImageByName';
-import { selectResultsById } from 'features/gallery/store/resultsSlice';
 import {
   clearInitialImage,
   initialImageSelected,
@@ -16,15 +15,13 @@ import { DragEvent, useCallback, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { ImageType } from 'services/api';
 import ImageToImageOverlay from 'common/components/ImageToImageOverlay';
+import { initialImageSelector } from 'features/parameters/store/generationSelectors';
 
 const initialImagePreviewSelector = createSelector(
-  [(state: RootState) => state],
-  (state) => {
-    const { initialImage } = state.generation;
-    const image = selectResultsById(state, initialImage as string);
-
+  [initialImageSelector],
+  (initialImage) => {
     return {
-      initialImage: image,
+      initialImage,
     };
   },
   { memoizeOptions: { resultEqualityCheck: isEqual } }
@@ -71,7 +68,7 @@ const InitialImagePreview = () => {
         return;
       }
 
-      dispatch(initialImageSelected(image.name));
+      dispatch(initialImageSelected({ name, type }));
     },
     [getImageByNameAndType, dispatch]
   );
diff --git a/invokeai/frontend/web/src/features/parameters/store/generationSelectors.ts b/invokeai/frontend/web/src/features/parameters/store/generationSelectors.ts
index ce3c9c4e1e..dbf5eec791 100644
--- a/invokeai/frontend/web/src/features/parameters/store/generationSelectors.ts
+++ b/invokeai/frontend/web/src/features/parameters/store/generationSelectors.ts
@@ -1,10 +1,6 @@
 import { createSelector } from '@reduxjs/toolkit';
 import { RootState } from 'app/store/store';
-import { gallerySelector } from 'features/gallery/store/gallerySelectors';
-import {
-  selectResultsById,
-  selectResultsEntities,
-} from 'features/gallery/store/resultsSlice';
+import { selectResultsById } from 'features/gallery/store/resultsSlice';
 import { selectUploadsById } from 'features/gallery/store/uploadsSlice';
 import { isEqual } from 'lodash-es';
 
@@ -25,11 +21,14 @@ export const mayGenerateMultipleImagesSelector = createSelector(
 export const initialImageSelector = createSelector(
   [(state: RootState) => state, generationSelector],
   (state, generation) => {
-    const { initialImage: initialImageName } = generation;
+    const { initialImage } = generation;
 
-    return (
-      selectResultsById(state, initialImageName as string) ??
-      selectUploadsById(state, initialImageName as string)
-    );
+    if (initialImage?.type === 'results') {
+      return selectResultsById(state, initialImage.name);
+    }
+
+    if (initialImage?.type === 'uploads') {
+      return selectUploadsById(state, initialImage.name);
+    }
   }
 );
diff --git a/invokeai/frontend/web/src/features/parameters/store/generationSlice.ts b/invokeai/frontend/web/src/features/parameters/store/generationSlice.ts
index f303491b2b..5dbe6b2c28 100644
--- a/invokeai/frontend/web/src/features/parameters/store/generationSlice.ts
+++ b/invokeai/frontend/web/src/features/parameters/store/generationSlice.ts
@@ -5,13 +5,19 @@ import { getPromptAndNegative } from 'common/util/getPromptAndNegative';
 import promptToString from 'common/util/promptToString';
 import { seedWeightsToString } from 'common/util/seedWeightPairs';
 import { clamp } from 'lodash-es';
+import { ImageField, ImageType } from 'services/api';
+
+export type InitialImage = {
+  name: string;
+  type: ImageType;
+};
 
 export interface GenerationState {
   cfgScale: number;
   height: number;
   img2imgStrength: number;
   infillMethod: string;
-  initialImage?: InvokeAI._Image | string; // can be an Image or url
+  initialImage?: InitialImage; // can be an Image or url
   iterations: number;
   maskPath: string;
   perlin: number;
@@ -345,7 +351,7 @@ export const generationSlice = createSlice({
     setVerticalSymmetrySteps: (state, action: PayloadAction<number>) => {
       state.verticalSymmetrySteps = action.payload;
     },
-    initialImageSelected: (state, action: PayloadAction<string>) => {
+    initialImageSelected: (state, action: PayloadAction<InitialImage>) => {
       state.initialImage = action.payload;
       state.isImageToImageEnabled = true;
     },