From 3b9d4533f9fc048c5009d7469649525601180980 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Sat, 27 May 2023 22:39:07 -0500
Subject: [PATCH] Preserve mask when sending output into image upload

---
 src/lib/components/ComfyQueue.svelte          |  4 +-
 src/lib/components/MaskCanvas.svelte          | 66 ++++++++++++++-----
 src/lib/nodes/ComfyPickImageNode.ts           |  4 +-
 src/lib/nodes/widgets/ComfyImageUploadNode.ts |  1 +
 src/lib/widgets/ImageUploadWidget.svelte      | 12 ++--
 5 files changed, 59 insertions(+), 28 deletions(-)

diff --git a/src/lib/components/ComfyQueue.svelte b/src/lib/components/ComfyQueue.svelte
index 9d8bd75..8aa3cdf 100644
--- a/src/lib/components/ComfyQueue.svelte
+++ b/src/lib/components/ComfyQueue.svelte
@@ -146,10 +146,10 @@
      if (entry.extraData?.workflowTitle != null) {
          message = `${entry.extraData.workflowTitle}`
      }
+
      if (subgraphs) {
          const subgraphsString = subgraphs.join(', ')
-         if (subgraphsString.length > 0)
-             message += ` (${subgraphsString})`
+         message += ` (${subgraphsString})`
      }
 
      let submessage = `Nodes: ${Object.keys(entry.prompt).length}`
diff --git a/src/lib/components/MaskCanvas.svelte b/src/lib/components/MaskCanvas.svelte
index 6ad0cb6..4463d7a 100644
--- a/src/lib/components/MaskCanvas.svelte
+++ b/src/lib/components/MaskCanvas.svelte
@@ -2,7 +2,19 @@
  export type MaskCanvasData = {
      hasMask: boolean,
      maskCanvas: HTMLCanvasElement | null,
+     curLineGroup: LineGroup,
+     redoCurLines: LineGroup,
  }
+ export type LinePoint = {
+     x: number,
+     y: number
+ }
+ export interface Line {
+     size?: number,
+     points: LinePoint[]
+ }
+ export type LineGroup = Line[];
+
 </script>
 
 <script lang="ts">
@@ -16,6 +28,7 @@
  const dispatch = createEventDispatcher<{
      change: MaskCanvasData;
      release: MaskCanvasData;
+     loaded: MaskCanvasData
  }>();
 
  let canvasCursor: string | undefined = undefined;
@@ -87,8 +100,8 @@
      isImageLoaded = false;
      original = null;
      renders = []
-     curLineGroup = [];
-     redoCurLines = []
+     // curLineGroup = [];
+     // redoCurLines = []
      imageWidth = 512;
      imageHeight = 512;
      scale = 1.0;
@@ -96,12 +109,14 @@
  }
 
  let loadedFileURL: string | null = null
+ let dispatchLoaded: boolean = false;
  $: if (fileURL !== loadedFileURL) {
      clearState();
      if (fileURL) {
          loadImage(fileURL).then(i => {
              original = i;
              isImageLoaded = true;
+             dispatchLoaded = true;
          })
                            .catch(i => {
                                isImageLoaded = false;
@@ -119,6 +134,28 @@
      [imageWidth, imageHeight] = getCurrentWidthAndHeight(isImageLoaded, original)
      scale = initScale(imageWidth, imageHeight)
      initImagePos()
+
+     // in case mask strokes were preserved after new image load
+     // (use case: sending an inpainted image back while reusing the same mask)
+     tick().then(() => {
+         loaded();
+     })
+ }
+
+ function loaded() {
+     if (!dispatchLoaded)
+         return;
+
+     dispatchLoaded = false;
+     redrawCurLines()
+     hasMask = curLineGroup.length > 0;
+     console.warn("[MaskCanvas] LOADED", maskCanvas, hasMask)
+     dispatch("loaded", {
+         hasMask,
+         maskCanvas,
+         curLineGroup,
+         redoCurLines
+     })
  }
 
  $: hasMask = curLineGroup.length > 0;
@@ -143,21 +180,13 @@
      minScale = scale;
  }
 
- type LinePoint = {
-     x: number,
-     y: number
- }
- interface Line {
-     size?: number,
-     points: LinePoint[]
- }
- type LineGroup = Line[];
-
  function drawOnCurrentRender(lineGroup: LineGroup) {
      draw(lineGroup)
      dispatch("change", {
          hasMask,
-         maskCanvas
+         maskCanvas,
+         curLineGroup,
+         redoCurLines
      })
  }
 
@@ -214,6 +243,10 @@
      })
  }
 
