From 7053347559263c6bb209332cc240b2110847b666 Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Thu, 27 Jul 2023 07:09:51 +1200
Subject: [PATCH] fix: Metadata Not Being Saved

---
 .../buildCanvasImageToImageGraph.ts           | 16 +++++++
 .../graphBuilders/buildCanvasInpaintGraph.ts  | 42 ++++++++++++++++++-
 .../buildCanvasTextToImageGraph.ts            | 16 +++++++
 .../buildLinearImageToImageGraph.ts           | 16 +++++++
 .../buildLinearSDXLImageToImageGraph.ts       | 27 +++++++-----
 .../buildLinearSDXLTextToImageGraph.ts        | 27 +++++++-----
 .../buildLinearTextToImageGraph.ts            | 16 +++++++
 7 files changed, 136 insertions(+), 24 deletions(-)

diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasImageToImageGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasImageToImageGraph.ts
index 81e99e0dd0..2981c090d8 100644
--- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasImageToImageGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasImageToImageGraph.ts
@@ -335,5 +335,21 @@ export const buildCanvasImageToImageGraph = (
     addWatermarkerToGraph(state, graph);
   }
 
+  if (
+    !state.system.shouldUseNSFWChecker &&
+    !state.system.shouldUseWatermarker
+  ) {
+    graph.edges.push({
+      source: {
+        node_id: METADATA_ACCUMULATOR,
+        field: 'metadata',
+      },
+      destination: {
+        node_id: LATENTS_TO_IMAGE,
+        field: 'metadata',
+      },
+    });
+  }
+
   return graph;
 };
diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasInpaintGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasInpaintGraph.ts
index 4154c1b5eb..8462b32e63 100644
--- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasInpaintGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasInpaintGraph.ts
@@ -8,20 +8,21 @@ import {
   RangeOfSizeInvocation,
 } from 'services/api/types';
 import { addLoRAsToGraph } from './addLoRAsToGraph';
+import { addNSFWCheckerToGraph } from './addNSFWCheckerToGraph';
 import { addVAEToGraph } from './addVAEToGraph';
+import { addWatermarkerToGraph } from './addWatermarkerToGraph';
 import {
   CLIP_SKIP,
   INPAINT,
   INPAINT_GRAPH,
   ITERATE,
   MAIN_MODEL_LOADER,
+  METADATA_ACCUMULATOR,
   NEGATIVE_CONDITIONING,
   POSITIVE_CONDITIONING,
   RANDOM_INT,
   RANGE_OF_SIZE,
 } from './constants';
