feat: replace hardcoded CANVAS_IMAGE_PREVIEW_NODE_TYPES with declarative flag

Replace the hardcoded set of node type names with a declarative
canvas_image_preview boolean on node definitions. Custom nodes can now
opt in by setting canvas_image_preview: true in their node def.

- Update supportsVirtualCanvasImagePreview() to check nodeData flag
- Add canvas_image_preview to zComfyNodeDef schema and ComfyNodeDefImpl
- Add canvas_image_preview to LGraphNode.nodeData type
- Create Comfy.CanvasImagePreview extension to set flag for core nodes
- Update promotionUtils tests to use subclass nodes with flag

Amp-Thread-ID: https://ampcode.com/threads/T-019d2deb-7154-745a-a267-491d6c75c666
Co-authored-by: Amp <amp@ampcode.com>

src/composables/node/canvasImagePreviewTypes.test.ts

@@ -0,0 +1,36 @@
+import { describe, expect, it } from 'vitest'
+
+import { LGraphNode } from '@/lib/litegraph/src/litegraph'
+
+import { supportsVirtualCanvasImagePreview } from './canvasImagePreviewTypes'
+
+function createNodeWithFlag(flag?: boolean): LGraphNode {
+  const node = new LGraphNode('TestNode')
+  const ctor = node.constructor as typeof LGraphNode
+  ctor.nodeData = flag !== undefined ? { canvas_image_preview: flag } : {}
+  return node
+}
+
+describe('supportsVirtualCanvasImagePreview', () => {
+  it('returns true when nodeData.canvas_image_preview is true', () => {
+    const node = createNodeWithFlag(true)
+    expect(supportsVirtualCanvasImagePreview(node)).toBe(true)
+  })
+
+  it('returns false when nodeData.canvas_image_preview is false', () => {
+    const node = createNodeWithFlag(false)
+    expect(supportsVirtualCanvasImagePreview(node)).toBe(false)
+  })
+
+  it('returns false when nodeData.canvas_image_preview is not set', () => {
+    const node = createNodeWithFlag()
+    expect(supportsVirtualCanvasImagePreview(node)).toBe(false)
+  })
+
+  it('returns false when nodeData is undefined', () => {
+    const node = new LGraphNode('TestNode')
+    const ctor = node.constructor as typeof LGraphNode
+    ctor.nodeData = undefined
+    expect(supportsVirtualCanvasImagePreview(node)).toBe(false)
+  })
+})

src/composables/node/canvasImagePreviewTypes.ts

@@ -1,12 +1,10 @@
 import type { LGraphNode } from '@/lib/litegraph/src/litegraph'
 
 export const CANVAS_IMAGE_PREVIEW_WIDGET = '$$canvas-image-preview'
-const CANVAS_IMAGE_PREVIEW_NODE_TYPES = new Set([
-  'PreviewImage',
-  'SaveImage',
-  'GLSLShader'
-])
 
 export function supportsVirtualCanvasImagePreview(node: LGraphNode): boolean {
-  return CANVAS_IMAGE_PREVIEW_NODE_TYPES.has(node.type)
+  return (
+    (node.constructor as typeof LGraphNode).nodeData?.canvas_image_preview ===
+    true
+  )
 }

src/core/graph/subgraph/promotionUtils.test.ts

@@ -139,10 +139,14 @@ describe('pruneDisconnected', () => {
     expect(warnSpy).toHaveBeenCalledOnce()
   })
 
