Update link_fixer from rgthree

This commit is contained in:
Chenlei Hu
2025-05-07 14:44:28 -04:00
parent bb1ac32ccd
commit fd01f13fbf
2 changed files with 508 additions and 379 deletions

View File

@@ -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<ValidationResult> {
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
}
}

View File

@@ -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<T = ISerialisedGraph | LGraph> {
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<G> | 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<G> {
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<ISerialisedNode | LGraphNode>
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<G> {
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<SerialisedLLinkArray | LLink> = []
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<LGraph, LGraphNode> {
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
}
}