Fix extension node def manipulation (#2873)

src/scripts/app.ts

@@ -19,7 +19,6 @@ import {
   type NodeId,
   validateComfyWorkflow
 } from '@/schemas/comfyWorkflowSchema'
-import { transformNodeDefV1ToV2 } from '@/schemas/nodeDef/migration'
 import {
   type ComfyNodeDef as ComfyNodeDefV2,
   isComboInputSpec,
@@ -889,9 +888,7 @@ export class ComfyApp {
     nodeDefStore.updateNodeDefs(nodeDefArray)
   }
 
-  async #getNodeDefs(): Promise<
-    Record<string, ComfyNodeDefV1 & ComfyNodeDefV2>
-  > {
+  async #getNodeDefs(): Promise<Record<string, ComfyNodeDefV1>> {
     const translateNodeDef = (def: ComfyNodeDefV1): ComfyNodeDefV1 => ({
       ...def,
       display_name: st(
@@ -911,7 +908,7 @@ export class ComfyApp {
       await api.getNodeDefs({
         validate: useSettingStore().get('Comfy.Validation.NodeDefs')
       }),
-      (def) => new ComfyNodeDefImpl(translateNodeDef(def))
+      (def) => translateNodeDef(def)
     )
   }
 
@@ -947,15 +944,7 @@ export class ComfyApp {
   }
 
   async registerNodeDef(nodeId: string, nodeDef: ComfyNodeDefV1) {
-    return await useLitegraphService().registerNodeDef(
-      nodeId,
-      isComfyNodeDefV2(nodeDef)
-        ? nodeDef
-        : {
-            ...(nodeDef as ComfyNodeDefV1),
-            ...transformNodeDefV1ToV2(nodeDef)
-          }
-    )
+    return await useLitegraphService().registerNodeDef(nodeId, nodeDef)
   }
 
   async registerNodesFromDefs(defs: Record<string, ComfyNodeDefV1>) {
@@ -1606,31 +1595,26 @@ export class ComfyApp {
       useToastStore().add(requestToastMessage)
     }
 
-    const defs: Record<string, ComfyNodeDefV1 & ComfyNodeDefV2> =
-      await this.#getNodeDefs()
-
-    for (const [nodeId, nodeDef] of Object.entries(defs)) {
-      this.registerNodeDef(nodeId, nodeDef)
+    const defs = await this.#getNodeDefs()
+    for (const nodeId in defs) {
+      this.registerNodeDef(nodeId, defs[nodeId])
     }
-
     for (const node of this.graph.nodes) {
       const def = defs[node.type]
       // Allow primitive nodes to handle refresh
       node.refreshComboInNode?.(defs)
 
-      if (!def) continue
+      if (!def?.input) continue
 
-      // Update combo options in combo widgets
       for (const widget of node.widgets) {
-        const inputSpec = def.inputs[widget.name]
-        if (
-          inputSpec &&
-          isComboInputSpec(inputSpec) &&
-          widget.type === 'combo'
-        ) {
-          widget.options.values = inputSpec.options.map((o) =>
-            typeof o === 'string' ? o : o.toString()
-          )
+        if (widget.type === 'combo') {
+          if (def['input'].required?.[widget.name] !== undefined) {
+            // @ts-expect-error InputSpec is not typed correctly
+            widget.options.values = def['input'].required[widget.name][0]
+          } else if (def['input'].optional?.[widget.name] !== undefined) {
+            // @ts-expect-error InputSpec is not typed correctly
+            widget.options.values = def['input'].optional[widget.name][0]
+          }
         }
       }
     }

src/services/litegraphService.ts

@@ -20,6 +20,7 @@ import { $el } from '@/scripts/ui'
 import { calculateImageGrid, createImageHost } from '@/scripts/ui/imagePreview'
 import { useCanvasStore } from '@/stores/graphStore'
 import { useNodeOutputStore } from '@/stores/imagePreviewStore'
+import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
 import { useToastStore } from '@/stores/toastStore'
 import { normalizeI18nKey } from '@/utils/formatUtil'
 import { is_all_same_aspect_ratio } from '@/utils/imageUtil'
@@ -35,14 +36,11 @@ export const useLitegraphService = () => {
   const toastStore = useToastStore()
   const canvasStore = useCanvasStore()
 
-  async function registerNodeDef(
-    nodeId: string,
-    nodeDef: ComfyNodeDefV2 & ComfyNodeDefV1
-  ) {
+  async function registerNodeDef(nodeId: string, nodeDefV1: ComfyNodeDefV1) {
     const node = class ComfyNode extends LGraphNode {
-      static comfyClass?: string = nodeDef.name
-      static title?: string = nodeDef.display_name || nodeDef.name
-      static nodeData?: ComfyNodeDefV1 & ComfyNodeDefV2 = nodeDef
+      static comfyClass?: string
+      static title?: string
+      static nodeData?: ComfyNodeDefV1 & ComfyNodeDefV2
       static category?: string
 
       constructor(title?: string) {
@@ -170,7 +168,6 @@ export const useLitegraphService = () => {
         super.configure(data)
       }
     }
-    node.prototype.comfyClass = nodeDef.name
 
     addNodeContextMenuHandler(node)
     addDrawBackgroundHandler(node)
@@ -179,12 +176,18 @@ export const useLitegraphService = () => {
     await extensionService.invokeExtensionsAsync(
       'beforeRegisterNodeDef',
       node,
-      nodeDef // Receives V1 NodeDef
+      nodeDefV1 // Receives V1 NodeDef, and potentially make modifications to it
     )
 
+    const nodeDef = new ComfyNodeDefImpl(nodeDefV1)
+    node.comfyClass = nodeDef.name
+    node.prototype.comfyClass = nodeDef.name
+    node.nodeData = nodeDef
     LiteGraph.registerNodeType(nodeId, node)
-    // Note: Do not move this to the class definition, it will be overwritten
+    // Note: Do not following assignments before `LiteGraph.registerNodeType`
+    // because `registerNodeType` will overwrite the assignments.
     node.category = nodeDef.category
+    node.title = nodeDef.display_name || nodeDef.name
   }
 
   /**

src/stores/nodeDefStore.ts

@@ -70,6 +70,12 @@ export class ComfyNodeDefImpl implements ComfyNodeDefV1, ComfyNodeDefV2 {
   readonly nodeSource: NodeSource
 
   constructor(obj: ComfyNodeDefV1) {
+    /**
+     * Assign extra fields to `this` for compatibility with group node feature.
+     * TODO: Remove this once group node feature is removed.
+     */
+    Object.assign(this, obj)
+
     // Initialize V1 fields
     this.name = obj.name
     this.display_name = obj.display_name