-  it('keeps virtual canvas preview promotions for PreviewImage nodes', () => {
+  it('keeps virtual canvas preview promotions for canvas_image_preview nodes', () => {
     const subgraph = createTestSubgraph()
     const subgraphNode = createTestSubgraphNode(subgraph)
-    const interiorNode = new LGraphNode('PreviewImage')
+
+    class PreviewImageNode extends LGraphNode {
+      static override nodeData = { canvas_image_preview: true }
+    }
+    const interiorNode = new PreviewImageNode('PreviewImage')
     interiorNode.type = 'PreviewImage'
     subgraphNode.subgraph.add(interiorNode)
 
@@ -168,8 +172,12 @@ describe('pruneDisconnected', () => {
 })
 
 describe('getPromotableWidgets', () => {
-  it('adds virtual canvas preview widget for PreviewImage nodes', () => {
-    const node = new LGraphNode('PreviewImage')
+  class CanvasPreviewNode extends LGraphNode {
+    static override nodeData = { canvas_image_preview: true }
+  }
+
+  it('adds virtual canvas preview widget when canvas_image_preview is true', () => {
+    const node = new CanvasPreviewNode('PreviewImage')
     node.type = 'PreviewImage'
 
     const widgets = getPromotableWidgets(node)
@@ -179,29 +187,7 @@ describe('getPromotableWidgets', () => {
     ).toBe(true)
   })
 
-  it('adds virtual canvas preview widget for SaveImage nodes', () => {
-    const node = new LGraphNode('SaveImage')
-    node.type = 'SaveImage'
-
-    const widgets = getPromotableWidgets(node)
-
-    expect(
-      widgets.some((widget) => widget.name === CANVAS_IMAGE_PREVIEW_WIDGET)
-    ).toBe(true)
-  })
-
-  it('adds virtual canvas preview widget for GLSLShader nodes', () => {
-    const node = new LGraphNode('GLSLShader')
-    node.type = 'GLSLShader'
-
-    const widgets = getPromotableWidgets(node)
-
-    expect(
-      widgets.some((widget) => widget.name === CANVAS_IMAGE_PREVIEW_WIDGET)
-    ).toBe(true)
-  })
-
-  it('does not add virtual canvas preview widget for non-image nodes', () => {
+  it('does not add virtual canvas preview widget for non-canvas_image_preview nodes', () => {
     const node = new LGraphNode('TextNode')
     node.addOutput('TEXT', 'STRING')
 
@@ -211,17 +197,6 @@ describe('getPromotableWidgets', () => {
       widgets.some((widget) => widget.name === CANVAS_IMAGE_PREVIEW_WIDGET)
     ).toBe(false)
   })
-
-  it('does not add virtual canvas preview widget for ImageInvert nodes', () => {
-    const node = new LGraphNode('ImageInvert')
-    node.type = 'ImageInvert'
-
-    const widgets = getPromotableWidgets(node)
-
-    expect(
-      widgets.some((widget) => widget.name === CANVAS_IMAGE_PREVIEW_WIDGET)
-    ).toBe(false)
-  })
 })
 
 describe('promoteRecommendedWidgets', () => {
@@ -250,10 +225,14 @@ describe('promoteRecommendedWidgets', () => {
     expect(updatePreviewsMock).not.toHaveBeenCalled()
   })
 
-  it('eagerly promotes virtual preview widget for CANVAS_IMAGE_PREVIEW nodes', () => {
+  it('eagerly promotes virtual preview widget for canvas_image_preview nodes', () => {
     const subgraph = createTestSubgraph()
     const subgraphNode = createTestSubgraphNode(subgraph)
-    const glslNode = new LGraphNode('GLSLShader')
+
+    class GLSLShaderNode extends LGraphNode {
+      static override nodeData = { canvas_image_preview: true }
+    }
+    const glslNode = new GLSLShaderNode('GLSLShader')
     glslNode.type = 'GLSLShader'
     subgraph.add(glslNode)
 
@@ -269,12 +248,16 @@ describe('promoteRecommendedWidgets', () => {
     expect(updatePreviewsMock).not.toHaveBeenCalled()
   })
 
-  it('registers $$canvas-image-preview on configure for GLSLShader in saved workflow', () => {
+  it('registers $$canvas-image-preview on configure for canvas_image_preview node in saved workflow', () => {
     // Simulate loading a saved workflow where proxyWidgets does NOT contain
     // the $$canvas-image-preview entry (e.g. blueprint authored before the
     // promotion system, or old workflow save).
     const subgraph = createTestSubgraph()
-    const glslNode = new LGraphNode('GLSLShader')
+
+    class GLSLShaderNode extends LGraphNode {
+      static override nodeData = { canvas_image_preview: true }
+    }
+    const glslNode = new GLSLShaderNode('GLSLShader')
     glslNode.type = 'GLSLShader'
     subgraph.add(glslNode)
 

src/extensions/core/canvasImagePreview.ts

@@ -0,0 +1,16 @@
+import { app } from '@/scripts/app'
+
+const CANVAS_IMAGE_PREVIEW_NODES = new Set([
+  'PreviewImage',
+  'SaveImage',
+  'GLSLShader'
+])
+
+app.registerExtension({
+  name: 'Comfy.CanvasImagePreview',
+  beforeRegisterNodeDef(_nodeType, nodeData) {
+    if (CANVAS_IMAGE_PREVIEW_NODES.has(nodeData.name)) {
+      nodeData.canvas_image_preview = true
+    }
+  }
+})

src/extensions/core/index.ts

@@ -1,5 +1,6 @@
 import { isCloud, isNightly } from '@/platform/distribution/types'
 
+import './canvasImagePreview'
 import './clipspace'
 import './contextMenuFilter'
 import './customWidgets'

src/lib/litegraph/src/LGraphNode.ts

@@ -247,6 +247,7 @@ export class LGraphNode
     output_node?: boolean
     api_node?: boolean
     name?: string
+    canvas_image_preview?: boolean
   }
 
   static resizeHandleSize = 15

src/schemas/nodeDefSchema.ts

@@ -314,7 +314,9 @@ export const zComfyNodeDef = z.object({
   /** Category for the Essentials tab. If set, the node appears in Essentials. */
   essentials_category: z.string().optional(),
   /** Whether the blueprint is a global/installed blueprint (not user-created). */
-  isGlobal: z.boolean().optional()
+  isGlobal: z.boolean().optional(),
+  /** Whether the node supports canvas image preview rendering. */
+  canvas_image_preview: z.boolean().optional()
 })
 
 export const zAutogrowOptions = z.object({

src/stores/nodeDefStore.ts

@@ -90,6 +90,8 @@ export class ComfyNodeDefImpl
   readonly essentials_category?: string
   /** Whether the blueprint is a global/installed blueprint (not user-created). */
   readonly isGlobal?: boolean
+  /** Whether the node supports canvas image preview rendering. */
+  readonly canvas_image_preview?: boolean
   readonly isCoreNode: boolean
 
   // V2 fields
@@ -169,6 +171,7 @@ export class ComfyNodeDefImpl
         ) ?? obj.essentials_category)
       : undefined
     this.isGlobal = obj.isGlobal
+    this.canvas_image_preview = obj.canvas_image_preview
     this.isCoreNode = CORE_NODE_MODULES.includes(
       this.python_module.split('.')[0]
     )