-import { addNSFWCheckerToGraph } from './addNSFWCheckerToGraph';
-import { addWatermarkerToGraph } from './addWatermarkerToGraph';
 
 /**
  * Builds the Canvas tab's Inpaint graph.
@@ -251,6 +252,27 @@ export const buildCanvasInpaintGraph = (
     (graph.nodes[RANGE_OF_SIZE] as RangeOfSizeInvocation).start = seed;
   }
 
+  // add metadata accumulator, which is only mostly populated - some fields are added later
+  graph.nodes[METADATA_ACCUMULATOR] = {
+    id: METADATA_ACCUMULATOR,
+    type: 'metadata_accumulator',
+    generation_mode: 'txt2img',
+    cfg_scale,
+    height,
+    width,
+    positive_prompt: '', // set in addDynamicPromptsToGraph
+    negative_prompt: negativePrompt,
+    model,
+    seed: 0, // set in addDynamicPromptsToGraph
+    steps,
+    rand_device: 'cpu',
+    scheduler,
+    vae: undefined, // option; set in addVAEToGraph
+    controlnets: [], // populated in addControlNetToLinearGraph
+    loras: [], // populated in addLoRAsToGraph
+    clip_skip: clipSkip,
+  };
+
   // NSFW & watermark - must be last thing added to graph
   if (state.system.shouldUseNSFWChecker) {
     // must add before watermarker!
@@ -262,5 +284,21 @@ export const buildCanvasInpaintGraph = (
     addWatermarkerToGraph(state, graph, INPAINT);
   }
 
+  if (
+    !state.system.shouldUseNSFWChecker &&
+    !state.system.shouldUseWatermarker
+  ) {
+    graph.edges.push({
+      source: {
+        node_id: METADATA_ACCUMULATOR,
+        field: 'metadata',
+      },
+      destination: {
+        node_id: INPAINT,
+        field: 'metadata',
+      },
+    });
+  }
+
   return graph;
 };
diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasTextToImageGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasTextToImageGraph.ts
index 597a643367..55e11e4f37 100644
--- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasTextToImageGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasTextToImageGraph.ts
@@ -240,5 +240,21 @@ export const buildCanvasTextToImageGraph = (
     addWatermarkerToGraph(state, graph);
   }
 
+  if (
+    !state.system.shouldUseNSFWChecker &&
+    !state.system.shouldUseWatermarker
+  ) {
+    graph.edges.push({
+      source: {
+        node_id: METADATA_ACCUMULATOR,
+        field: 'metadata',
+      },
+      destination: {
+        node_id: LATENTS_TO_IMAGE,
+        field: 'metadata',
+      },
+    });
+  }
+
   return graph;
 };
diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearImageToImageGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearImageToImageGraph.ts
index d78e2e1356..e726a02a5c 100644
--- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearImageToImageGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearImageToImageGraph.ts
@@ -345,5 +345,21 @@ export const buildLinearImageToImageGraph = (
     addWatermarkerToGraph(state, graph);
   }
 
+  if (
+    !state.system.shouldUseNSFWChecker &&
+    !state.system.shouldUseWatermarker
+  ) {
+    graph.edges.push({
+      source: {
+        node_id: METADATA_ACCUMULATOR,
+        field: 'metadata',
+      },
+      destination: {
+        node_id: LATENTS_TO_IMAGE,
+        field: 'metadata',
+      },
+    });
+  }
+
   return graph;
 };
diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearSDXLImageToImageGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearSDXLImageToImageGraph.ts
index a260dbc467..3697b96bea 100644
--- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearSDXLImageToImageGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearSDXLImageToImageGraph.ts
@@ -353,17 +353,6 @@ export const buildLinearSDXLImageToImageGraph = (
     negative_style_prompt: negativeStylePrompt,
   };
 
-  graph.edges.push({
-    source: {
-      node_id: METADATA_ACCUMULATOR,
-      field: 'metadata',
-    },
-    destination: {
-      node_id: LATENTS_TO_IMAGE,
-      field: 'metadata',
-    },
-  });
-
   // Add Refiner if enabled
   if (shouldUseSDXLRefiner) {
     addSDXLRefinerToGraph(state, graph, SDXL_LATENTS_TO_LATENTS);
@@ -383,5 +372,21 @@ export const buildLinearSDXLImageToImageGraph = (
     addWatermarkerToGraph(state, graph);
   }
 
+  if (
+    !state.system.shouldUseNSFWChecker &&
+    !state.system.shouldUseWatermarker
+  ) {
+    graph.edges.push({
+      source: {
+        node_id: METADATA_ACCUMULATOR,
+        field: 'metadata',
+      },
+      destination: {
+        node_id: LATENTS_TO_IMAGE,
+        field: 'metadata',
+      },
+    });
+  }
+
   return graph;
 };
diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearSDXLTextToImageGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearSDXLTextToImageGraph.ts
index c10e7831d3..1669de5341 100644
--- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearSDXLTextToImageGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearSDXLTextToImageGraph.ts
@@ -235,17 +235,6 @@ export const buildLinearSDXLTextToImageGraph = (
     negative_style_prompt: negativeStylePrompt,
   };
 
-  graph.edges.push({
-    source: {
-      node_id: METADATA_ACCUMULATOR,
-      field: 'metadata',
-    },
-    destination: {
-      node_id: LATENTS_TO_IMAGE,
-      field: 'metadata',
-    },
-  });
-
   // Add Refiner if enabled
   if (shouldUseSDXLRefiner) {
     addSDXLRefinerToGraph(state, graph, SDXL_TEXT_TO_LATENTS);
@@ -265,5 +254,21 @@ export const buildLinearSDXLTextToImageGraph = (
     addWatermarkerToGraph(state, graph);
   }
 
+  if (
+    !state.system.shouldUseNSFWChecker &&
+    !state.system.shouldUseWatermarker
+  ) {
+    graph.edges.push({
+      source: {
+        node_id: METADATA_ACCUMULATOR,
+        field: 'metadata',
+      },
+      destination: {
+        node_id: LATENTS_TO_IMAGE,
+        field: 'metadata',
+      },
+    });
+  }
+
   return graph;
 };
diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearTextToImageGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearTextToImageGraph.ts
index 94f35feba0..08aa355d4d 100644
--- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearTextToImageGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearTextToImageGraph.ts
@@ -229,5 +229,21 @@ export const buildLinearTextToImageGraph = (
     addWatermarkerToGraph(state, graph);
   }
 
+  if (
+    !state.system.shouldUseNSFWChecker &&
+    !state.system.shouldUseWatermarker
+  ) {
+    graph.edges.push({
+      source: {
+        node_id: METADATA_ACCUMULATOR,
+        field: 'metadata',
+      },
+      destination: {
+        node_id: LATENTS_TO_IMAGE,
+        field: 'metadata',
+      },
+    });
+  }
+
   return graph;
 };