From 59e20964a07819125139222ea00e32a797fe3b95 Mon Sep 17 00:00:00 2001
From: Christian Byrne <cbyrne@comfy.org>
Date: Fri, 11 Apr 2025 21:50:30 +0800
Subject: [PATCH] [Manager] Add fallback to infer node pack when metadata is
 missing from workflow (#3396)

---
 src/composables/nodePack/useWorkflowPacks.ts | 66 +++++++++++++++-----
 src/stores/comfyRegistryStore.ts             |  9 +++
 2 files changed, 61 insertions(+), 14 deletions(-)

diff --git a/src/composables/nodePack/useWorkflowPacks.ts b/src/composables/nodePack/useWorkflowPacks.ts
index b1dc3ec06..f68614ac1 100644
--- a/src/composables/nodePack/useWorkflowPacks.ts
+++ b/src/composables/nodePack/useWorkflowPacks.ts
@@ -1,10 +1,11 @@
 import { LGraphNode } from '@comfyorg/litegraph'
-import { computed, onUnmounted } from 'vue'
+import { computed, onUnmounted, ref } from 'vue'
 
 import { useNodePacks } from '@/composables/nodePack/useNodePacks'
 import { ComfyWorkflowJSON } from '@/schemas/comfyWorkflowSchema'
 import { app } from '@/scripts/app'
-import { UseNodePacksOptions } from '@/types/comfyManagerTypes'
+import { useComfyRegistryStore } from '@/stores/comfyRegistryStore'
+import { SelectedVersion, UseNodePacksOptions } from '@/types/comfyManagerTypes'
 import type { components } from '@/types/comfyRegistryTypes'
 
 type WorkflowPack = {
@@ -21,7 +22,11 @@ const CORE_NODES_PACK_NAME = 'comfy-core'
  * associated node packs from the registry
  */
 export const useWorkflowPacks = (options: UseNodePacksOptions = {}) => {
-  const getWorkflowNodeId = (node: LGraphNode): string | undefined => {
+  const { search } = useComfyRegistryStore()
+
+  const workflowPacks = ref<WorkflowPack[]>([])
+
+  const getWorkflowNodePackId = (node: LGraphNode): string | undefined => {
     if (typeof node.properties?.cnr_id === 'string') {
       return node.properties.cnr_id
     }
@@ -31,26 +36,56 @@ export const useWorkflowPacks = (options: UseNodePacksOptions = {}) => {
     return undefined
   }
 
-  const workflowNodeToPack = (node: LGraphNode): WorkflowPack | undefined => {
-    const id = getWorkflowNodeId(node)
-    if (!id) return undefined
-    if (id === CORE_NODES_PACK_NAME) return undefined
+  /**
+   * Infer the pack for a node by searching the registry for packs that have nodes
+   * with the same name.
+   */
+  const inferPack = async (
+    node: LGraphNode
+  ): Promise<WorkflowPack | undefined> => {
+    const nodeName = node.type
+    const searchResult = await search.call({
+      comfy_node_search: nodeName,
+      limit: 1
+    })
+    if (searchResult?.nodes?.length) {
+      const pack = searchResult.nodes[0]
+      return {
+        id: pack.id,
+        version: pack.latest_version?.version ?? SelectedVersion.NIGHTLY
+      }
+    }
+  }
+
+  /**
+   * Map a workflow node to its pack using the node pack metadata.
+   * If the node pack metadata is not available, fallback to searching the
+   * registry for packs that have nodes with the same name.
+   */
+  const workflowNodeToPack = async (
+    node: LGraphNode
+  ): Promise<WorkflowPack | undefined> => {
+    const packId = getWorkflowNodePackId(node)
+    if (!packId) return inferPack(node) // Fallback
+    if (packId === CORE_NODES_PACK_NAME) return undefined
 
     const version =
       typeof node.properties.ver === 'string' ? node.properties.ver : undefined
 
     return {
-      id,
+      id: packId,
       version
     }
   }
 
-  const workflowPacks = computed<WorkflowPack[]>(() => {
+  /**
+   * Get the node packs for all nodes in the workflow.
+   */
+  const getWorkflowPacks = async () => {
     if (!app.graph?.nodes?.length) return []
-    return app.graph.nodes
-      .map(workflowNodeToPack)
-      .filter((pack) => pack !== undefined)
-  })
+    const packs = await Promise.all(app.graph.nodes.map(workflowNodeToPack))
+    workflowPacks.value = packs.filter((pack) => pack !== undefined)
+  }
 
   const packsToUniqueIds = (packs: WorkflowPack[]) =>
     packs.reduce((acc, pack) => {
@@ -80,7 +115,10 @@ export const useWorkflowPacks = (options: UseNodePacksOptions = {}) => {
     isLoading,
     isReady,
     workflowPacks: nodePacks,
-    startFetchWorkflowPacks: startFetch,
+    startFetchWorkflowPacks: async () => {
+      await getWorkflowPacks() // Parse the packs from the workflow nodes
+      startFetch() // Fetch the packs infos from the registry
+    },
     filterWorkflowPack
   }
 }
diff --git a/src/stores/comfyRegistryStore.ts b/src/stores/comfyRegistryStore.ts
index 12eb929c3..a5a54bbb4 100644
--- a/src/stores/comfyRegistryStore.ts
+++ b/src/stores/comfyRegistryStore.ts
@@ -93,6 +93,14 @@ export const useComfyRegistryStore = defineStore('comfyRegistry', () => {
     ComfyNode[]
   >(registryService.getNodeDefs, { maxSize: PACK_BY_ID_CACHE_SIZE })
 
+  /**
+   * Search for packs by pack name or node names
+   */
+  const search = useCachedRequest<
+    operations['searchNodes']['parameters']['query'],
+    ListPacksResult
+  >(registryService.search, { maxSize: PACK_LIST_CACHE_SIZE })
+
   /**
    * Clear all cached data
    */
@@ -120,6 +128,7 @@ export const useComfyRegistryStore = defineStore('comfyRegistry', () => {
       cancel: () => getPacksByIdController?.abort()
     },
     getNodeDefs,
+    search,
 
     clearCache,
     cancelRequests,