From 929d0f5bdf1d495589cd7694181c9a366d700301 Mon Sep 17 00:00:00 2001
From: bymyself <abolkonsky.rem@gmail.com>
Date: Fri, 21 Feb 2025 06:25:44 -0700
Subject: [PATCH] move patch to setup hook

---
 js/workflow-metadata.js | 113 +++++++++++-----------------------------
 1 file changed, 31 insertions(+), 82 deletions(-)

diff --git a/js/workflow-metadata.js b/js/workflow-metadata.js
index 4d159f2f..9bf29285 100644
--- a/js/workflow-metadata.js
+++ b/js/workflow-metadata.js
@@ -42,100 +42,49 @@ class WorkflowMetadataExtension {
     return await res.json();
   }
 
-  /**
-   * Set node versions for the given graph
-   * @param {LGraph} graph The graph to process
-   * @param {Object} workflow The serialized workflow object
-   */
-  setGraphNodeVersions(graph, workflow) {
-    if (!graph?.nodes || !workflow?.nodes) return;
-
-    // Create a map of workflow nodes by ID
-    const workflowNodesById = {};
-    for (const node of workflow.nodes) {
-      if (node.id != null) {
-        workflowNodesById[node.id] = node;
-      }
-    }
-
-    // Process each graph node and find its corresponding workflow node by ID
-    for (const graphNode of graph.nodes) {
-      const workflowNode = workflowNodesById[graphNode.id];
-      if (!workflowNode) continue;
-      this.setNodeVersion(graphNode, workflowNode);
-    }
-  }
-
-  /**
-   * Set version information for a single node
-   * @param {Object} graphNode The graph node
-   * @param {Object} workflowNode The workflow node
-   */
-  setNodeVersion(graphNode, workflowNode) {
-    const nodeProperties = (workflowNode.properties ??= {});
-    const nodeData = graphNode.constructor?.nodeData;
-
-    // The node is missing or is a frontend only node
-    // (node was not constructed in registerNodes closure where nodeData is set)
-    if (!nodeData) {
-      return;
-    }
-
-    const modules = nodeData.python_module.split(".");
-    const moduleType = modules[0];
-
-    if (moduleType === "custom_nodes") {
-      this.setCustomNodeVersion(modules[1], nodeProperties);
-    } else if (["nodes", "comfy_extras"].includes(moduleType)) {
-      this.setCoreNodeVersion(nodeProperties);
-    } else {
-      console.warn(`Unknown node source: ${nodeData.python_module}`);
-    }
-  }
-
-  /**
-   * Set version for custom nodes
-   * @private
-   */
-  setCustomNodeVersion(nodePackageName, nodeProperties) {
-    const nodeInfo =
-      this.installedNodes[nodePackageName] ??
-      this.installedNodes[nodePackageName.toLowerCase()];
-
-    if (nodeInfo) {
-      if (nodeInfo.cnr_id === "comfy-core") return; // reserved package name
-      // Preserve workflow cnr_id and version if they exist
-      nodeProperties.cnr_id ??= nodeInfo.cnr_id;
-      nodeProperties.version ??= nodeInfo.ver;
-    }
-  }
-
-  /**
-   * Set version for core nodes
-   * @private
-   */
-  setCoreNodeVersion(nodeProperties) {
-    nodeProperties.cnr_id ??= "comfy-core";
-    nodeProperties.version ??= this.comfyCoreVersion;
-  }
-
   async init() {
-    const extension = this;
     this.installedNodes = await this.getInstalledNodes();
     this.comfyCoreVersion = (await api.getSystemStats()).system.comfyui_version;
+  }
 
-    // Attach metadata when app.graphToPrompt is called.
+  setup() {
+    const extension = this;
     const originalSerialize = LGraph.prototype.serialize;
+
     LGraph.prototype.serialize = function () {
       const workflow = originalSerialize.apply(this, arguments);
-
-      // Add metadata to the nodes
       const graph = this;
+
       try {
-        extension.setGraphNodeVersions(graph, workflow);
+        // Add version metadata to each node in the workflow
+        for (const node of workflow.nodes) {
+          const graphNode = graph.getNodeById(node.id);
+          if (!graphNode?.constructor?.nodeData?.python_module) continue;
+
+          const nodeProperties = (node.properties ??= {});
+          const modules =
+            graphNode.constructor.nodeData.python_module.split(".");
+          const moduleType = modules[0];
+
+          if (moduleType === "custom_nodes") {
+            const nodePackageName = modules[1];
+            const { cnr_id, ver } =
+              extension.installedNodes[nodePackageName] ??
+              extension.installedNodes[nodePackageName.toLowerCase()] ??
+              {};
+
+            if (cnr_id === "comfy-core") continue; // reserved name
+            if (cnr_id) nodeProperties.cnr_id = cnr_id;
+            if (ver) nodeProperties.ver = ver;
+          } else if (["nodes", "comfy_extras"].includes(moduleType)) {
+            nodeProperties.cnr_id = "comfy-core";
+            nodeProperties.ver = extension.comfyCoreVersion;
+          }
+        }
       } catch (e) {
         console.error(e);
       }
+
       return workflow;
     };
   }