From de19128e1bafe58a32ae409bf0445465d22003af Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Fri, 5 Apr 2024 22:49:55 +0530
Subject: [PATCH] fix: Update SDXL Refiner graphs to use Gradient Mask

---
 .../nodes/util/graph/addSDXLRefinerToGraph.ts | 45 ++++++-------------
 .../nodes/util/graph/addVAEToGraph.ts         | 22 ---------
 2 files changed, 14 insertions(+), 53 deletions(-)

diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/addSDXLRefinerToGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/addSDXLRefinerToGraph.ts
index 24a83d9d19..3fe4d13cd3 100644
--- a/invokeai/frontend/web/src/features/nodes/util/graph/addSDXLRefinerToGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graph/addSDXLRefinerToGraph.ts
@@ -1,16 +1,15 @@
 import type { RootState } from 'app/store/store';
 import { fetchModelConfigWithTypeGuard } from 'features/metadata/util/modelFetchingHelpers';
-import {
-  type CreateDenoiseMaskInvocation,
-  type ImageDTO,
-  isRefinerMainModelModelConfig,
-  type NonNullableGraph,
-  type SeamlessModeInvocation,
+import type {
+  CreateGradientMaskInvocation,
+  ImageDTO,
+  NonNullableGraph,
+  SeamlessModeInvocation,
 } from 'services/api/types';
+import { isRefinerMainModelModelConfig } from 'services/api/types';
 
 import {
   CANVAS_OUTPUT,
-  INPAINT_IMAGE_RESIZE_UP,
   LATENTS_TO_IMAGE,
   MASK_COMBINE,
   MASK_RESIZE_UP,
@@ -47,15 +46,15 @@ export const addSDXLRefinerToGraph = async (
     refinerStart,
   } = state.sdxl;
 
+  const { canvasCoherenceEdgeSize, canvasCoherenceMinDenoise, canvasCoherenceMode } = state.generation;
+
   if (!refinerModel) {
     return;
   }
 
-  const { seamlessXAxis, seamlessYAxis, vaePrecision } = state.generation;
+  const { seamlessXAxis, seamlessYAxis } = state.generation;
   const { boundingBoxScaleMethod } = state.canvas;
 
-  const fp32 = vaePrecision === 'fp32';
-
   const isUsingScaledDimensions = ['auto', 'manual'].includes(boundingBoxScaleMethod);
   const modelConfig = await fetchModelConfigWithTypeGuard(refinerModel.key, isRefinerMainModelModelConfig);
 
@@ -215,30 +214,14 @@ export const addSDXLRefinerToGraph = async (
 
   if (graph.id === SDXL_CANVAS_INPAINT_GRAPH || graph.id === SDXL_CANVAS_OUTPAINT_GRAPH) {
     graph.nodes[SDXL_REFINER_INPAINT_CREATE_MASK] = {
-      type: 'create_denoise_mask',
+      type: 'create_gradient_mask',
       id: SDXL_REFINER_INPAINT_CREATE_MASK,
       is_intermediate: true,
-      fp32,
+      edge_radius: canvasCoherenceEdgeSize,
+      coherence_mode: canvasCoherenceMode,
+      minimum_denoise: canvasCoherenceMinDenoise,
     };
 
-    if (isUsingScaledDimensions) {
-      graph.edges.push({
-        source: {
-          node_id: INPAINT_IMAGE_RESIZE_UP,
-          field: 'image',
-        },
-        destination: {
-          node_id: SDXL_REFINER_INPAINT_CREATE_MASK,
-          field: 'image',
-        },
-      });
-    } else {
-      graph.nodes[SDXL_REFINER_INPAINT_CREATE_MASK] = {
-        ...(graph.nodes[SDXL_REFINER_INPAINT_CREATE_MASK] as CreateDenoiseMaskInvocation),
-        image: canvasInitImage,
-      };
-    }
-
     if (graph.id === SDXL_CANVAS_INPAINT_GRAPH) {
       if (isUsingScaledDimensions) {
         graph.edges.push({
@@ -253,7 +236,7 @@ export const addSDXLRefinerToGraph = async (
         });
       } else {
         graph.nodes[SDXL_REFINER_INPAINT_CREATE_MASK] = {
-          ...(graph.nodes[SDXL_REFINER_INPAINT_CREATE_MASK] as CreateDenoiseMaskInvocation),
+          ...(graph.nodes[SDXL_REFINER_INPAINT_CREATE_MASK] as CreateGradientMaskInvocation),
           mask: canvasMaskImage,
         };
       }
diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/addVAEToGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/addVAEToGraph.ts
index 2e1366a01e..242fc15298 100644
--- a/invokeai/frontend/web/src/features/nodes/util/graph/addVAEToGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graph/addVAEToGraph.ts
@@ -17,7 +17,6 @@ import {
   SDXL_CANVAS_OUTPAINT_GRAPH,
   SDXL_CANVAS_TEXT_TO_IMAGE_GRAPH,
   SDXL_IMAGE_TO_IMAGE_GRAPH,
-  SDXL_REFINER_INPAINT_CREATE_MASK,
   SDXL_REFINER_SEAMLESS,
   SDXL_TEXT_TO_IMAGE_GRAPH,
   SEAMLESS,
@@ -166,27 +165,6 @@ export const addVAEToGraph = async (
     );
   }
 
-  if (refinerModel) {
-    if (graph.id === SDXL_CANVAS_INPAINT_GRAPH || graph.id === SDXL_CANVAS_OUTPAINT_GRAPH) {
-      graph.edges.push({
-        source: {
-          node_id: isSeamlessEnabled
-            ? isUsingRefiner
-              ? SDXL_REFINER_SEAMLESS
-              : SEAMLESS
-            : isAutoVae
-              ? modelLoaderNodeId
-              : VAE_LOADER,
-          field: 'vae',
-        },
-        destination: {
-          node_id: SDXL_REFINER_INPAINT_CREATE_MASK,
-          field: 'vae',
-        },
-      });
-    }
-  }
-
   if (vae) {
     upsertMetadata(graph, { vae });
   }