mirror of
https://github.com/Comfy-Org/ComfyUI_frontend.git
synced 2026-01-27 03:19:56 +00:00
478 lines
16 KiB
TypeScript
478 lines
16 KiB
TypeScript
/**
|
|
* 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
|
|
*
|
|
* MIT License
|
|
*
|
|
* Copyright (c) 2023 Regis Gaughan, III (rgthree)
|
|
*
|
|
* Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
* of this software and associated documentation files (the "Software"), to deal
|
|
* in the Software without restriction, including without limitation the rights
|
|
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
* copies of the Software, and to permit persons to whom the Software is
|
|
* furnished to do so, subject to the following conditions:
|
|
*
|
|
* The above copyright notice and this permission notice shall be included in all
|
|
* copies or substantial portions of the Software.
|
|
*
|
|
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
* 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 { NodeId } from '@comfyorg/litegraph/dist/LGraphNode'
|
|
import type { SerialisedLLinkArray } from '@comfyorg/litegraph/dist/LLink'
|
|
import type {
|
|
ISerialisedGraph,
|
|
ISerialisedNode
|
|
} from '@comfyorg/litegraph/dist/types/serialisation'
|
|
|
|
export interface BadLinksData<T = ISerialisedGraph | LGraph> {
|
|
hasBadLinks: boolean
|
|
fixed: boolean
|
|
graph: T
|
|
patched: number
|
|
deleted: number
|
|
}
|
|
|
|
enum IoDirection {
|
|
INPUT,
|
|
OUTPUT
|
|
}
|
|
|
|
function getNodeById(graph: ISerialisedGraph | LGraph, id: NodeId) {
|
|
if ((graph as LGraph).getNodeById) {
|
|
return (graph as LGraph).getNodeById(id)
|
|
}
|
|
graph = graph as ISerialisedGraph
|
|
return graph.nodes.find((node: ISerialisedNode) => node.id == id)!
|
|
}
|
|
|
|
function extendLink(link: SerialisedLLinkArray) {
|
|
return {
|
|
link: link,
|
|
id: link[0],
|
|
origin_id: link[1],
|
|
origin_slot: link[2],
|
|
target_id: link[3],
|
|
target_slot: link[4],
|
|
type: link[5]
|
|
}
|
|
}
|
|
|
|
/**
|
|
* 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.
|
|
*
|
|
* 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.
|
|
*/
|
|
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)
|
|
}
|
|
}
|
|
}
|
|
|
|
const patchedNodeSlots: {
|
|
[nodeId: string]: {
|
|
inputs?: { [slot: number]: number | null }
|
|
outputs?: {
|
|
[slots: number]: {
|
|
links: number[]
|
|
changes: { [linkId: number]: 'ADD' | 'REMOVE' }
|
|
}
|
|
}
|
|
}
|
|
} = {}
|
|
|
|
const data: {
|
|
patchedNodes: Array<ISerialisedNode | LGraphNode>
|
|
deletedLinks: number[]
|
|
} = {
|
|
patchedNodes: [],
|
|
deletedLinks: []
|
|
}
|
|
|
|
/**
|
|
* Internal patch node. We keep track of changes in patchedNodeSlots in case we're in a dry run.
|
|
*/
|
|
async function patchNodeSlot(
|
|
node: ISerialisedNode | LGraphNode,
|
|
ioDir: IoDirection,
|
|
slot: number,
|
|
linkId: number,
|
|
op: 'ADD' | 'REMOVE'
|
|
) {
|
|
patchedNodeSlots[node.id] = patchedNodeSlots[node.id] || {}
|
|
const patchedNode = 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.`
|
|
)
|
|
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)
|
|
}
|
|
}
|
|
}
|
|
data.patchedNodes.push(node)
|
|
return true
|
|
}
|
|
|
|
/**
|
|
* Internal to check if a node (or patched data) has a linkId.
|
|
*/
|
|
function nodeHasLinkId(
|
|
node: ISerialisedNode | LGraphNode,
|
|
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']) {
|
|
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.')
|
|
}
|
|
has = patchedHasIt
|
|
} else {
|
|
has = !!nodeHasIt
|
|
}
|
|
} else {
|
|
const nodeHasIt = node.outputs?.[slot]?.links?.includes(linkId)
|
|
if (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.')
|
|
}
|
|
has = !!patchedHasIt
|
|
} else {
|
|
has = !!nodeHasIt
|
|
}
|
|
}
|
|
return has
|
|
}
|
|
|
|
/**
|
|
* Internal to check if a node (or patched data) has a linkId.
|
|
*/
|
|
function nodeHasAnyLink(
|
|
node: ISerialisedNode | LGraphNode,
|
|
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']) {
|
|
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.')
|
|
}
|
|
hasAny = patchedHasAny
|
|
} else {
|
|
hasAny = !!nodeHasAny
|
|
}
|
|
} else {
|
|
const nodeHasAny = node.outputs?.[slot]?.links?.length
|
|
if (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.')
|
|
}
|
|
hasAny = !!patchedHasAny
|
|
} else {
|
|
hasAny = !!nodeHasAny
|
|
}
|
|
}
|
|
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
|
|
}
|
|
|
|
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')
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// 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)
|
|
}
|
|
}
|
|
// 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)
|
|
}
|
|
}
|
|
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
|
|
}
|
|
}
|