From f2c4e567e4b36f29938f39bae3891d3f071e8f5e Mon Sep 17 00:00:00 2001
From: Yiximail <yiximail@gmail.com>
Date: Tue, 20 May 2025 16:23:39 +0800
Subject: [PATCH] Add a button to selection toolbox to open mask editor (#3603)

Co-authored-by: bymyself <cbyrne@comfy.org>
---
 src/components/graph/SelectionToolbox.vue     |  13 +--
 .../selectionToolbox/MaskEditorButton.vue     |  35 ++++++
 src/extensions/core/maskeditor.ts             | 100 +++++++++++-------
 src/locales/en/commands.json                  |   3 +
 src/locales/es/commands.json                  |   3 +
 src/locales/fr/commands.json                  |   3 +
 src/locales/ja/commands.json                  |   3 +
 src/locales/ko/commands.json                  |   3 +
 src/locales/ru/commands.json                  |   3 +
 src/locales/zh/commands.json                  |   3 +
 10 files changed, 124 insertions(+), 45 deletions(-)
 create mode 100644 src/components/graph/selectionToolbox/MaskEditorButton.vue

diff --git a/src/components/graph/SelectionToolbox.vue b/src/components/graph/SelectionToolbox.vue
index d6d6e66cb..4ff1eee92 100644
--- a/src/components/graph/SelectionToolbox.vue
+++ b/src/components/graph/SelectionToolbox.vue
@@ -10,6 +10,7 @@
     <ColorPickerButton />
     <BypassButton />
     <PinButton />
+    <MaskEditorButton />
     <DeleteButton />
     <RefreshButton />
     <ExtensionCommandButton
@@ -24,18 +25,18 @@
 import Panel from 'primevue/panel'
 import { computed } from 'vue'
 
+import BypassButton from '@/components/graph/selectionToolbox/BypassButton.vue'
 import ColorPickerButton from '@/components/graph/selectionToolbox/ColorPickerButton.vue'
+import DeleteButton from '@/components/graph/selectionToolbox/DeleteButton.vue'
 import ExecuteButton from '@/components/graph/selectionToolbox/ExecuteButton.vue'
+import ExtensionCommandButton from '@/components/graph/selectionToolbox/ExtensionCommandButton.vue'
+import MaskEditorButton from '@/components/graph/selectionToolbox/MaskEditorButton.vue'
+import PinButton from '@/components/graph/selectionToolbox/PinButton.vue'
+import RefreshButton from '@/components/graph/selectionToolbox/RefreshButton.vue'
 import { useExtensionService } from '@/services/extensionService'
 import { type ComfyCommandImpl, useCommandStore } from '@/stores/commandStore'
 import { useCanvasStore } from '@/stores/graphStore'
 
-import BypassButton from './selectionToolbox/BypassButton.vue'
-import DeleteButton from './selectionToolbox/DeleteButton.vue'
-import ExtensionCommandButton from './selectionToolbox/ExtensionCommandButton.vue'
-import PinButton from './selectionToolbox/PinButton.vue'
-import RefreshButton from './selectionToolbox/RefreshButton.vue'
-
 const commandStore = useCommandStore()
 const canvasStore = useCanvasStore()
 const extensionService = useExtensionService()
diff --git a/src/components/graph/selectionToolbox/MaskEditorButton.vue b/src/components/graph/selectionToolbox/MaskEditorButton.vue
new file mode 100644
index 000000000..9336e58bd
--- /dev/null
+++ b/src/components/graph/selectionToolbox/MaskEditorButton.vue
@@ -0,0 +1,35 @@
+<template>
+  <Button
+    v-show="isSingleImageNode"
+    v-tooltip.top="{
+      value: t('commands.Comfy_MaskEditor_OpenMaskEditor.label'),
+      showDelay: 1000
+    }"
+    severity="secondary"
+    text
+    icon="pi pi-pencil"
+    @click="openMaskEditor"
+  />
+</template>
+
+<script setup lang="ts">
+import Button from 'primevue/button'
+import { computed } from 'vue'
+
+import { t } from '@/i18n'
+import { useCommandStore } from '@/stores/commandStore'
+import { useCanvasStore } from '@/stores/graphStore'
+import { isImageNode, isLGraphNode } from '@/utils/litegraphUtil'
+
+const commandStore = useCommandStore()
+const canvasStore = useCanvasStore()
+
+const isSingleImageNode = computed(() => {
+  const nodes = canvasStore.selectedItems.filter(isLGraphNode)
+  return nodes.length === 1 && nodes.some(isImageNode)
+})
+
+const openMaskEditor = () => {
+  void commandStore.execute('Comfy.MaskEditor.OpenMaskEditor')
+}
+</script>
diff --git a/src/extensions/core/maskeditor.ts b/src/extensions/core/maskeditor.ts
index 39b62308b..d980e714a 100644
--- a/src/extensions/core/maskeditor.ts
+++ b/src/extensions/core/maskeditor.ts
@@ -4912,6 +4912,45 @@ class KeyboardManager {
   }
 }
 
