Skip to content

Commit

Permalink
[Refactor] Move app.graphToPrompt to executionUtil (#2664)
Browse files Browse the repository at this point in the history
  • Loading branch information
@huchenlei
huchenlei authored Feb 21, 2025
1 parent cea5a4a commit ec94811
Show file tree
Hide file tree
Showing 3 changed files with 212 additions and 171 deletions.
175 changes: 5 additions & 170 deletions src/scripts/app.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ import {
} from '@/types/comfyWorkflow'
import { ExtensionManager } from '@/types/extensionTypes'
import { ColorAdjustOptions, adjustColor } from '@/utils/colorUtil'
import { graphToPrompt } from '@/utils/executionUtil'
import { isImageNode } from '@/utils/litegraphUtil'
import { deserialiseAndCreate } from '@/utils/vintageClipboard'

Expand Down Expand Up @@ -1226,176 +1227,11 @@ export class ComfyApp {
return graph.serialize({ sortNodes })
}

/**
* Converts the current graph workflow for sending to the API.
* Note: Node widgets are updated before serialization to prepare queueing.
* @returns The workflow and node links
*/
async graphToPrompt(graph = this.graph, clean = true) {
for (const outerNode of graph.computeExecutionOrder(false)) {
if (outerNode.widgets) {
for (const widget of outerNode.widgets) {
// Allow widgets to run callbacks before a prompt has been queued
// e.g. random seed before every gen
widget.beforeQueued?.()
}
}

const innerNodes = outerNode.getInnerNodes
? outerNode.getInnerNodes()
: [outerNode]
for (const node of innerNodes) {
if (node.isVirtualNode) {
// Don't serialize frontend only nodes but let them make changes
if (node.applyToGraph) {
node.applyToGraph()
}
}
}
}

const workflow = this.serializeGraph(graph)

// Remove localized_name from the workflow
for (const node of workflow.nodes) {
for (const slot of node.inputs) {
delete slot.localized_name
}
for (const slot of node.outputs) {
delete slot.localized_name
}
}

const output = {}
// Process nodes in order of execution
for (const outerNode of graph.computeExecutionOrder(false)) {
const skipNode =
outerNode.mode === LGraphEventMode.NEVER ||
outerNode.mode === LGraphEventMode.BYPASS
const innerNodes =
!skipNode && outerNode.getInnerNodes
? outerNode.getInnerNodes()
: [outerNode]
for (const node of innerNodes) {
if (node.isVirtualNode) {
continue
}

if (
node.mode === LGraphEventMode.NEVER ||
node.mode === LGraphEventMode.BYPASS
) {
// Don't serialize muted nodes
continue
}

const inputs = {}
const widgets = node.widgets

// Store all widget values
if (widgets) {
for (let i = 0; i < widgets.length; i++) {
const widget = widgets[i]
if (!widget.options || widget.options.serialize !== false) {
inputs[widget.name] = widget.serializeValue
? await widget.serializeValue(node, i)
: widget.value
}
}
}

// Store all node links
for (let i = 0; i < node.inputs.length; i++) {
let parent = node.getInputNode(i)
if (parent) {
let link = node.getInputLink(i)
while (
parent.mode === LGraphEventMode.BYPASS ||
parent.isVirtualNode
) {
let found = false
if (parent.isVirtualNode) {
link = parent.getInputLink(link.origin_slot)
if (link) {
parent = parent.getInputNode(link.target_slot)
if (parent) {
found = true
}
}
} else if (link && parent.mode === LGraphEventMode.BYPASS) {
let all_inputs = [link.origin_slot]
if (parent.inputs) {
// @ts-expect-error convert list of strings to list of numbers
all_inputs = all_inputs.concat(Object.keys(parent.inputs))
for (let parent_input in all_inputs) {
// @ts-expect-error assign string to number
parent_input = all_inputs[parent_input]
if (
parent.inputs[parent_input]?.type === node.inputs[i].type
) {
// @ts-expect-error convert string to number
link = parent.getInputLink(parent_input)
if (link) {
// @ts-expect-error convert string to number
parent = parent.getInputNode(parent_input)
}
found = true
break
}
}
}
}

if (!found) {
break
}
}

if (link) {
if (parent?.updateLink) {
link = parent.updateLink(link)
}
if (link) {
inputs[node.inputs[i].name] = [
String(link.origin_id),
// @ts-expect-error link.origin_slot is already number.
parseInt(link.origin_slot)
]
}
}
}
}

const node_data = {
inputs,
class_type: node.comfyClass
}

// Ignored by the backend.
node_data['_meta'] = {
title: node.title
}

output[String(node.id)] = node_data
}
}

// Remove inputs connected to removed nodes
if (clean) {
for (const o in output) {
for (const i in output[o].inputs) {
if (
Array.isArray(output[o].inputs[i]) &&
output[o].inputs[i].length === 2 &&
!output[output[o].inputs[i][0]]
) {
delete output[o].inputs[i]
}
}
}
}

return { workflow, output }
return graphToPrompt(graph, {
clean,
sortNodes: useSettingStore().get('Comfy.Workflow.SortNodeIdOnSave')
})
}

#formatPromptError(error) {
Expand Down Expand Up @@ -1444,7 +1280,6 @@ export class ComfyApp {
const p = await this.graphToPrompt()

try {
// @ts-expect-error Discrepancies between zod and litegraph - in progress
const res = await api.queuePrompt(number, p)
this.lastNodeErrors = res.node_errors
if (this.lastNodeErrors.length > 0) {
Expand Down
24 changes: 23 additions & 1 deletion src/types/comfyWorkflow.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import { fromZodError } from 'zod-validation-error'
// innerNode.id = `${this.node.id}:${i}`
// Remove it after GroupNode is redesigned.
export const zNodeId = z.union([z.number().int(), z.string()])
export const zNodeInputName = z.string()
export type NodeId = z.infer<typeof zNodeId>
export const zSlotIndex = z.union([
z.number().int(),
Expand Down Expand Up @@ -96,7 +97,7 @@ const zNodeOutput = z

const zNodeInput = z
.object({
name: z.string(),
name: zNodeInputName,
type: zDataType,
link: z.number().nullable().optional(),
slot_index: zSlotIndex.optional()
Expand Down Expand Up @@ -251,3 +252,24 @@ export async function validateComfyWorkflow(
onError(`Invalid workflow against zod schema:\n${error}`)
return null
}

/**
* API format workflow for direct API usage.
*/
const zNodeInputValue = z.union([
// For widget values (can be any type)
z.any(),
// For node links [nodeId, slotIndex]
z.tuple([zNodeId, zSlotIndex])
])

const zNodeData = z.object({
inputs: z.record(zNodeInputName, zNodeInputValue),
class_type: z.string(),
_meta: z.object({
title: z.string()
})
})

export const zComfyApiWorkflow = z.record(zNodeId, zNodeData)
export type ComfyApiWorkflow = z.infer<typeof zComfyApiWorkflow>
Loading

0 comments on commit ec94811

Please sign in to comment.