From 476d6df1ca8cae0904fc6089ccfb115c5efe824b Mon Sep 17 00:00:00 2001
From: Terry Jia <terryjia88@gmail.com>
Date: Tue, 14 Oct 2025 20:42:36 -0400
Subject: [PATCH] fix mask editor bug under vueNodes (#5953)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary

fix mask editor issues on vueNodes

┆Issue is synchronized with this [Notion
page](https://www.notion.so/PR-5953-fix-mask-editor-bug-under-vueNodes-2856d73d3650810aa8a2e1a94c4d97a6)
by [Unito](https://www.unito.io)
---
 .../vueNodes/components/ImagePreview.vue      | 15 +++++++++++++
 src/scripts/app.ts                            | 11 ++++++++--
 src/stores/imagePreviewStore.ts               | 21 +++++++++++++++++++
 .../vueNodes/components/ImagePreview.spec.ts  |  3 +--
 4 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/src/renderer/extensions/vueNodes/components/ImagePreview.vue b/src/renderer/extensions/vueNodes/components/ImagePreview.vue
index aa7c27100..52f7c8a0b 100644
--- a/src/renderer/extensions/vueNodes/components/ImagePreview.vue
+++ b/src/renderer/extensions/vueNodes/components/ImagePreview.vue
@@ -32,6 +32,7 @@
       <!-- Main Image -->
       <img
         v-else
+        ref="currentImageEl"
         :src="currentImageUrl"
         :alt="imageAltText"
         class="block size-full object-contain"
@@ -118,6 +119,7 @@ import { computed, ref, watch } from 'vue'
 import { useI18n } from 'vue-i18n'
 
 import { downloadFile } from '@/base/common/downloadUtil'
+import { app } from '@/scripts/app'
 import { useCommandStore } from '@/stores/commandStore'
 import { useNodeOutputStore } from '@/stores/imagePreviewStore'
 
@@ -143,6 +145,8 @@ const actualDimensions = ref<string | null>(null)
 const imageError = ref(false)
 const isLoading = ref(false)
 
+const currentImageEl = ref<HTMLImageElement>()
+
 // Computed values
 const currentImageUrl = computed(() => props.imageUrls[currentIndex.value])
 const hasMultipleImages = computed(() => props.imageUrls.length > 1)
@@ -182,7 +186,18 @@ const handleImageError = () => {
   actualDimensions.value = null
 }
 
+// In vueNodes mode, we need to set them manually before opening the mask editor.
+const setupNodeForMaskEditor = () => {
+  if (!props.nodeId || !currentImageEl.value) return
+  const node = app.rootGraph?.getNodeById(props.nodeId)
+  if (!node) return
+  node.imageIndex = currentIndex.value
+  node.imgs = [currentImageEl.value]
+  app.canvas?.select(node)
+}
+
 const handleEditMask = () => {
+  setupNodeForMaskEditor()
   void commandStore.execute('Comfy.MaskEditor.OpenMaskEditor')
 }
 
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index c962fd1eb..043c76862 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -380,11 +380,15 @@ export class ComfyApp {
     const paintedIndex = selectedIndex + 1
     const combinedIndex = selectedIndex + 2
 
+    // for vueNodes mode
+    const images =
+      node.images ?? useNodeOutputStore().getNodeOutputs(node)?.images
+
     ComfyApp.clipspace = {
       widgets: widgets,
       imgs: imgs,
       original_imgs: orig_imgs,
-      images: node.images,
+      images: images,
       selectedIndex: selectedIndex,
       img_paste_mode: 'selected', // reset to default im_paste_mode state on copy action
       paintedIndex: paintedIndex,
@@ -411,7 +415,8 @@ export class ComfyApp {
           ComfyApp.clipspace.imgs[ComfyApp.clipspace.combinedIndex].src
       }
       if (ComfyApp.clipspace.imgs && node.imgs) {
-        if (node.images && ComfyApp.clipspace.images) {
+        // Update node.images even if it's initially undefined (vueNodes mode)
+        if (ComfyApp.clipspace.images) {
           if (ComfyApp.clipspace['img_paste_mode'] == 'selected') {
             node.images = [
               ComfyApp.clipspace.images[ComfyApp.clipspace['selectedIndex']]
@@ -513,6 +518,8 @@ export class ComfyApp {
       }
 
       app.graph.setDirtyCanvas(true)
+
+      useNodeOutputStore().updateNodeImages(node)
     }
   }
 
diff --git a/src/stores/imagePreviewStore.ts b/src/stores/imagePreviewStore.ts
index 848eb0ae7..300c532c9 100644
--- a/src/stores/imagePreviewStore.ts
+++ b/src/stores/imagePreviewStore.ts
@@ -331,6 +331,26 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => {
     nodeOutputs.value = outputs
   }
 
+  function updateNodeImages(node: LGraphNode) {
+    if (!node.images?.length) return
+
+    const nodeLocatorId = nodeIdToNodeLocatorId(node.id)
+
+    if (nodeLocatorId) {
+      const existingOutputs = app.nodeOutputs[nodeLocatorId]
+
+      if (existingOutputs) {
+        const updatedOutputs = {
+          ...existingOutputs,
+          images: node.images
+        }
+
+        app.nodeOutputs[nodeLocatorId] = updatedOutputs
+        nodeOutputs.value[nodeLocatorId] = updatedOutputs
+      }
+    }
+  }
+
   function resetAllOutputsAndPreviews() {
     app.nodeOutputs = {}
     nodeOutputs.value = {}
@@ -349,6 +369,7 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => {
     setNodeOutputsByExecutionId,
     setNodePreviewsByExecutionId,
     setNodePreviewsByNodeId,
+    updateNodeImages,
 
     // Cleanup
     revokePreviewsByExecutionId,
diff --git a/tests-ui/tests/renderer/extensions/vueNodes/components/ImagePreview.spec.ts b/tests-ui/tests/renderer/extensions/vueNodes/components/ImagePreview.spec.ts
index ccf61ebcf..5b9d658df 100644
--- a/tests-ui/tests/renderer/extensions/vueNodes/components/ImagePreview.spec.ts
+++ b/tests-ui/tests/renderer/extensions/vueNodes/components/ImagePreview.spec.ts
@@ -203,9 +203,8 @@ describe('ImagePreview', () => {
     await navigationDots[1].trigger('click')
     await nextTick()
 
-    // After clicking, component shows loading state (Skeleton), not img
+    // After clicking, component shows loading state (Skeleton)
     expect(wrapper.find('skeleton-stub').exists()).toBe(true)
-    expect(wrapper.find('img').exists()).toBe(false)
 
     // Simulate image load event to clear loading state
     const component = wrapper.vm as any