From 273dfd86abd1a0372179e7a9d65009ad54941a07 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Mon, 28 Apr 2025 13:33:55 +1000
Subject: [PATCH] fix(ui): upscale builder

---
 .../listeners/enqueueRequestedUpscale.ts             | 12 ++++++++++--
 .../util/graph/buildMultidiffusionUpscaleGraph.ts    | 10 +++++++---
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedUpscale.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedUpscale.ts
index 2b1281fe40..f7fc198db3 100644
--- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedUpscale.ts
+++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedUpscale.ts
@@ -18,9 +18,17 @@ export const addEnqueueRequestedUpscale = (startAppListening: AppStartListening)
       const state = getState();
       const { prepend } = action.payload;
 
-      const { g, noise, posCond } = await buildMultidiffusionUpscaleGraph(state);
+      const { g, seedFieldIdentifier, positivePromptFieldIdentifier } = await buildMultidiffusionUpscaleGraph(state);
 
-      const batchConfig = prepareLinearUIBatch(state, g, prepend, noise, posCond, 'upscaling', 'gallery');
+      const batchConfig = prepareLinearUIBatch(
+        state,
+        g,
+        prepend,
+        seedFieldIdentifier,
+        positivePromptFieldIdentifier,
+        'upscaling',
+        'gallery'
+      );
 
       const req = dispatch(queueApi.endpoints.enqueueBatch.initiate(batchConfig, enqueueMutationFixedCacheKeyOptions));
       try {
diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/buildMultidiffusionUpscaleGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/buildMultidiffusionUpscaleGraph.ts
index c63f6fc582..d344ff439a 100644
--- a/invokeai/frontend/web/src/features/nodes/util/graph/buildMultidiffusionUpscaleGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graph/buildMultidiffusionUpscaleGraph.ts
@@ -1,9 +1,9 @@
 import type { RootState } from 'app/store/store';
 import { getPrefixedId } from 'features/controlLayers/konva/util';
 import { fetchModelConfigWithTypeGuard } from 'features/metadata/util/modelFetchingHelpers';
+import type { FieldIdentifier } from 'features/nodes/types/field';
 import { addSDXLLoRAs } from 'features/nodes/util/graph/generation/addSDXLLoRAs';
 import { Graph } from 'features/nodes/util/graph/generation/Graph';
-import type { Invocation } from 'services/api/types';
 import { isNonRefinerMainModelConfig, isSpandrelImageToImageModelConfig } from 'services/api/types';
 import { assert } from 'tsafe';
 
@@ -12,7 +12,7 @@ import { getBoardField, selectPresetModifiedPrompts } from './graphBuilderUtils'
 
 export const buildMultidiffusionUpscaleGraph = async (
   state: RootState
-): Promise<{ g: Graph; noise: Invocation<'noise'>; posCond: Invocation<'compel' | 'sdxl_compel_prompt'> }> => {
+): Promise<{ g: Graph; seedFieldIdentifier: FieldIdentifier; positivePromptFieldIdentifier: FieldIdentifier }> => {
   const {
     model,
     upscaleCfgScale: cfg_scale,
@@ -243,5 +243,9 @@ export const buildMultidiffusionUpscaleGraph = async (
 
   g.addEdge(controlNetCollector, 'collection', tiledMultidiffusion, 'control');
 
-  return { g, noise, posCond };
+  return {
+    g,
+    seedFieldIdentifier: { nodeId: noise.id, fieldName: 'seed' },
+    positivePromptFieldIdentifier: { nodeId: posCond.id, fieldName: 'prompt' },
+  };
 };