diff --git a/src/composables/useWorkflowValidation.ts b/src/composables/useWorkflowValidation.ts index 31c9a9556..86fbfbb5b 100644 --- a/src/composables/useWorkflowValidation.ts +++ b/src/composables/useWorkflowValidation.ts @@ -3,7 +3,7 @@ import type { ISerialisedGraph } from '@comfyorg/litegraph/dist/types/serialisat import type { ComfyWorkflowJSON } from '@/schemas/comfyWorkflowSchema' import { validateComfyWorkflow } from '@/schemas/comfyWorkflowSchema' import { useToastStore } from '@/stores/toastStore' -import { fixBadLinks } from '@/utils/linkFixer' +import { WorkflowLinkFixer } from '@/utils/linkFixer' export interface ValidationResult { graphData: ComfyWorkflowJSON | null @@ -27,7 +27,6 @@ export function useWorkflowValidation() { ): Promise { const { silent = false } = options - let linksFixes let validatedData: ComfyWorkflowJSON | null = null // First do schema validation @@ -44,20 +43,17 @@ export function useWorkflowValidation() { // Collect all logs in an array const logs: string[] = [] // Then validate and fix links if schema validation passed - const linkValidation = fixBadLinks( - validatedGraphData as unknown as ISerialisedGraph, - { - fix: true, - silent, - logger: { - log: (message: string) => { - logs.push(message) - } - } - } + const fixer = WorkflowLinkFixer.create( + validatedGraphData as unknown as ISerialisedGraph ) + fixer.logger = { + log: (...args: any[]) => { + logs.push(...args) + } + } + const checkBadLinksResult = fixer.check() - if (!silent && logs.length > 0) { + if (!silent && checkBadLinksResult.hasBadLinks) { toastStore.add({ severity: 'warn', summary: 'Workflow Validation', @@ -65,27 +61,18 @@ export function useWorkflowValidation() { }) } - // If links were fixed, notify the user - if (linkValidation.fixed) { - if (!silent) { - toastStore.add({ - severity: 'success', - summary: 'Workflow Links Fixed', - detail: `Fixed ${linkValidation.patched} node connections and removed ${linkValidation.deleted} invalid links.` - }) - } - } - - validatedData = linkValidation.graph as unknown as ComfyWorkflowJSON - linksFixes = { - patched: linkValidation.patched, - deleted: linkValidation.deleted + const fixBadLinksResult = fixer.fix() + if (!fixBadLinksResult.hasBadLinks && !silent) { + toastStore.add({ + severity: 'success', + summary: 'Workflow Links Fixed' + }) + validatedData = fixBadLinksResult.graph as unknown as ComfyWorkflowJSON } } return { - graphData: validatedData, - linksFixes + graphData: validatedData } } diff --git a/src/utils/linkFixer.ts b/src/utils/linkFixer.ts index 218b1c962..d27a5a5fa 100644 --- a/src/utils/linkFixer.ts +++ b/src/utils/linkFixer.ts @@ -1,6 +1,6 @@ /** * This code is adapted from rgthree-comfy's link_fixer.ts - * @see https://github.com/rgthree/rgthree-comfy/blob/b84f39c7c224de765de0b54c55b967329011819d/src_web/common/link_fixer.ts + * @see https://github.com/rgthree/rgthree-comfy/blob/main/src_web/common/link_fixer.ts * * MIT License * @@ -24,20 +24,31 @@ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE * SOFTWARE. */ -import type { LGraph, LGraphNode, LLink } from '@comfyorg/litegraph' +import type { + INodeOutputSlot, + ISlotType, + LGraph, + LGraphNode, + LLink +} from '@comfyorg/litegraph' import type { NodeId } from '@comfyorg/litegraph/dist/LGraphNode' -import type { SerialisedLLinkArray } from '@comfyorg/litegraph/dist/LLink' +import type { + LinkId, + SerialisedLLinkArray +} from '@comfyorg/litegraph/dist/LLink' import type { ISerialisedGraph, ISerialisedNode } from '@comfyorg/litegraph/dist/types/serialisation' +/** + * The bad links data returned from either a fixer `check()`, or the results of a `fix()` call. + */ export interface BadLinksData { hasBadLinks: boolean - fixed: boolean graph: T - patched: number - deleted: number + patches: number + deletes: number } enum IoDirection { @@ -45,183 +56,468 @@ enum IoDirection { OUTPUT } -function getNodeById(graph: ISerialisedGraph | LGraph, id: NodeId) { - if ((graph as LGraph).getNodeById) { - return (graph as LGraph).getNodeById(id) +/** + * Data interface that mimics a nodes `inputs` and `outputs` holding the _to be_ mutated node data + * during a check. + */ +interface PatchedNodeSlots { + [nodeId: string]: { + inputs?: { [slot: number]: number | null } + outputs?: { + [slots: number]: { + links: number[] + changes: { [linkId: number]: 'ADD' | 'REMOVE' } + } + } } - graph = graph as ISerialisedGraph - return graph.nodes.find((node: ISerialisedNode) => node.id == id)! } -function extendLink(link: SerialisedLLinkArray) { - return { - link: link, +/** + * Link data derived from either a ISerialisedGraph or LGraph `links` property. + */ +interface LinkData { + id: LinkId + origin_id: NodeId + origin_slot: number + target_id: NodeId + target_slot: number + type: ISlotType +} + +/** + * Returns a list of links data for the given links type; either from an LGraph or SerializedGraph. + */ +function getLinksData( + links: ISerialisedGraph['links'] | LGraph['links'] | { [key: string]: LLink } +): LinkData[] { + if (links instanceof Map) { + const data: LinkData[] = [] + for (const llink of links.values()) { + if (!llink) continue + data.push(llink) + } + return data + } + // This is apparently marked deprecated in ComfyUI but who knows if we would get stale data in + // here that's not a map (handled above). Go ahead and handle it anyway. + if (!Array.isArray(links)) { + const data: LinkData[] = [] + for (const key in links) { + // eslint-disable-next-line no-prototype-builtins + const llink = (links.hasOwnProperty(key) && links[key]) || null + if (!llink) continue + data.push(llink) + } + return data + } + return links.map((link: SerialisedLLinkArray) => ({ id: link[0], origin_id: link[1], origin_slot: link[2], target_id: link[3], target_slot: link[4], type: link[5] - } + })) } +/** The instruction data for fixing a node's inputs or outputs. */ +interface WorkflowLinkFixerNodeInstruction { + node: ISerialisedNode | LGraphNode + op: 'REMOVE' | 'ADD' + dir: IoDirection + slot: number + linkId: number + linkIdToUse: number | null +} + +/** The instruction data for fixing a link from a workflow links. */ +interface WorkflowLinkFixerLinksInstruction { + op: 'DELETE' + linkId: number + reason: string +} + +type WorkflowLinkFixerInstruction = + | WorkflowLinkFixerNodeInstruction + | WorkflowLinkFixerLinksInstruction + /** - * Takes a ISerialisedGraph or live LGraph and inspects the links and nodes to ensure the linking - * makes logical sense. Can apply fixes when passed the `fix` argument as true. + * The WorkflowLinkFixer for either ISerialisedGraph or a live LGraph. * - * Note that fixes are a best-effort attempt. Seems to get it correct in most cases, but there is a - * chance it correct an anomoly that results in placing an incorrect link (say, if there were two - * links in the data). Users should take care to not overwrite work until manually checking the - * result. + * Use `WorkflowLinkFixer.create(graph: ISerialisedGraph | LGraph)` to create a new instance. */ -export function fixBadLinks( - graph: ISerialisedGraph | LGraph, - options: { - fix?: boolean - silent?: boolean - logger?: { log: (...args: any[]) => void } - } = {} -): BadLinksData { - const { fix = false, silent = false, logger: _logger = console } = options - const logger = { - log: (...args: any[]) => { - if (!silent) { - _logger.log(...args) - } +export abstract class WorkflowLinkFixer< + G extends ISerialisedGraph | LGraph, + N extends ISerialisedNode | LGraphNode +> { + silent: boolean = false + checkedData: BadLinksData | null = null + logger: { log: (...args: any[]) => void } = console + + protected graph: G + protected patchedNodeSlots: PatchedNodeSlots = {} + protected instructions: WorkflowLinkFixerInstruction[] = [] + + /** + * Creates the WorkflowLinkFixer for the given graph type. + */ + static create(graph: ISerialisedGraph): WorkflowLinkFixerSerialized + static create(graph: LGraph): WorkflowLinkFixerGraph + static create( + graph: ISerialisedGraph | LGraph + ): WorkflowLinkFixerSerialized | WorkflowLinkFixerGraph { + if (typeof (graph as LGraph).getNodeById === 'function') { + return new WorkflowLinkFixerGraph(graph as LGraph) } + return new WorkflowLinkFixerSerialized(graph as ISerialisedGraph) } - const patchedNodeSlots: { - [nodeId: string]: { - inputs?: { [slot: number]: number | null } - outputs?: { - [slots: number]: { - links: number[] - changes: { [linkId: number]: 'ADD' | 'REMOVE' } + protected constructor(graph: G) { + this.graph = graph + } + + abstract getNodeById(id: NodeId): N | null + abstract deleteGraphLink(id: LinkId): true | string + + /** + * Checks the current graph data for any bad links. + */ + check(force: boolean = false): BadLinksData { + if (this.checkedData && !force) { + return { ...this.checkedData } + } + this.instructions = [] + this.patchedNodeSlots = {} + + const instructions: (WorkflowLinkFixerInstruction | null)[] = [] + + const links: LinkData[] = getLinksData(this.graph.links) + links.reverse() + for (const link of links) { + if (!link) continue + + const originNode = this.getNodeById(link.origin_id) + const originHasLink = () => + this.nodeHasLinkId( + originNode!, + IoDirection.OUTPUT, + link.origin_slot, + link.id + ) + const patchOrigin = (op: 'ADD' | 'REMOVE', id = link.id) => + this.getNodePatchInstruction( + originNode!, + IoDirection.OUTPUT, + link.origin_slot, + id, + op + ) + + const targetNode = this.getNodeById(link.target_id) + const targetHasLink = () => + this.nodeHasLinkId( + targetNode!, + IoDirection.INPUT, + link.target_slot, + link.id + ) + const targetHasAnyLink = () => + this.nodeHasAnyLink(targetNode!, IoDirection.INPUT, link.target_slot) + const patchTarget = (op: 'ADD' | 'REMOVE', id = link.id) => + this.getNodePatchInstruction( + targetNode!, + IoDirection.INPUT, + link.target_slot, + id, + op + ) + + const originLog = `origin(${link.origin_id}).outputs[${link.origin_slot}].links` + const targetLog = `target(${link.target_id}).inputs[${link.target_slot}].link` + + if (!originNode || !targetNode) { + if (!originNode && !targetNode) { + // This can fall through and continue; we remove it after this loop. + } else if (!originNode && targetNode) { + this.log( + `Link ${link.id} is funky... ` + + `origin ${link.origin_id} does not exist, but target ${link.target_id} does.` + ) + if (targetHasLink()) { + this.log( + ` > [PATCH] ${targetLog} does have link, will remove the inputs' link first.` + ) + instructions.push(patchTarget('REMOVE', -1)) + } + } else if (!targetNode && originNode) { + this.log( + `Link ${link.id} is funky... ` + + `target ${link.target_id} does not exist, but origin ${link.origin_id} does.` + ) + if (originHasLink()) { + this.log( + ` > [PATCH] Origin's links' has ${link.id}; will remove the link first.` + ) + instructions.push(patchOrigin('REMOVE')) + } + } + continue + } + + if (targetHasLink() || originHasLink()) { + if (!originHasLink()) { + this.log( + `${link.id} is funky... ${originLog} does NOT contain it, but ${targetLog} does.` + ) + this.log( + ` > [PATCH] Attempt a fix by adding this ${link.id} to ${originLog}.` + ) + instructions.push(patchOrigin('ADD')) + } else if (!targetHasLink()) { + this.log( + `${link.id} is funky... ${targetLog} is NOT correct (is ${ + targetNode.inputs?.[link.target_slot]?.link + }), but ${originLog} contains it` + ) + if (!targetHasAnyLink()) { + this.log( + ` > [PATCH] ${targetLog} is not defined, will set to ${link.id}.` + ) + let instruction = patchTarget('ADD') + if (!instruction) { + this.log( + ` > [PATCH] Nvm, ${targetLog} already patched. Removing ${link.id} from ${originLog}.` + ) + instruction = patchOrigin('REMOVE') + } + instructions.push(instruction) + } else { + this.log( + ` > [PATCH] ${targetLog} is defined, removing ${link.id} from ${originLog}.` + ) + instructions.push(patchOrigin('REMOVE')) + } } } } - } = {} - const data: { - patchedNodes: Array - deletedLinks: number[] - } = { - patchedNodes: [], - deletedLinks: [] + // Now that we've cleaned up the inputs, outputs, run through it looking for dangling links., + for (const link of links) { + if (!link) continue + const originNode = this.getNodeById(link.origin_id) + const targetNode = this.getNodeById(link.target_id) + if (!originNode && !targetNode) { + instructions.push({ + op: 'DELETE', + linkId: link.id, + reason: `Both nodes #${link.origin_id} & #${link.target_id} are removed` + }) + } + // Now that we've manipulated the linking, check again if they both exist. + if ( + (!originNode || + !this.nodeHasLinkId( + originNode, + IoDirection.OUTPUT, + link.origin_slot, + link.id + )) && + (!targetNode || + !this.nodeHasLinkId( + targetNode, + IoDirection.INPUT, + link.target_slot, + link.id + )) + ) { + instructions.push({ + op: 'DELETE', + linkId: link.id, + reason: + `both origin node #${link.origin_id} ` + + `${!originNode ? 'is removed' : `is missing link id output slot ${link.origin_slot}`}` + + `and target node #${link.target_id} ` + + `${!targetNode ? 'is removed' : `is missing link id input slot ${link.target_slot}`}.` + }) + continue + } + } + + this.instructions = instructions.filter((i) => !!i) + this.checkedData = { + hasBadLinks: !!this.instructions.length, + graph: this.graph, + patches: this.instructions.filter( + (i) => !!(i as WorkflowLinkFixerNodeInstruction).node + ).length, + deletes: this.instructions.filter((i) => i.op === 'DELETE').length + } + return { ...this.checkedData } } /** - * Internal patch node. We keep track of changes in patchedNodeSlots in case we're in a dry run. + * Fixes a checked graph by running through the instructions generated during the check run. Also + * double-checks for inconsistencies after the fix, recursively calling itself up to five times + * before giving up. */ - function patchNodeSlot( - node: ISerialisedNode | LGraphNode, + fix(force: boolean = false, times?: number): BadLinksData { + if (!this.checkedData || force) { + this.check(force) + } + let patches = 0 + let deletes = 0 + for (const instruction of this.instructions) { + if ((instruction as WorkflowLinkFixerNodeInstruction).node) { + const { node, slot, linkIdToUse, dir, op } = + instruction as WorkflowLinkFixerNodeInstruction + if (dir == IoDirection.INPUT) { + node.inputs = node.inputs || [] + const oldValue = node.inputs[slot]?.link + node.inputs[slot]!.link = linkIdToUse + this.log( + `Node #${node.id}: Set link ${linkIdToUse} to input slot ${slot} (was ${oldValue})` + ) + } else if (op === 'ADD' && linkIdToUse != null) { + node.outputs = node.outputs || [] + node.outputs[slot] = node.outputs[slot] || ({} as INodeOutputSlot) + node.outputs[slot].links = node.outputs[slot].links || [] + node.outputs[slot].links.push(linkIdToUse) + this.log( + `Node #${node.id}: Add link ${linkIdToUse} to output slot #${slot}` + ) + } else if (op === 'REMOVE' && linkIdToUse != null) { + const linkIdIndex = node.outputs![slot]!.links!.indexOf(linkIdToUse) + node.outputs![slot]!.links!.splice(linkIdIndex, 1) + this.log( + `Node #${node.id}: Remove link ${linkIdToUse} from output slot #${slot}` + ) + } else { + throw new Error('Unhandled Node Instruction') + } + patches++ + } else if (instruction.op === 'DELETE') { + const wasDeleted = this.deleteGraphLink(instruction.linkId) + if (wasDeleted === true) { + this.log( + `Link #${instruction.linkId}: Removed workflow link b/c ${instruction.reason}` + ) + } else { + this.log(`Error Link #${instruction.linkId}: ${wasDeleted}`) + } + deletes++ + } else { + throw new Error('Unhandled Instruction') + } + } + + const newCheck = this.check(force) + times = times == null ? 5 : times + let newFix = null + // If we still have bad links, then recurse (up to five times). + if (newCheck.hasBadLinks && times > 0) { + newFix = this.fix(true, times - 1) + } + + return { + hasBadLinks: newFix?.hasBadLinks ?? newCheck.hasBadLinks, + graph: this.graph, + patches: patches + (newFix?.patches ?? 0), + deletes: deletes + (newFix?.deletes ?? 0) + } + } + + /** Logs if not silent. */ + protected log(...args: any[]) { + if (this.silent) return + this.logger.log(...args) + } + + /** + * Patches a node for a check run, returning the instruction that would be made. + */ + private getNodePatchInstruction( + node: N, ioDir: IoDirection, slot: number, linkId: number, op: 'ADD' | 'REMOVE' - ) { - patchedNodeSlots[node.id] = patchedNodeSlots[node.id] || {} - const patchedNode = patchedNodeSlots[node.id]! + ): WorkflowLinkFixerNodeInstruction | null { + this.patchedNodeSlots[node.id] = this.patchedNodeSlots[node.id] || {} + const patchedNode = this.patchedNodeSlots[node.id]! if (ioDir == IoDirection.INPUT) { patchedNode['inputs'] = patchedNode['inputs'] || {} // We can set to null (delete), so undefined means we haven't set it at all. if (patchedNode['inputs']![slot] !== undefined) { - logger.log( - ` > Already set ${node.id}.inputs[${slot}] to ${patchedNode[ - 'inputs' - ]![slot]!} Skipping.` + this.log( + ` > Already set ${node.id}.inputs[${slot}] to ${patchedNode['inputs']![slot]!} Skipping.` ) - return false - } - const linkIdToSet = op === 'REMOVE' ? null : linkId - patchedNode['inputs']![slot] = linkIdToSet - if (fix) { - // node.inputs[slot]!.link = linkIdToSet; - } - } else { - patchedNode['outputs'] = patchedNode['outputs'] || {} - patchedNode['outputs']![slot] = patchedNode['outputs']![slot] || { - links: [...(node.outputs?.[slot]?.links || [])], - changes: {} - } - if (patchedNode['outputs']![slot]!['changes']![linkId] !== undefined) { - logger.log( - ` > Already set ${node.id}.outputs[${slot}] to ${ - patchedNode['inputs']![slot] - }! Skipping.` - ) - return false - } - patchedNode['outputs']![slot]!['changes']![linkId] = op - if (op === 'ADD') { - const linkIdIndex = - patchedNode['outputs']![slot]!['links'].indexOf(linkId) - if (linkIdIndex !== -1) { - logger.log( - ` > Hmmm.. asked to add ${linkId} but it is already in list...` - ) - return false - } - patchedNode['outputs']![slot]!['links'].push(linkId) - if (fix) { - node.outputs = node.outputs || [] - node.outputs[slot] = node.outputs[slot] || ({} as any) - node.outputs[slot]!.links = node.outputs[slot]!.links || [] - node.outputs[slot]!.links!.push(linkId) - } - } else { - const linkIdIndex = - patchedNode['outputs']![slot]!['links'].indexOf(linkId) - if (linkIdIndex === -1) { - logger.log( - ` > Hmmm.. asked to remove ${linkId} but it doesn't exist...` - ) - return false - } - patchedNode['outputs']![slot]!['links'].splice(linkIdIndex, 1) - if (fix) { - node.outputs?.[slot]!.links!.splice(linkIdIndex, 1) - } + return null } + const linkIdToUse = op === 'REMOVE' ? null : linkId + patchedNode['inputs']![slot] = linkIdToUse + return { node, dir: ioDir, op, slot, linkId, linkIdToUse } } - data.patchedNodes.push(node) - return true + + patchedNode['outputs'] = patchedNode['outputs'] || {} + patchedNode['outputs']![slot] = patchedNode['outputs']![slot] || { + links: [...(node.outputs?.[slot]?.links || [])], + changes: {} + } + if (patchedNode['outputs']![slot]!['changes']![linkId] !== undefined) { + this.log( + ` > Already set ${node.id}.outputs[${slot}] to ${patchedNode['inputs']![slot]}! Skipping.` + ) + return null + } + patchedNode['outputs']![slot]!['changes']![linkId] = op + if (op === 'ADD') { + const linkIdIndex = + patchedNode['outputs']![slot]!['links'].indexOf(linkId) + if (linkIdIndex !== -1) { + this.log( + ` > Hmmm.. asked to add ${linkId} but it is already in list...` + ) + return null + } + patchedNode['outputs']![slot]!['links'].push(linkId) + return { node, dir: ioDir, op, slot, linkId, linkIdToUse: linkId } + } + + const linkIdIndex = patchedNode['outputs']![slot]!['links'].indexOf(linkId) + if (linkIdIndex === -1) { + this.log(` > Hmmm.. asked to remove ${linkId} but it doesn't exist...`) + return null + } + patchedNode['outputs']![slot]!['links'].splice(linkIdIndex, 1) + return { node, dir: ioDir, op, slot, linkId, linkIdToUse: linkId } } - /** - * Internal to check if a node (or patched data) has a linkId. - */ - function nodeHasLinkId( - node: ISerialisedNode | LGraphNode, + /** Checks if a node (or patched data) has a linkId. */ + private nodeHasLinkId( + node: N, ioDir: IoDirection, slot: number, linkId: number ) { - // Patched data should be canonical. We can double check if fixing too. let has = false if (ioDir === IoDirection.INPUT) { const nodeHasIt = node.inputs?.[slot]?.link === linkId - if (patchedNodeSlots[node.id]?.['inputs']) { + if (this.patchedNodeSlots[node.id]?.['inputs']) { const patchedHasIt = - patchedNodeSlots[node.id]!['inputs']![slot] === linkId - // If we're fixing, double check that node matches. - if (fix && nodeHasIt !== patchedHasIt) { - throw Error('Error. Expected node to match patched data.') - } + this.patchedNodeSlots[node.id]!['inputs']![slot] === linkId has = patchedHasIt } else { - has = !!nodeHasIt + has = nodeHasIt } } else { const nodeHasIt = node.outputs?.[slot]?.links?.includes(linkId) - if (patchedNodeSlots[node.id]?.['outputs']?.[slot]?.['changes'][linkId]) { + if ( + this.patchedNodeSlots[node.id]?.['outputs']?.[slot]?.['changes'][linkId] + ) { const patchedHasIt = - patchedNodeSlots[node.id]!['outputs']![slot]?.links.includes(linkId) - // If we're fixing, double check that node matches. - if (fix && nodeHasIt !== patchedHasIt) { - throw Error('Error. Expected node to match patched data.') - } + this.patchedNodeSlots[node.id]!['outputs']![slot]?.links.includes( + linkId + ) has = !!patchedHasIt } else { has = !!nodeHasIt @@ -230,38 +526,24 @@ export function fixBadLinks( return has } - /** - * Internal to check if a node (or patched data) has a linkId. - */ - function nodeHasAnyLink( - node: ISerialisedNode | LGraphNode, - ioDir: IoDirection, - slot: number - ) { + /** Checks if a node (or patched data) has a linkId. */ + private nodeHasAnyLink(node: N, ioDir: IoDirection, slot: number) { // Patched data should be canonical. We can double check if fixing too. let hasAny = false if (ioDir === IoDirection.INPUT) { const nodeHasAny = node.inputs?.[slot]?.link != null - if (patchedNodeSlots[node.id]?.['inputs']) { + if (this.patchedNodeSlots[node.id]?.['inputs']) { const patchedHasAny = - patchedNodeSlots[node.id]!['inputs']![slot] != null - // If we're fixing, double check that node matches. - if (fix && nodeHasAny !== patchedHasAny) { - throw Error('Error. Expected node to match patched data.') - } + this.patchedNodeSlots[node.id]!['inputs']![slot] != null hasAny = patchedHasAny } else { hasAny = !!nodeHasAny } } else { const nodeHasAny = node.outputs?.[slot]?.links?.length - if (patchedNodeSlots[node.id]?.['outputs']?.[slot]?.['changes']) { + if (this.patchedNodeSlots[node.id]?.['outputs']?.[slot]?.['changes']) { const patchedHasAny = - patchedNodeSlots[node.id]!['outputs']![slot]?.links.length - // If we're fixing, double check that node matches. - if (fix && nodeHasAny !== patchedHasAny) { - throw Error('Error. Expected node to match patched data.') - } + this.patchedNodeSlots[node.id]!['outputs']![slot]?.links.length hasAny = !!patchedHasAny } else { hasAny = !!nodeHasAny @@ -269,209 +551,69 @@ export function fixBadLinks( } return hasAny } +} - let links: Array = [] - if (!Array.isArray(graph.links)) { - links = Object.values(graph.links).reduce((acc, v) => { - acc[v.id] = v - return acc - }, links) - } else { - links = graph.links +/** + * A WorkflowLinkFixer for serialized data. + */ +class WorkflowLinkFixerSerialized extends WorkflowLinkFixer< + ISerialisedGraph, + ISerialisedNode +> { + constructor(graph: ISerialisedGraph) { + super(graph) } - const linksReverse = [...links] - linksReverse.reverse() - for (const l of linksReverse) { - if (!l) continue - const link = - (l as LLink).origin_slot != null - ? (l as LLink) - : extendLink(l as SerialisedLLinkArray) - - const originNode = getNodeById(graph, link.origin_id) - const originHasLink = () => - nodeHasLinkId(originNode!, IoDirection.OUTPUT, link.origin_slot, link.id) - const patchOrigin = (op: 'ADD' | 'REMOVE', id = link.id) => - patchNodeSlot(originNode!, IoDirection.OUTPUT, link.origin_slot, id, op) - - const targetNode = getNodeById(graph, link.target_id) - const targetHasLink = () => - nodeHasLinkId(targetNode!, IoDirection.INPUT, link.target_slot, link.id) - const targetHasAnyLink = () => - nodeHasAnyLink(targetNode!, IoDirection.INPUT, link.target_slot) - const patchTarget = (op: 'ADD' | 'REMOVE', id = link.id) => - patchNodeSlot(targetNode!, IoDirection.INPUT, link.target_slot, id, op) - - const originLog = `origin(${link.origin_id}).outputs[${link.origin_slot}].links` - const targetLog = `target(${link.target_id}).inputs[${link.target_slot}].link` - - if (!originNode || !targetNode) { - if (!originNode && !targetNode) { - logger.log( - `Link ${link.id} is invalid, ` + - `both origin ${link.origin_id} and target ${link.target_id} do not exist` - ) - } else if (!originNode) { - logger.log( - `Link ${link.id} is funky... ` + - `origin ${link.origin_id} does not exist, but target ${link.target_id} does.` - ) - if (targetHasLink()) { - logger.log( - ` > [PATCH] ${targetLog} does have link, will remove the inputs' link first.` - ) - patchTarget('REMOVE', -1) - } - } else if (!targetNode) { - logger.log( - `Link ${link.id} is funky... ` + - `target ${link.target_id} does not exist, but origin ${link.origin_id} does.` - ) - if (originHasLink()) { - logger.log( - ` > [PATCH] Origin's links' has ${link.id}; will remove the link first.` - ) - patchOrigin('REMOVE') - } - } - continue - } - - if (targetHasLink() || originHasLink()) { - if (!originHasLink()) { - logger.log( - `${link.id} is funky... ${originLog} does NOT contain it, but ${targetLog} does.` - ) - - logger.log( - ` > [PATCH] Attempt a fix by adding this ${link.id} to ${originLog}.` - ) - patchOrigin('ADD') - } else if (!targetHasLink()) { - logger.log( - `${link.id} is funky... ${targetLog} is NOT correct (is ${ - targetNode.inputs?.[link.target_slot]?.link - }), but ${originLog} contains it` - ) - if (!targetHasAnyLink()) { - logger.log( - ` > [PATCH] ${targetLog} is not defined, will set to ${link.id}.` - ) - let patched = patchTarget('ADD') - if (!patched) { - logger.log( - ` > [PATCH] Nvm, ${targetLog} already patched. Removing ${link.id} from ${originLog}.` - ) - patched = patchOrigin('REMOVE') - } - } else { - logger.log( - ` > [PATCH] ${targetLog} is defined, removing ${link.id} from ${originLog}.` - ) - patchOrigin('REMOVE') - } - } - } + getNodeById(id: NodeId) { + return this.graph.nodes.find((node) => Number(node.id) === id) ?? null } - // Now that we've cleaned up the inputs, outputs, run through it looking for dangling links., - for (const l of linksReverse) { - if (!l) continue - const link = - (l as LLink).origin_slot != null - ? (l as LLink) - : extendLink(l as SerialisedLLinkArray) - const originNode = getNodeById(graph, link.origin_id) - const targetNode = getNodeById(graph, link.target_id) - // Now that we've manipulated the linking, check again if they both exist. - if ( - (!originNode || - !nodeHasLinkId( - originNode, - IoDirection.OUTPUT, - link.origin_slot, - link.id - )) && - (!targetNode || - !nodeHasLinkId( - targetNode, - IoDirection.INPUT, - link.target_slot, - link.id - )) - ) { - logger.log( - `${link.id} is def invalid; BOTH origin node ${link.origin_id} ${ - !originNode ? 'is removed' : `doesn't have ${link.id}` - } and ${link.origin_id} target node ${ - !targetNode ? 'is removed' : `doesn't have ${link.id}` - }.` - ) - data.deletedLinks.push(link.id) - continue - } - } - - // If we're fixing, then we've been patching along the way. Now go through and actually delete - // the zombie links from `app.graph.links` - if (fix) { - for (let i = data.deletedLinks.length - 1; i >= 0; i--) { - logger.log(`Deleting link #${data.deletedLinks[i]}.`) - if ((graph as LGraph).getNodeById) { - delete graph.links[data.deletedLinks[i]!] - } else { - graph = graph as ISerialisedGraph - // Sometimes we got objects for links if passed after ComfyUI's loadGraphData modifies the - // data. We make a copy now, but can handle the bastardized objects just in case. - const idx = graph.links.findIndex( - (l) => - l && - (l[0] === data.deletedLinks[i] || - (l as any).id === data.deletedLinks[i]) - ) - if (idx === -1) { - logger.log(`INDEX NOT FOUND for #${data.deletedLinks[i]}`) - } - logger.log(`splicing ${idx} from links`) - graph.links.splice(idx, 1) - } - } + override fix(force: boolean = false, times?: number) { + const ret = super.fix(force, times) // If we're a serialized graph, we can filter out the links because it's just an array. - if (!(graph as LGraph).getNodeById) { - graph.links = (graph as ISerialisedGraph).links.filter((l) => !!l) + this.graph.links = this.graph.links.filter((l) => !!l) + return ret + } + + deleteGraphLink(id: LinkId) { + // Sometimes we got objects instead of serializzed array for links if passed after ComfyUI's + // loadGraphData modifies the data. Let's find the id handling the bastardized objects just in + // case. + const idx = this.graph.links.findIndex( + (l) => l && (l[0] === id || (l as any).id === id) + ) + if (idx === -1) { + return `Link #${id} not found in workflow links.` } - } - if (!data.patchedNodes.length && !data.deletedLinks.length) { - return { - hasBadLinks: false, - fixed: false, - graph, - patched: data.patchedNodes.length, - deleted: data.deletedLinks.length - } - } - - logger.log( - `${fix ? 'Made' : 'Would make'} ${data.patchedNodes.length || 'no'} node link patches, and ${ - data.deletedLinks.length || 'no' - } stale link removals.` - ) - - let hasBadLinks: boolean = !!( - data.patchedNodes.length || data.deletedLinks.length - ) - // If we're fixing, then let's run it again to see if there are no more bad links. - if (fix && !silent) { - const rerun = fixBadLinks(graph, { fix: false, silent: true }) - hasBadLinks = rerun.hasBadLinks - } - - return { - hasBadLinks, - fixed: !!hasBadLinks && fix, - graph, - patched: data.patchedNodes.length, - deleted: data.deletedLinks.length + this.graph.links.splice(idx, 1) + return true + } +} + +/** + * A WorkflowLinkFixer for live LGraph data. + */ +class WorkflowLinkFixerGraph extends WorkflowLinkFixer { + constructor(graph: LGraph) { + super(graph) + } + + getNodeById(id: NodeId) { + return this.graph.getNodeById(id) ?? null + } + + deleteGraphLink(id: LinkId) { + if (this.graph.links instanceof Map) { + if (!this.graph.links.has(id)) { + return `Link #${id} not found in workflow links.` + } + this.graph.links.delete(id) + return true + } + if (this.graph.links[id] == null) { + return `Link #${id} not found in workflow links.` + } + delete this.graph.links[id] + return true } }