+ function redrawCurLines() {
+     drawOnCurrentRender(curLineGroup || [])
+ }
+
  $: if (canvas && original) {
      console.warn("INITCANVAS", imageWidth, imageHeight, original.src)
      maskCanvas = document.createElement("canvas");
@@ -222,7 +255,7 @@
      maskCanvas.height = imageHeight;
      canvas.width = imageWidth;
      canvas.height = imageHeight;
-     drawOnCurrentRender([])
+     redrawCurLines() // no react on curLineGroup
  }
 
  function getCurrentWidthAndHeight(isImageLoaded: boolean, original: HTMLImageElement | null) {
@@ -421,11 +454,12 @@
          return;
 
      isDrawing = false;
-     dispatch("release", { hasMask, maskCanvas })
+     dispatch("release", { hasMask, maskCanvas, curLineGroup, redoCurLines })
  }
 
  function dispatchRelease() {
-     dispatch("release", { hasMask, maskCanvas })
+     updateMaskImage()
+     dispatch("release", { hasMask, maskCanvas, curLineGroup, redoCurLines })
  }
 </script>
 
diff --git a/src/lib/nodes/ComfyPickImageNode.ts b/src/lib/nodes/ComfyPickImageNode.ts
index c314a79..8404463 100644
--- a/src/lib/nodes/ComfyPickImageNode.ts
+++ b/src/lib/nodes/ComfyPickImageNode.ts
@@ -99,9 +99,9 @@ export default class ComfyPickImageNode extends ComfyGraphNode {
     override onExecute() {
         const data = this.getInputData(0)
         let index = this.getInputData(1);
-        if (this.properties.imageTagFilter && Array.isArray(data))
+        if (this.properties.imageTagFilter != "" && Array.isArray(data))
             index = data.findIndex(i => i.tags?.includes(this.properties.imageTagFilter))
-        else
+        else if (index == null)
             index = 0;
         this.setValue(data, index);
 
diff --git a/src/lib/nodes/widgets/ComfyImageUploadNode.ts b/src/lib/nodes/widgets/ComfyImageUploadNode.ts
index eee73b4..1782b18 100644
--- a/src/lib/nodes/widgets/ComfyImageUploadNode.ts
+++ b/src/lib/nodes/widgets/ComfyImageUploadNode.ts
@@ -6,6 +6,7 @@ import ImageUploadWidget from "$lib/widgets/ImageUploadWidget.svelte";
 import type { ComfyWidgetProperties } from "./ComfyWidgetNode";
 import ComfyWidgetNode from "./ComfyWidgetNode";
 import { get, writable, type Writable } from "svelte/store";
+import { type LineGroup } from "$lib/components/MaskCanvas.svelte"
 
 export interface ComfyImageUploadNodeProperties extends ComfyWidgetProperties {
     maskCount: number
diff --git a/src/lib/widgets/ImageUploadWidget.svelte b/src/lib/widgets/ImageUploadWidget.svelte
index 2bc6c6a..b06bc3a 100644
--- a/src/lib/widgets/ImageUploadWidget.svelte
+++ b/src/lib/widgets/ImageUploadWidget.svelte
@@ -5,8 +5,6 @@
  import Row from "$lib/components/gradio/app/Row.svelte";
  import { writable, type Writable } from "svelte/store";
  import { Button } from "@gradio/button";
-
- import ImageUpload from "$lib/components/ImageUpload.svelte";
  import {
      type ComfyBoxImageMetadata,
      comfyFileToComfyBoxMetadata,
@@ -16,13 +14,12 @@
      convertComfyOutputToComfyURL,
      batchUploadBlobsToComfyUI,
      canvasToBlob,
-
 	 basename
-
  } from "$lib/utils";
+ import ImageUpload from "$lib/components/ImageUpload.svelte";
  import notify from "$lib/notify";
  import { ImageViewer } from "$lib/ImageViewer";
- import MaskCanvas, { type MaskCanvasData } from "$lib/components/MaskCanvas.svelte";
+ import MaskCanvas, { type LineGroup, type MaskCanvasData } from "$lib/components/MaskCanvas.svelte";
  import type { ComfyImageUploadNode } from "$lib/nodes/widgets";
 	import { tick } from "svelte";
 
@@ -65,7 +62,7 @@
 
  async function onMaskReleased(e: CustomEvent<MaskCanvasData>) {
      const data = e.detail;
-     if (data.maskCanvas && data.hasMask) {
+     if (data.maskCanvas != null && data.hasMask) {
          await saveMask(data.maskCanvas)
      }
  }
@@ -203,7 +200,6 @@
 
  $: canEdit = status === "empty" || status === "uploaded";
 
-
  function onChange(e: CustomEvent<ComfyImageLocation[]>) {
  }
 </script>
@@ -229,7 +225,7 @@
             {#if _value && canMask}
                 {@const comfyURL = convertComfyOutputToComfyURL(_value[0])}
                 <div class="mask-canvas-wrapper" style:display={editMask ? "block" : "none"}>
-                    <MaskCanvas bind:this={maskCanvasComp} fileURL={comfyURL} on:release={onMaskReleased} />
+                    <MaskCanvas bind:this={maskCanvasComp} fileURL={comfyURL} on:release={onMaskReleased} on:loaded={onMaskReleased} />
                 </div>
             {/if}
             <div style:display={(canMask && editMask) ? "none" : "block"}>