From d1715972e337a6fafc75da6f585fad4d92b4cd54 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Fri, 16 Aug 2024 12:16:41 -0400
Subject: [PATCH] Register frontend only node defs (#468)

* Register frontend only node defs

* nit

* nit
---
 src/scripts/app.ts         | 29 +++++++++++++++++++++++++++--
 src/stores/nodeDefStore.ts | 10 +++++-----
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index 674d76053..79591b1a3 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -1964,12 +1964,37 @@ export class ComfyApp {
     // Load node definitions from the backend
     const defs = await api.getNodeDefs()
     await this.registerNodesFromDefs(defs)
+    await this.#invokeExtensionsAsync('registerCustomNodes')
     if (this.vueAppReady) {
+      // Frontend only nodes registered by custom nodes.
+      // Example: https://github.com/rgthree/rgthree-comfy/blob/dd534e5384be8cf0c0fa35865afe2126ba75ac55/src_web/comfyui/fast_groups_bypasser.ts#L10
+      const rawDefs = Object.fromEntries(
+        Object.entries(LiteGraph.registered_node_types).map(([name, _]) => [
+          name,
+          {
+            name,
+            display_name: name,
+            category: '__frontend_only__',
+            input: { required: {}, optional: {} },
+            output: [],
+            output_name: [],
+            output_is_list: [],
+            python_module: 'custom_nodes.frontend_only',
+            description: `Frontend only node for ${name}`
+          }
+        ])
+      )
+
+      const allNodeDefs = {
+        ...rawDefs,
+        ...defs,
+        ...SYSTEM_NODE_DEFS
+      }
+
       const nodeDefStore = useNodeDefStore()
-      nodeDefStore.updateNodeDefs([...Object.values(defs), ...SYSTEM_NODE_DEFS])
+      nodeDefStore.updateNodeDefs(Object.values(allNodeDefs))
       nodeDefStore.updateWidgets(this.widgets)
     }
-    await this.#invokeExtensionsAsync('registerCustomNodes')
   }
 
   getWidgetType(inputData, inputName) {
diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts
index 03715f1c5..6ca8de839 100644
--- a/src/stores/nodeDefStore.ts
+++ b/src/stores/nodeDefStore.ts
@@ -190,8 +190,8 @@ export class ComfyNodeDefImpl {
   }
 }
 
-export const SYSTEM_NODE_DEFS: ComfyNodeDef[] = [
-  {
+export const SYSTEM_NODE_DEFS: Record<string, ComfyNodeDef> = {
+  PrimitiveNode: {
     name: 'PrimitiveNode',
     display_name: 'Primitive',
     category: 'utils',
@@ -202,7 +202,7 @@ export const SYSTEM_NODE_DEFS: ComfyNodeDef[] = [
     python_module: 'nodes',
     description: 'Primitive values like numbers, strings, and booleans.'
   },
-  {
+  Reroute: {
     name: 'Reroute',
     display_name: 'Reroute',
     category: 'utils',
@@ -213,7 +213,7 @@ export const SYSTEM_NODE_DEFS: ComfyNodeDef[] = [
     python_module: 'nodes',
     description: 'Reroute the connection to another node.'
   },
-  {
+  Note: {
     name: 'Note',
     display_name: 'Note',
     category: 'utils',
@@ -224,7 +224,7 @@ export const SYSTEM_NODE_DEFS: ComfyNodeDef[] = [
     python_module: 'nodes',
     description: 'Node that add notes to your project'
   }
-]
+}
 
 function sortedTree(node: TreeNode): TreeNode {
   // Create a new node with the same label and data