[Feature] Enhanced MaskEditor to an Image Canvas (#4361)

Co-authored-by: duckcomfy <a@a.a>

src/extensions/core/maskEditorLayerFilenames.ts

@@ -0,0 +1,29 @@
+export interface ImageLayerFilenames {
+  maskedImage: string
+  paint: string
+  paintedImage: string
+  paintedMaskedImage: string
+}
+
+const paintedMaskedImagePrefix = 'clipspace-painted-masked-'
+
+export const imageLayerFilenamesByTimestamp = (
+  timestamp: number
+): ImageLayerFilenames => ({
+  maskedImage: `clipspace-mask-${timestamp}.png`,
+  paint: `clipspace-paint-${timestamp}.png`,
+  paintedImage: `clipspace-painted-${timestamp}.png`,
+  paintedMaskedImage: `${paintedMaskedImagePrefix}${timestamp}.png`
+})
+
+export const imageLayerFilenamesIfApplicable = (
+  inputImageFilename: string
+): ImageLayerFilenames | undefined => {
+  const isPaintedMaskedImageFilename = inputImageFilename.startsWith(
+    paintedMaskedImagePrefix
+  )
+  if (!isPaintedMaskedImageFilename) return undefined
+  const suffix = inputImageFilename.slice(paintedMaskedImagePrefix.length)
+  const timestamp = parseInt(suffix.split('.')[0], 10)
+  return imageLayerFilenamesByTimestamp(timestamp)
+}

src/scripts/app.ts

@@ -116,6 +116,8 @@ type Clipspace = {
   images?: any[] | null
   selectedIndex: number
   img_paste_mode: string
+  paintedIndex: number
+  combinedIndex: number
 }
 
 export class ComfyApp {
@@ -357,13 +359,18 @@ export class ComfyApp {
       selectedIndex = node.imageIndex
     }
 
+    const paintedIndex = selectedIndex + 1
+    const combinedIndex = selectedIndex + 2
+
     ComfyApp.clipspace = {
       widgets: widgets,
       imgs: imgs,
       original_imgs: orig_imgs,
       images: node.images,
       selectedIndex: selectedIndex,
-      img_paste_mode: 'selected' // reset to default im_paste_mode state on copy action
+      img_paste_mode: 'selected', // reset to default im_paste_mode state on copy action
+      paintedIndex: paintedIndex,
+      combinedIndex: combinedIndex
     }
 
     ComfyApp.clipspace_return_node = null
@@ -376,6 +383,8 @@ export class ComfyApp {
   static pasteFromClipspace(node: LGraphNode) {
     if (ComfyApp.clipspace) {
       // image paste
+      const combinedImgSrc =
+        ComfyApp.clipspace.imgs?.[ComfyApp.clipspace.combinedIndex].src
       if (ComfyApp.clipspace.imgs && node.imgs) {
         if (node.images && ComfyApp.clipspace.images) {
           if (ComfyApp.clipspace['img_paste_mode'] == 'selected') {
@@ -409,6 +418,28 @@ export class ComfyApp {
         }
       }
 
+      // Paste the RGB canvas if paintedindex exists
+      if (
+        ComfyApp.clipspace.imgs?.[ComfyApp.clipspace.paintedIndex] &&
+        node.imgs
+      ) {
+        const paintedImg = new Image()
+        paintedImg.src =
+          ComfyApp.clipspace.imgs[ComfyApp.clipspace.paintedIndex].src
+        node.imgs.push(paintedImg) // Add the RGB canvas to the node's images
+      }
+
+      // Store only combined image inside the node if it exists
+      if (
+        ComfyApp.clipspace.imgs?.[ComfyApp.clipspace.combinedIndex] &&
+        node.imgs &&
+        combinedImgSrc
+      ) {
+        const combinedImg = new Image()
+        combinedImg.src = combinedImgSrc
+        node.imgs = [combinedImg]
+      }
+
       if (node.widgets) {
         if (ComfyApp.clipspace.images) {
           const clip_image =

src/services/litegraphService.ts

@@ -787,7 +787,7 @@ export const useLitegraphService = () => {
 
         if (isImageNode(this)) {
           options.push({
-            content: 'Open in MaskEditor',
+            content: 'Open in MaskEditor | Image Canvas',
             callback: () => {
               ComfyApp.copyToClipspace(this)
               // @ts-expect-error fixme ts strict error

src/utils/colorUtil.ts

@@ -40,7 +40,7 @@ function rgbToHsl({ r, g, b }: RGB): HSL {
   return { h, s, l }
 }
 
-function hexToRgb(hex: string): RGB {
+export function hexToRgb(hex: string): RGB {
   let r = 0,
     g = 0,
     b = 0

tests-ui/tests/maskeditor.test.ts

@@ -0,0 +1,19 @@
+import { describe, expect, it } from 'vitest'
+
+import { imageLayerFilenamesIfApplicable } from '@/extensions/core/maskEditorLayerFilenames'
+
+describe('imageLayerFilenamesIfApplicable', () => {
+  // In case the naming scheme changes, this test will ensure CI fails if developers forget to support the old naming scheme. (Causing MaskEditor to lose layer data for previously-saved images.)
+  it('should support all past layer naming schemes to preserve backward compatibility', async () => {
+    const dummyTimestamp = 1234567890
+    const inputToSupport = `clipspace-painted-masked-${dummyTimestamp}.png`
+    const expectedOutput = {
+      maskedImage: `clipspace-mask-${dummyTimestamp}.png`,
+      paint: `clipspace-paint-${dummyTimestamp}.png`,
+      paintedImage: `clipspace-painted-${dummyTimestamp}.png`,
+      paintedMaskedImage: inputToSupport
+    }
+    const actualOutput = imageLayerFilenamesIfApplicable(inputToSupport)
+    expect(actualOutput).toEqual(expectedOutput)
+  })
+})