From b46b20210dd1d2dceae8042581afc8444a996a55 Mon Sep 17 00:00:00 2001
From: chainchompa <chainchompa@gmail.com>
Date: Tue, 23 Jul 2024 11:53:43 -0400
Subject: [PATCH] handle simple upscale models on modelsLoaded

---
 .../listeners/modelsLoaded.ts                 | 23 ++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelsLoaded.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelsLoaded.ts
index 2ace69c54e..5d7213bb57 100644
--- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelsLoaded.ts
+++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelsLoaded.ts
@@ -10,7 +10,7 @@ import { heightChanged, widthChanged } from 'features/controlLayers/store/contro
 import { loraRemoved } from 'features/lora/store/loraSlice';
 import { calculateNewSize } from 'features/parameters/components/ImageSize/calculateNewSize';
 import { modelChanged, vaeSelected } from 'features/parameters/store/generationSlice';
-import { upscaleModelChanged } from 'features/parameters/store/upscaleSlice';
+import { simpleUpscaleModelChanged, upscaleModelChanged } from 'features/parameters/store/upscaleSlice';
 import { zParameterModel, zParameterVAEModel } from 'features/parameters/types/parameterSchemas';
 import { getIsSizeOptimal, getOptimalDimension } from 'features/parameters/util/optimalDimension';
 import { refinerModelChanged } from 'features/sdxl/store/sdxlSlice';
@@ -43,6 +43,7 @@ export const addModelsLoadedListener = (startAppListening: AppStartListening) =>
       handleLoRAModels(models, state, dispatch, log);
       handleControlAdapterModels(models, state, dispatch, log);
       handleSpandrelImageToImageModels(models, state, dispatch, log);
+      handleSimpleUpscaleModels(models, state, dispatch, log);
     },
   });
 };
@@ -204,3 +205,23 @@ const handleSpandrelImageToImageModels: ModelHandler = (models, state, dispatch,
 
   dispatch(upscaleModelChanged(null));
 };
+
+const handleSimpleUpscaleModels: ModelHandler = (models, state, dispatch, _log) => {
+  const currentSimpleUpscaleModel = state.upscale.simpleUpscaleModel;
+  const upscaleModels = models.filter(isSpandrelImageToImageModelConfig);
+
+  if (currentSimpleUpscaleModel) {
+    const isCurrentSimpleUpscaleModelAvailable = upscaleModels.some((m) => m.key === currentSimpleUpscaleModel.key);
+    if (isCurrentSimpleUpscaleModelAvailable) {
+      return;
+    }
+  }
+
+  const firstModel = upscaleModels[0];
+  if (firstModel) {
+    dispatch(simpleUpscaleModelChanged(firstModel));
+    return;
+  }
+
+  dispatch(simpleUpscaleModelChanged(null));
+};