diff --git a/src/extensions/core/widgetInputs.ts b/src/extensions/core/widgetInputs.ts index 25567376..181dbded 100644 --- a/src/extensions/core/widgetInputs.ts +++ b/src/extensions/core/widgetInputs.ts @@ -334,6 +334,26 @@ class PrimitiveNode extends LGraphNode { } } + isValidWidgetLink( + originSlot: number, + targetNode: LGraphNode, + targetWidget: IWidget + ) { + const config2 = getConfig.call(targetNode, targetWidget.name) ?? [ + targetWidget.type, + targetWidget.options || {} + ] + if (!isConvertibleWidget(targetWidget, config2)) return false + + const output = this.outputs[originSlot] + if (!(output.widget?.[CONFIG] ?? output.widget?.[GET_CONFIG]())) { + // No widget defined for this primitive yet so allow it + return true + } + + return !!mergeIfValid.call(this, output, config2) + } + #isValidConnection(input: INodeInputSlot, forceUpdate?: boolean) { // Only allow connections where the configs match const output = this.outputs[0] @@ -448,7 +468,11 @@ function showWidget(widget) { } } -function convertToInput(node: LGraphNode, widget: IWidget, config: InputSpec) { +export function convertToInput( + node: LGraphNode, + widget: IWidget, + config: InputSpec +) { hideWidget(node, widget) const { type } = getWidgetType(config) @@ -456,7 +480,7 @@ function convertToInput(node: LGraphNode, widget: IWidget, config: InputSpec) { // Add input and store widget config for creating on primitive node const sz = node.size const inputIsOptional = !!widget.options?.inputIsOptional - node.addInput(widget.name, type, { + const input = node.addInput(widget.name, type, { // @ts-expect-error GET_CONFIG is not defined in LiteGraph widget: { name: widget.name, [GET_CONFIG]: () => config }, // @ts-expect-error LiteGraph.SlotShape is not typed. @@ -469,6 +493,7 @@ function convertToInput(node: LGraphNode, widget: IWidget, config: InputSpec) { // Restore original size but grow if needed node.setSize([Math.max(sz[0], node.size[0]), Math.max(sz[1], node.size[1])]) + return input } function convertToWidget(node, widget) { diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 3a4a90e1..33e4646a 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -52,6 +52,7 @@ import { ModelStore, useModelStore } from '@/stores/modelStore' import type { ToastMessageOptions } from 'primevue/toast' import { useWorkspaceStore } from '@/stores/workspaceStateStore' import { useExecutionStore } from '@/stores/executionStore' +import { IWidget } from '@comfyorg/litegraph' export const ANIM_PREVIEW_WIDGET = '$$comfy_animation_preview' @@ -1694,6 +1695,52 @@ export class ComfyApp { } } + #addWidgetLinkHandling() { + app.canvas.getWidgetLinkType = function (widget, node) { + const nodeDefStore = useNodeDefStore() + const nodeDef = nodeDefStore.nodeDefsByName[node.type] + const input = nodeDef.input.getInput(widget.name) + return input?.type + } + + type ConnectingWidgetLink = { + subType: 'connectingWidgetLink' + widget: IWidget + node: LGraphNode + link: { node: LGraphNode; slot: number } + } + + document.addEventListener( + 'litegraph:canvas', + async (e: CustomEvent) => { + if (e.detail.subType === 'connectingWidgetLink') { + const { convertToInput } = await import( + '@/extensions/core/widgetInputs' + ) + + const { node, link, widget } = e.detail + if (!node || !link || !widget) return + + const nodeData = node.constructor.nodeData + if (!nodeData) return + const all = { + ...nodeData?.input?.required, + ...nodeData?.input?.optional + } + const inputSpec = all[widget.name] + if (!inputSpec) return + + const input = convertToInput(node, widget, inputSpec) + if (!input) return + + const originNode = link.node + + originNode.connect(link.slot, node, node.inputs.lastIndexOf(input)) + } + } + ) + } + #addAfterConfigureHandler() { const app = this // @ts-expect-error @@ -1915,6 +1962,7 @@ export class ComfyApp { this.#addCopyHandler() this.#addPasteHandler() this.#addKeyboardHandler() + this.#addWidgetLinkHandling() await this.#invokeExtensionsAsync('setup') }