custom_node provided blueprints (#6172)

Frontend implementation for a system to allow custom_nodes to provide a
set of subgraph blueprints.

Requires comfyanonymous/ComfyUI#10438, but handles gracefully in unavailable.

┆Issue is synchronized with this [Notion
page](https://www.notion.so/PR-6172-Early-POC-custom_node-provided-blueprints-2926d73d3650814982ecd43f12abd873)
by [Unito](https://www.unito.io)

---------

Co-authored-by: Christian Byrne <[email protected]>

Author: AustinMroz

src/platform/workflow/management/stores/workflowStore.ts

@@ -83,11 +83,11 @@ export class ComfyWorkflow extends UserFile {
    * @param force Whether to force loading the content even if it is already loaded.
    * @returns this
    */
-  override async load({
-    force = false
-  }: { force?: boolean } = {}): Promise<LoadedComfyWorkflow> {
+  override async load({ force = false }: { force?: boolean } = {}): Promise<
+    this & LoadedComfyWorkflow
+  > {
     await super.load({ force })
-    if (!force && this.isLoaded) return this as LoadedComfyWorkflow
+    if (!force && this.isLoaded) return this as this & LoadedComfyWorkflow
 
     if (!this.originalContent) {
       throw new Error('[ASSERT] Workflow content should be loaded')
@@ -100,7 +100,7 @@ export class ComfyWorkflow extends UserFile {
         /* initialState= */ JSON.parse(this.originalContent)
       )
     )
-    return this as LoadedComfyWorkflow
+    return this as this & LoadedComfyWorkflow
   }
 
   override unload(): void {

src/scripts/api.ts

@@ -204,6 +204,12 @@ type SimpleApiEvents = keyof PickNevers<ApiEventTypes>
 /** Keys (names) of API events that pass a {@link CustomEvent} `detail` object. */
 type ComplexApiEvents = keyof NeverNever<ApiEventTypes>
 
+export type GlobalSubgraphData = {
+  name: string
+  info: { node_pack: string }
+  data: string | Promise<string>
+}
+
 function addHeaderEntry(headers: HeadersInit, key: string, value: string) {
   if (Array.isArray(headers)) {
     headers.push([key, value])
@@ -1118,6 +1124,22 @@ export class ComfyApi extends EventTarget {
     return resp.json()
   }
 
+  async getGlobalSubgraphData(id: string): Promise<string> {
+    const resp = await api.fetchApi('/global_subgraphs/' + id)
+    if (resp.status !== 200) return ''
+    const subgraph: GlobalSubgraphData = await resp.json()
+    return subgraph?.data ?? ''
+  }
+  async getGlobalSubgraphs(): Promise<Record<string, GlobalSubgraphData>> {
+    const resp = await api.fetchApi('/global_subgraphs')
+    if (resp.status !== 200) return {}
+    const subgraphs: Record<string, GlobalSubgraphData> = await resp.json()
+    for (const [k, v] of Object.entries(subgraphs)) {
+      if (!v.data) v.data = this.getGlobalSubgraphData(k)
+    }
+    return subgraphs
+  }
+
   async getLogs(): Promise<string> {
     return (await axios.get(this.internalURL('/logs'))).data
   }

src/scripts/app.ts

@@ -775,7 +775,8 @@ export class ComfyApp {
     this.canvasElRef.value = canvasEl
 
     await useWorkspaceStore().workflow.syncWorkflows()
-    await useSubgraphStore().fetchSubgraphs()
+    //Doesn't need to block. Blueprints will load async
+    void useSubgraphStore().fetchSubgraphs()
     await useExtensionService().loadExtensions()
 
     this.addProcessKeyHandler()

src/stores/subgraphStore.ts

@@ -23,6 +23,7 @@ import type {
   InputSpec
 } from '@/schemas/nodeDefSchema'
 import { api } from '@/scripts/api'
+import type { GlobalSubgraphData } from '@/scripts/api'
 import { useDialogService } from '@/services/dialogService'
 import { useExecutionStore } from '@/stores/executionStore'
 import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
@@ -106,9 +107,9 @@ export const useSubgraphStore = defineStore('subgraph', () => {
       useSubgraphStore().updateDef(await this.load())
       return ret
     }
-    override async load({
-      force = false
-    }: { force?: boolean } = {}): Promise<LoadedComfyWorkflow> {
+    override async load({ force = false }: { force?: boolean } = {}): Promise<
+      this & LoadedComfyWorkflow
+    > {
       if (!force && this.isLoaded) return await super.load({ force })
       const loaded = await super.load({ force })
       const st = loaded.activeState
@@ -147,20 +148,46 @@ export const useSubgraphStore = defineStore('subgraph', () => {
       modified: number
       size: number
     }): Promise<void> {
-      const name = options.path.slice(0, -'.json'.length)
       options.path = SubgraphBlueprint.basePath + options.path
       const bp = await new SubgraphBlueprint(options, true).load()
       useWorkflowStore().attachWorkflow(bp)
-      const nodeDef = convertToNodeDef(bp)
-
-      subgraphDefCache.value.set(name, nodeDef)
-      subgraphCache[name] = bp
+      registerNodeDef(bp)
+    }
+    async function loadInstalledBlueprints() {
+      async function loadGlobalBlueprint([k, v]: [string, GlobalSubgraphData]) {
+        const path = SubgraphBlueprint.basePath + v.name + '.json'
+        const blueprint = new SubgraphBlueprint({
+          path,
+          modified: Date.now(),
+          size: -1
+        })
+        blueprint.originalContent = blueprint.content = await v.data
+        blueprint.filename = v.name
+        useWorkflowStore().attachWorkflow(blueprint)
+        const loaded = await blueprint.load()
+        registerNodeDef(
+          loaded,
+          {
+            python_module: v.info.node_pack,
+            display_name: v.name
+          },
+          k
+        )
+      }
+      const subgraphs = await api.getGlobalSubgraphs()
+      await Promise.allSettled(
+        Object.entries(subgraphs).map(loadGlobalBlueprint)
+      )
     }
 
-    const res = (
+    const userSubs = (
       await api.listUserDataFullInfo(SubgraphBlueprint.basePath)
     ).filter((f) => f.path.endsWith('.json'))
-    const settled = await Promise.allSettled(res.map(loadBlueprint))
+    const settled = await Promise.allSettled([
+      ...userSubs.map(loadBlueprint),
+      loadInstalledBlueprints()
+    ])
+
     const errors = settled.filter((i) => 'reason' in i).map((i) => i.reason)
     errors.forEach((e) => console.error('Failed to load subgraph blueprint', e))
     if (errors.length > 0) {
@@ -172,8 +199,11 @@ export const useSubgraphStore = defineStore('subgraph', () => {
       })
     }
   }
-  function convertToNodeDef(workflow: LoadedComfyWorkflow): ComfyNodeDefImpl {
-    const name = workflow.filename
+  function registerNodeDef(
+    workflow: LoadedComfyWorkflow,
+    overrides: Partial<ComfyNodeDefV1> = {},
+    name: string = workflow.filename
+  ) {
     const subgraphNode = workflow.changeTracker.initialState.nodes[0]
     if (!subgraphNode) throw new Error('Invalid Subgraph Blueprint')
     subgraphNode.inputs ??= []
@@ -197,10 +227,12 @@ export const useSubgraphStore = defineStore('subgraph', () => {
       description,
       category: 'Subgraph Blueprints',
       output_node: false,
-      python_module: 'blueprint'
+      python_module: 'blueprint',
+      ...overrides
     }
     const nodeDefImpl = new ComfyNodeDefImpl(nodedefv1)
-    return nodeDefImpl
+    subgraphDefCache.value.set(name, nodeDefImpl)
+    subgraphCache[name] = workflow
   }
   async function publishSubgraph() {
     const canvas = canvasStore.getCanvas()
@@ -252,8 +284,7 @@ export const useSubgraphStore = defineStore('subgraph', () => {
     await workflow.save()
     //add to files list?
     useWorkflowStore().attachWorkflow(loadedWorkflow)
-    subgraphDefCache.value.set(name, convertToNodeDef(loadedWorkflow))
-    subgraphCache[name] = loadedWorkflow
+    registerNodeDef(loadedWorkflow)
     useToastStore().add({
       severity: 'success',
       summary: t('subgraphStore.publishSuccess'),
@@ -262,7 +293,7 @@ export const useSubgraphStore = defineStore('subgraph', () => {
     })
   }
   function updateDef(blueprint: LoadedComfyWorkflow) {
-    subgraphDefCache.value.set(blueprint.filename, convertToNodeDef(blueprint))
+    registerNodeDef(blueprint)
   }
   async function editBlueprint(nodeType: string) {
     const name = nodeType.slice(typePrefix.length)