ComfyUI_frontend/src/services/subgraphService.ts

import {
  type ExportedSubgraph,
  type ExportedSubgraphInstance,
  type Subgraph
} from '@comfyorg/litegraph'

import type { ComfyWorkflowJSON } from '@/schemas/comfyWorkflowSchema'
import type { ComfyNodeDef as ComfyNodeDefV1 } from '@/schemas/nodeDefSchema'
import { app as comfyApp } from '@/scripts/app'
import { useNodeDefStore } from '@/stores/nodeDefStore'

import { useLitegraphService } from './litegraphService'

export const useSubgraphService = () => {
  /** Loads a single subgraph definition and registers it with the node def store */
  const deserialiseSubgraph = (
    subgraph: Subgraph,
    exportedSubgraph: ExportedSubgraph
  ) => {
    const { id, name } = exportedSubgraph

    const nodeDef: ComfyNodeDefV1 = {
      input: { required: {} },
      output: [],
      output_is_list: [],
      output_name: [],
      output_tooltips: [],
      name: id,
      display_name: name,
      description: `Subgraph node for ${name}`,
      category: 'subgraph',
      output_node: false,
      python_module: 'nodes'
    }

    useNodeDefStore().addNodeDef(nodeDef)

    const instanceData: ExportedSubgraphInstance = {
      id: -1,
      type: exportedSubgraph.id,
      pos: [0, 0],
      size: [100, 100],
      inputs: [],
      outputs: [],
      flags: {},
      order: 0,
      mode: 0
    }

    useLitegraphService().registerSubgraphNodeDef(
      nodeDef,
      subgraph,
      instanceData
    )
  }

  /** Loads all exported subgraph definitionsfrom workflow */
  const loadSubgraphs = (graphData: ComfyWorkflowJSON) => {
    if (!graphData.definitions?.subgraphs) return

    for (const subgraphData of graphData.definitions.subgraphs) {
      const subgraph =
        comfyApp.graph.subgraphs.get(subgraphData.id) ??
        comfyApp.graph.createSubgraph(subgraphData as ExportedSubgraph)

      // @ts-expect-error Zod
      deserialiseSubgraph(subgraph, subgraphData)
    }
  }

  /** Registers a new subgraph (e.g. user converted from nodes) */
  const registerNewSubgraph = (
    subgraph: Subgraph,
    exportedSubgraph: ExportedSubgraph
  ) => {
    if (comfyApp.graph.subgraphs.has(subgraph.id)) {
      console.debug(`Subgraph ${subgraph.id} already registered`)
      return
    }

    deserialiseSubgraph(subgraph, exportedSubgraph)
  }

  return {
    loadSubgraphs,
    registerNewSubgraph
  }
}