mirror of
https://github.com/Comfy-Org/ComfyUI_frontend.git
synced 2026-04-22 23:39:45 +00:00
Add front end support for type matching (#6582)
This PR implements front end logic to handle MatchType inputs and outputs. See comfyanonymous/ComfyUI#10644 This allows for the implementation of nodes such as a "switch node" where input types change based on the connections made.  As part of this implementation, significant cleanup is being performed in the reroute code. Extra testing will be required to make sure these changes don't introduce regressions. ┆Issue is synchronized with this [Notion page](https://www.notion.so/PR-6582-Add-front-end-support-for-type-matching-2a16d73d36508189b042cd23f82a332e) by [Unito](https://www.unito.io)
This commit is contained in:
AustinMroz
@@ -4,6 +4,7 @@ import {
|
||||
LGraphNode,
|
||||
LiteGraph
|
||||
} from '@/lib/litegraph/src/litegraph'
|
||||
import type { ISlotType } from '@/lib/litegraph/src/interfaces'
|
||||
|
||||
import { app } from '../../scripts/app'
|
||||
import { getWidgetConfig, mergeIfValid, setWidgetConfig } from './widgetInputs'
|
||||
@@ -14,7 +15,7 @@ app.registerExtension({
|
||||
name: 'Comfy.RerouteNode',
|
||||
registerCustomNodes(app) {
|
||||
interface RerouteNode extends LGraphNode {
|
||||
__outputType?: string
|
||||
__outputType?: string | number
|
||||
}
|
||||
|
||||
class RerouteNode extends LGraphNode {
|
||||
@@ -22,8 +23,7 @@ app.registerExtension({
|
||||
static defaultVisibility = false
|
||||
|
||||
constructor(title?: string) {
|
||||
// @ts-expect-error fixme ts strict error
|
||||
super(title)
|
||||
super(title ?? '')
|
||||
if (!this.properties) {
|
||||
this.properties = {}
|
||||
}
|
||||
@@ -33,225 +33,198 @@ app.registerExtension({
|
||||
this.addInput('', '*')
|
||||
this.addOutput(this.properties.showOutputText ? '*' : '', '*')
|
||||
|
||||
this.onAfterGraphConfigured = function () {
|
||||
requestAnimationFrame(() => {
|
||||
// @ts-expect-error fixme ts strict error
|
||||
this.onConnectionsChange(LiteGraph.INPUT, null, true, null)
|
||||
})
|
||||
}
|
||||
|
||||
this.onConnectionsChange = (type, _index, connected) => {
|
||||
if (app.configuringGraph) return
|
||||
|
||||
// Prevent multiple connections to different types when we have no input
|
||||
if (connected && type === LiteGraph.OUTPUT) {
|
||||
// Ignore wildcard nodes as these will be updated to real types
|
||||
const types = new Set(
|
||||
// @ts-expect-error fixme ts strict error
|
||||
this.outputs[0].links
|
||||
.map((l) => app.graph.links[l].type)
|
||||
.filter((t) => t !== '*')
|
||||
)
|
||||
if (types.size > 1) {
|
||||
const linksToDisconnect = []
|
||||
// @ts-expect-error fixme ts strict error
|
||||
for (let i = 0; i < this.outputs[0].links.length - 1; i++) {
|
||||
// @ts-expect-error fixme ts strict error
|
||||
const linkId = this.outputs[0].links[i]
|
||||
const link = app.graph.links[linkId]
|
||||
linksToDisconnect.push(link)
|
||||
}
|
||||
for (const link of linksToDisconnect) {
|
||||
const node = app.graph.getNodeById(link.target_id)
|
||||
// @ts-expect-error fixme ts strict error
|
||||
node.disconnectInput(link.target_slot)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Find root input
|
||||
let currentNode: LGraphNode | null = this
|
||||
let updateNodes = []
|
||||
let inputType = null
|
||||
let inputNode = null
|
||||
while (currentNode) {
|
||||
updateNodes.unshift(currentNode)
|
||||
const linkId = currentNode.inputs[0].link
|
||||
if (linkId !== null) {
|
||||
const link = app.graph.links[linkId]
|
||||
if (!link) return
|
||||
const node = app.graph.getNodeById(link.origin_id)
|
||||
// @ts-expect-error fixme ts strict error
|
||||
const type = node.constructor.type
|
||||
if (type === 'Reroute') {
|
||||
if (node === this) {
|
||||
// We've found a circle
|
||||
currentNode.disconnectInput(link.target_slot)
|
||||
currentNode = null
|
||||
} else {
|
||||
// Move the previous node
|
||||
currentNode = node
|
||||
}
|
||||
} else {
|
||||
// We've found the end
|
||||
inputNode = currentNode
|
||||
// @ts-expect-error fixme ts strict error
|
||||
inputType = node.outputs[link.origin_slot]?.type ?? null
|
||||
break
|
||||
}
|
||||
} else {
|
||||
// This path has no input node
|
||||
currentNode = null
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// Find all outputs
|
||||
const nodes: LGraphNode[] = [this]
|
||||
let outputType = null
|
||||
while (nodes.length) {
|
||||
// @ts-expect-error fixme ts strict error
|
||||
currentNode = nodes.pop()
|
||||
const outputs =
|
||||
// @ts-expect-error fixme ts strict error
|
||||
(currentNode.outputs ? currentNode.outputs[0].links : []) || []
|
||||
if (outputs.length) {
|
||||
for (const linkId of outputs) {
|
||||
const link = app.graph.links[linkId]
|
||||
|
||||
// When disconnecting sometimes the link is still registered
|
||||
if (!link) continue
|
||||
|
||||
const node = app.graph.getNodeById(link.target_id)
|
||||
// @ts-expect-error fixme ts strict error
|
||||
const type = node.constructor.type
|
||||
|
||||
if (type === 'Reroute') {
|
||||
// Follow reroute nodes
|
||||
// @ts-expect-error fixme ts strict error
|
||||
nodes.push(node)
|
||||
updateNodes.push(node)
|
||||
} else {
|
||||
// We've found an output
|
||||
const nodeOutType =
|
||||
// @ts-expect-error fixme ts strict error
|
||||
node.inputs &&
|
||||
// @ts-expect-error fixme ts strict error
|
||||
node.inputs[link?.target_slot] &&
|
||||
// @ts-expect-error fixme ts strict error
|
||||
node.inputs[link.target_slot].type
|
||||
? // @ts-expect-error fixme ts strict error
|
||||
node.inputs[link.target_slot].type
|
||||
: null
|
||||
if (
|
||||
inputType &&
|
||||
// @ts-expect-error fixme ts strict error
|
||||
!LiteGraph.isValidConnection(inputType, nodeOutType)
|
||||
) {
|
||||
// The output doesnt match our input so disconnect it
|
||||
// @ts-expect-error fixme ts strict error
|
||||
node.disconnectInput(link.target_slot)
|
||||
} else {
|
||||
outputType = nodeOutType
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// No more outputs for this path
|
||||
}
|
||||
}
|
||||
|
||||
const displayType = inputType || outputType || '*'
|
||||
const color = LGraphCanvas.link_type_colors[displayType]
|
||||
|
||||
let widgetConfig
|
||||
let widgetType
|
||||
// Update the types of each node
|
||||
for (const node of updateNodes) {
|
||||
// If we dont have an input type we are always wildcard but we'll show the output type
|
||||
// This lets you change the output link to a different type and all nodes will update
|
||||
// @ts-expect-error fixme ts strict error
|
||||
node.outputs[0].type = inputType || '*'
|
||||
// @ts-expect-error fixme ts strict error
|
||||
node.__outputType = displayType
|
||||
// @ts-expect-error fixme ts strict error
|
||||
node.outputs[0].name = node.properties.showOutputText
|
||||
? displayType
|
||||
: ''
|
||||
// @ts-expect-error fixme ts strict error
|
||||
node.setSize(node.computeSize())
|
||||
|
||||
// @ts-expect-error fixme ts strict error
|
||||
for (const l of node.outputs[0].links || []) {
|
||||
const link = app.graph.links[l]
|
||||
if (link) {
|
||||
link.color = color
|
||||
|
||||
if (app.configuringGraph) continue
|
||||
const targetNode = app.graph.getNodeById(link.target_id)
|
||||
// @ts-expect-error fixme ts strict error
|
||||
const targetInput = targetNode.inputs?.[link.target_slot]
|
||||
if (targetInput?.widget) {
|
||||
const config = getWidgetConfig(targetInput)
|
||||
if (!widgetConfig) {
|
||||
widgetConfig = config[1] ?? {}
|
||||
widgetType = config[0]
|
||||
}
|
||||
|
||||
const merged = mergeIfValid(targetInput, [
|
||||
config[0],
|
||||
widgetConfig
|
||||
])
|
||||
if (merged.customConfig) {
|
||||
widgetConfig = merged.customConfig
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (const node of updateNodes) {
|
||||
if (widgetConfig && outputType) {
|
||||
// @ts-expect-error fixme ts strict error
|
||||
node.inputs[0].widget = { name: 'value' }
|
||||
// @ts-expect-error fixme ts strict error
|
||||
setWidgetConfig(node.inputs[0], [
|
||||
// @ts-expect-error fixme ts strict error
|
||||
widgetType ?? displayType,
|
||||
widgetConfig
|
||||
])
|
||||
} else {
|
||||
// @ts-expect-error fixme ts strict error
|
||||
setWidgetConfig(node.inputs[0], null)
|
||||
}
|
||||
}
|
||||
|
||||
if (inputNode) {
|
||||
// @ts-expect-error fixme ts strict error
|
||||
const link = app.graph.links[inputNode.inputs[0].link]
|
||||
if (link) {
|
||||
link.color = color
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
this.clone = function () {
|
||||
const cloned = RerouteNode.prototype.clone.apply(this)
|
||||
// @ts-expect-error fixme ts strict error
|
||||
cloned.removeOutput(0)
|
||||
// @ts-expect-error fixme ts strict error
|
||||
cloned.addOutput(this.properties.showOutputText ? '*' : '', '*')
|
||||
// @ts-expect-error fixme ts strict error
|
||||
cloned.setSize(cloned.computeSize())
|
||||
return cloned
|
||||
}
|
||||
|
||||
// This node is purely frontend and does not impact the resulting prompt so should not be serialized
|
||||
this.isVirtualNode = true
|
||||
}
|
||||
override onAfterGraphConfigured() {
|
||||
requestAnimationFrame(() => {
|
||||
this.onConnectionsChange(LiteGraph.INPUT, undefined, true)
|
||||
})
|
||||
}
|
||||
override clone(): LGraphNode | null {
|
||||
const cloned = super.clone()
|
||||
if (!cloned) return cloned
|
||||
cloned.removeOutput(0)
|
||||
cloned.addOutput(this.properties.showOutputText ? '*' : '', '*')
|
||||
cloned.setSize(cloned.computeSize())
|
||||
return cloned
|
||||
}
|
||||
override onConnectionsChange(
|
||||
type: ISlotType,
|
||||
_index: number | undefined,
|
||||
connected: boolean
|
||||
) {
|
||||
const { graph } = this
|
||||
if (!graph) return
|
||||
if (app.configuringGraph) return
|
||||
|
||||
// @ts-expect-error fixme ts strict error
|
||||
getExtraMenuOptions(_, options): IContextMenuValue[] {
|
||||
// Prevent multiple connections to different types when we have no input
|
||||
if (connected && type === LiteGraph.OUTPUT) {
|
||||
// Ignore wildcard nodes as these will be updated to real types
|
||||
const types = new Set(
|
||||
this.outputs[0].links
|
||||
?.map((l) => graph.links[l]?.type)
|
||||
?.filter((t) => t && t !== '*') ?? []
|
||||
)
|
||||
if (types.size > 1) {
|
||||
const linksToDisconnect = []
|
||||
for (const linkId of this.outputs[0].links ?? []) {
|
||||
const link = graph.links[linkId]
|
||||
linksToDisconnect.push(link)
|
||||
}
|
||||
linksToDisconnect.pop()
|
||||
for (const link of linksToDisconnect) {
|
||||
const node = graph.getNodeById(link.target_id)
|
||||
node?.disconnectInput(link.target_slot)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Find root input
|
||||
let currentNode: RerouteNode | null = this
|
||||
let updateNodes: RerouteNode[] = []
|
||||
let inputType = null
|
||||
let inputNode = null
|
||||
while (currentNode) {
|
||||
updateNodes.unshift(currentNode)
|
||||
const linkId = currentNode.inputs[0].link
|
||||
if (linkId !== null) {
|
||||
const link = graph.links[linkId]
|
||||
if (!link) return
|
||||
const node = graph.getNodeById(link.origin_id)
|
||||
if (!node) return
|
||||
if (node instanceof RerouteNode) {
|
||||
if (node === this) {
|
||||
// We've found a circle
|
||||
currentNode.disconnectInput(link.target_slot)
|
||||
currentNode = null
|
||||
} else {
|
||||
// Move the previous node
|
||||
currentNode = node
|
||||
}
|
||||
} else {
|
||||
// We've found the end
|
||||
inputNode = currentNode
|
||||
inputType = node.outputs[link.origin_slot]?.type ?? null
|
||||
break
|
||||
}
|
||||
} else {
|
||||
// This path has no input node
|
||||
currentNode = null
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// Find all outputs
|
||||
const nodes: RerouteNode[] = [this]
|
||||
let outputType = null
|
||||
while (nodes.length) {
|
||||
currentNode = nodes.pop()!
|
||||
const outputs = currentNode.outputs?.[0]?.links ?? []
|
||||
for (const linkId of outputs) {
|
||||
const link = graph.links[linkId]
|
||||
|
||||
// When disconnecting sometimes the link is still registered
|
||||
if (!link) continue
|
||||
|
||||
const node = graph.getNodeById(link.target_id)
|
||||
if (!node) continue
|
||||
if (node instanceof RerouteNode) {
|
||||
// Follow reroute nodes
|
||||
nodes.push(node)
|
||||
updateNodes.push(node)
|
||||
} else {
|
||||
// We've found an output
|
||||
const nodeInput = node.inputs[link.target_slot]
|
||||
const nodeOutType = nodeInput.type
|
||||
const keep =
|
||||
!inputType ||
|
||||
!nodeOutType ||
|
||||
LiteGraph.isValidConnection(inputType, nodeOutType)
|
||||
if (!keep) {
|
||||
// The output doesnt match our input so disconnect it
|
||||
node.disconnectInput(link.target_slot)
|
||||
continue
|
||||
}
|
||||
node.onConnectionsChange?.(
|
||||
LiteGraph.INPUT,
|
||||
link.target_slot,
|
||||
keep,
|
||||
link,
|
||||
nodeInput
|
||||
)
|
||||
outputType = node.inputs[link.target_slot].type
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const displayType = inputType || outputType || '*'
|
||||
const color = LGraphCanvas.link_type_colors[displayType]
|
||||
|
||||
let widgetConfig
|
||||
let widgetType
|
||||
// Update the types of each node
|
||||
for (const node of updateNodes) {
|
||||
// If we dont have an input type we are always wildcard but we'll show the output type
|
||||
// This lets you change the output link to a different type and all nodes will update
|
||||
node.outputs[0].type = inputType || '*'
|
||||
node.__outputType = displayType
|
||||
node.outputs[0].name = node.properties.showOutputText
|
||||
? `${displayType}`
|
||||
: ''
|
||||
node.setSize(node.computeSize())
|
||||
|
||||
for (const l of node.outputs[0].links || []) {
|
||||
const link = graph.links[l]
|
||||
if (!link) continue
|
||||
link.color = color
|
||||
|
||||
if (app.configuringGraph) continue
|
||||
const targetNode = graph.getNodeById(link.target_id)
|
||||
if (!targetNode) continue
|
||||
const targetInput = targetNode.inputs?.[link.target_slot]
|
||||
if (targetInput?.widget) {
|
||||
const config = getWidgetConfig(targetInput)
|
||||
if (!widgetConfig) {
|
||||
widgetConfig = config[1] ?? {}
|
||||
widgetType = config[0]
|
||||
}
|
||||
|
||||
const merged = mergeIfValid(targetInput, [
|
||||
config[0],
|
||||
widgetConfig
|
||||
])
|
||||
if (merged.customConfig) {
|
||||
widgetConfig = merged.customConfig
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (const node of updateNodes) {
|
||||
if (widgetConfig && outputType) {
|
||||
node.inputs[0].widget = { name: 'value' }
|
||||
setWidgetConfig(node.inputs[0], [
|
||||
widgetType ?? `${displayType}`,
|
||||
widgetConfig
|
||||
])
|
||||
} else {
|
||||
setWidgetConfig(node.inputs[0], undefined)
|
||||
}
|
||||
}
|
||||
|
||||
if (inputNode?.inputs?.[0]?.link) {
|
||||
const link = graph.links[inputNode.inputs[0].link]
|
||||
if (link) {
|
||||
link.color = color
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
override getExtraMenuOptions(
|
||||
_: unknown,
|
||||
options: (IContextMenuValue | null)[]
|
||||
): IContextMenuValue[] {
|
||||
options.unshift(
|
||||
{
|
||||
content:
|
||||
@@ -259,13 +232,12 @@ app.registerExtension({
|
||||
callback: () => {
|
||||
this.properties.showOutputText = !this.properties.showOutputText
|
||||
if (this.properties.showOutputText) {
|
||||
this.outputs[0].name =
|
||||
this.__outputType || (this.outputs[0].type as string)
|
||||
this.outputs[0].name = `${this.__outputType || this.outputs[0].type}`
|
||||
} else {
|
||||
this.outputs[0].name = ''
|
||||
}
|
||||
this.setSize(this.computeSize())
|
||||
app.graph.setDirtyCanvas(true, true)
|
||||
app.canvas.setDirty(true, true)
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -294,8 +266,7 @@ app.registerExtension({
|
||||
]
|
||||
}
|
||||
|
||||
// @ts-expect-error fixme ts strict error
|
||||
static setDefaultTextVisibility(visible) {
|
||||
static setDefaultTextVisibility(visible: boolean) {
|
||||
RerouteNode.defaultVisibility = visible
|
||||
if (visible) {
|
||||
localStorage['Comfy.RerouteNode.DefaultVisibility'] = 'true'
|
||||
|
||||
Reference in New Issue
Block a user