+// Function to open the mask editor
+function openMaskEditor(): void {
+  const useNewEditor = app.extensionManager.setting.get(
+    'Comfy.MaskEditor.UseNewEditor'
+  )
+  if (useNewEditor) {
+    const dlg = MaskEditorDialog.getInstance() as any
+    if (dlg?.isOpened && !dlg.isOpened()) {
+      dlg.show()
+    }
+  } else {
+    const dlg = MaskEditorDialogOld.getInstance() as any
+    if (dlg?.isOpened && !dlg.isOpened()) {
+      dlg.show()
+    }
+  }
+}
+
+// Check if the dialog is already opened
+function isOpened(): boolean {
+  const useNewEditor = app.extensionManager.setting.get(
+    'Comfy.MaskEditor.UseNewEditor'
+  )
+  if (useNewEditor) {
+    return MaskEditorDialog.instance?.isOpened?.() ?? false
+  } else {
+    return (MaskEditorDialogOld.instance as any)?.isOpened?.() ?? false
+  }
+}
+
+// Ensure boolean return type for context predicate
+const context_predicate = (): boolean => {
+  return !!(
+    ComfyApp.clipspace &&
+    ComfyApp.clipspace.imgs &&
+    ComfyApp.clipspace.imgs.length > 0
+  )
+}
+
 app.registerExtension({
   name: 'Comfy.MaskEditor',
   settings: [
@@ -4951,50 +4990,33 @@ app.registerExtension({
       experimental: true
     }
   ],
-  init(app) {
-    // Create function before assignment
-    function openMaskEditor(): void {
-      const useNewEditor = app.extensionManager.setting.get(
-        'Comfy.MaskEditor.UseNewEditor'
-      )
-      if (useNewEditor) {
-        const dlg = MaskEditorDialog.getInstance() as any
-        if (dlg?.isOpened && !dlg.isOpened()) {
-          dlg.show()
-        }
-      } else {
-        const dlg = MaskEditorDialogOld.getInstance() as any
-        if (dlg?.isOpened && !dlg.isOpened()) {
-          dlg.show()
-        }
+  commands: [
+    {
+      id: 'Comfy.MaskEditor.OpenMaskEditor',
+      icon: 'pi pi-pencil',
+      label: 'Open Mask Editor for Selected Node',
+      function: () => {
+        const selectedNodes = app.canvas.selected_nodes
+        if (!selectedNodes || Object.keys(selectedNodes).length !== 1) return
+
+        const selectedNode = selectedNodes[Object.keys(selectedNodes)[0]]
+        if (
+          !selectedNode.imgs?.length &&
+          selectedNode.previewMediaType !== 'image'
+        )
+          return
+
+        ComfyApp.copyToClipspace(selectedNode)
+        // @ts-expect-error clipspace_return_node is an extension property added at runtime
+        ComfyApp.clipspace_return_node = selectedNode
+        openMaskEditor()
       }
     }
-
-    // Check if the dialog is already opened
-    function isOpened(): boolean {
-      const useNewEditor = app.extensionManager.setting.get(
-        'Comfy.MaskEditor.UseNewEditor'
-      )
-      if (useNewEditor) {
-        return MaskEditorDialog.instance?.isOpened?.() ?? false
-      } else {
-        return (MaskEditorDialogOld.instance as any)?.isOpened?.() ?? false
-      }
-    }
-
-    // Assign the created function
+  ],
+  init() {
     ComfyApp.open_maskeditor = openMaskEditor
     ComfyApp.maskeditor_is_opended = isOpened
 
-    // Ensure boolean return type
-    const context_predicate = (): boolean => {
-      return !!(
-        ComfyApp.clipspace &&
-        ComfyApp.clipspace.imgs &&
-        ComfyApp.clipspace.imgs.length > 0
-      )
-    }
-
     ClipspaceDialog.registerButton(
       'MaskEditor',
       context_predicate,
diff --git a/src/locales/en/commands.json b/src/locales/en/commands.json
index e707fb93f..6083db1a6 100644
--- a/src/locales/en/commands.json
+++ b/src/locales/en/commands.json
@@ -71,6 +71,9 @@
   "Comfy_Canvas_ToggleSelectedNodes_Pin": {
     "label": "Pin/Unpin Selected Nodes"
   },
+  "Comfy_MaskEditor_OpenMaskEditor": {
+    "label": "Open Mask Editor for Selected Node"
+  },
   "Comfy_Canvas_ZoomIn": {
     "label": "Zoom In"
   },
diff --git a/src/locales/es/commands.json b/src/locales/es/commands.json
index 1294be59f..0b2352d0f 100644
--- a/src/locales/es/commands.json
+++ b/src/locales/es/commands.json
@@ -71,6 +71,9 @@
   "Comfy_Canvas_ToggleSelected_Pin": {
     "label": "Anclar/Desanclar elementos seleccionados"
   },
+  "Comfy_MaskEditor_OpenMaskEditor": {
+    "label": "Abrir editor de máscara para el nodo seleccionado"
+  },
   "Comfy_Canvas_ZoomIn": {
     "label": "Acercar"
   },
diff --git a/src/locales/fr/commands.json b/src/locales/fr/commands.json
index a38391210..8e335d6be 100644
--- a/src/locales/fr/commands.json
+++ b/src/locales/fr/commands.json
@@ -71,6 +71,9 @@
   "Comfy_Canvas_ToggleSelected_Pin": {
     "label": "Épingler/Désépingler les éléments sélectionnés"
   },
+  "Comfy_MaskEditor_OpenMaskEditor": {
+    "label": "Ouvrir l'éditeur de masque pour le nœud sélectionné"
+  },
   "Comfy_Canvas_ZoomIn": {
     "label": "Zoom avant"
   },
diff --git a/src/locales/ja/commands.json b/src/locales/ja/commands.json
index 75dfd93e3..0b3b00130 100644
--- a/src/locales/ja/commands.json
+++ b/src/locales/ja/commands.json
@@ -71,6 +71,9 @@
   "Comfy_Canvas_ToggleSelected_Pin": {
     "label": "選択したアイテムのピン留め/ピン留め解除"
   },
+  "Comfy_MaskEditor_OpenMaskEditor": {
+    "label": "選択したノードのマスクエディタを開く"
+  },
   "Comfy_Canvas_ZoomIn": {
     "label": "ズームイン"
   },
diff --git a/src/locales/ko/commands.json b/src/locales/ko/commands.json
index 624f0db1b..5b9945a85 100644
--- a/src/locales/ko/commands.json
+++ b/src/locales/ko/commands.json
@@ -71,6 +71,9 @@
   "Comfy_Canvas_ToggleSelected_Pin": {
     "label": "선택한 항목 고정/고정 해제"
   },
+  "Comfy_MaskEditor_OpenMaskEditor": {
+    "label": "선택한 노드 마스크 편집기 열기"
+  },
   "Comfy_Canvas_ZoomIn": {
     "label": "확대"
   },
diff --git a/src/locales/ru/commands.json b/src/locales/ru/commands.json
index e692760fc..5ee814024 100644
--- a/src/locales/ru/commands.json
+++ b/src/locales/ru/commands.json
@@ -71,6 +71,9 @@
   "Comfy_Canvas_ToggleSelected_Pin": {
     "label": "Закрепить/Открепить выбранных нод"
   },
+  "Comfy_MaskEditor_OpenMaskEditor": {
+    "label": "Открыть редактор масок для выбранной ноды"
+  },
   "Comfy_Canvas_ZoomIn": {
     "label": "Увеличить"
   },
diff --git a/src/locales/zh/commands.json b/src/locales/zh/commands.json
index 29fef86d1..e5f1a232e 100644
--- a/src/locales/zh/commands.json
+++ b/src/locales/zh/commands.json
@@ -71,6 +71,9 @@
   "Comfy_Canvas_ToggleSelected_Pin": {
     "label": "固定/取消固定选中项"
   },
+  "Comfy_MaskEditor_OpenMaskEditor": {
+    "label": "打开选中节点的遮罩编辑器"
+  },
   "Comfy_Canvas_ZoomIn": {
     "label": "放大"
   },