From e2141a81e2b25dc2a1cddfc31e5dff3c135be3ca Mon Sep 17 00:00:00 2001 From: filtered <176114999+webfiltered@users.noreply.github.com> Date: Sun, 11 Aug 2024 23:46:54 +1000 Subject: [PATCH] Add support for LiteGraph to convert to classes (#334) * Add support for LiteGraph to convert to classes * Fix large context menu search regression * Remove debug code * Fix regression from rename & prototype change * Fix super() calls to match LGraphNode --- src/extensions/core/contextMenuFilter.ts | 26 ++++++++++++------------ src/extensions/core/noteNode.ts | 13 ++++++------ src/extensions/core/rerouteNode.ts | 5 +++-- src/extensions/core/widgetInputs.ts | 9 ++++---- src/scripts/app.ts | 21 ++++++++++--------- src/scripts/domWidget.ts | 1 - src/types/litegraph-augmentation.d.ts | 11 ++++++++++ 7 files changed, 49 insertions(+), 37 deletions(-) diff --git a/src/extensions/core/contextMenuFilter.ts b/src/extensions/core/contextMenuFilter.ts index 7c17ca632..625f26c84 100644 --- a/src/extensions/core/contextMenuFilter.ts +++ b/src/extensions/core/contextMenuFilter.ts @@ -8,19 +8,19 @@ const ext = { init() { const ctxMenu = LiteGraph.ContextMenu // @ts-expect-error - // TODO Very hacky way to modify Litegraph behaviour. Fix this later. + // TODO Very hacky way to modify Litegraph behaviour. Fix ctx later. LiteGraph.ContextMenu = function (values, options) { - const ctx = ctxMenu.call(this, values, options) + const ctx = new ctxMenu(values, options) // If we are a dark menu (only used for combo boxes) then add a filter input - if (options?.className === 'dark' && values?.length > 10) { + if (options?.className === 'dark' && values?.length > 4) { const filter = document.createElement('input') filter.classList.add('comfy-context-menu-filter') filter.placeholder = 'Filter list' - this.root.prepend(filter) + ctx.root.prepend(filter) const items = Array.from( - this.root.querySelectorAll('.litemenu-entry') + ctx.root.querySelectorAll('.litemenu-entry') ) as HTMLElement[] let displayedItems = [...items] let itemCount = displayedItems.length @@ -61,16 +61,16 @@ const ext = { } const positionList = () => { - const rect = this.root.getBoundingClientRect() + const rect = ctx.root.getBoundingClientRect() // If the top is off-screen then shift the element with scaling applied if (rect.top < 0) { const scale = 1 - - this.root.getBoundingClientRect().height / - this.root.clientHeight - const shift = (this.root.clientHeight * scale) / 2 - this.root.style.top = -shift + 'px' + ctx.root.getBoundingClientRect().height / + ctx.root.clientHeight + const shift = (ctx.root.clientHeight * scale) / 2 + ctx.root.style.top = -shift + 'px' } } @@ -109,7 +109,7 @@ const ext = { selectedItem?.click() break case 'Escape': - this.close() + ctx.close() break } }) @@ -140,7 +140,7 @@ const ext = { let top = options.event.clientY - 10 const bodyRect = document.body.getBoundingClientRect() - const rootRect = this.root.getBoundingClientRect() + const rootRect = ctx.root.getBoundingClientRect() if ( bodyRect.height && top > bodyRect.height - rootRect.height - 10 @@ -148,7 +148,7 @@ const ext = { top = Math.max(0, bodyRect.height - rootRect.height - 10) } - this.root.style.top = top + 'px' + ctx.root.style.top = top + 'px' positionList() } }) diff --git a/src/extensions/core/noteNode.ts b/src/extensions/core/noteNode.ts index 7d74c42f1..d2cd46839 100644 --- a/src/extensions/core/noteNode.ts +++ b/src/extensions/core/noteNode.ts @@ -1,32 +1,32 @@ import { LiteGraph, LGraphCanvas } from '@comfyorg/litegraph' import { app } from '../../scripts/app' import { ComfyWidgets } from '../../scripts/widgets' +import { LGraphNode } from '@comfyorg/litegraph' // Node that add notes to your project app.registerExtension({ name: 'Comfy.NoteNode', registerCustomNodes() { - class NoteNode { + class NoteNode extends LGraphNode { static category: string color = LGraphCanvas.node_colors.yellow.color bgcolor = LGraphCanvas.node_colors.yellow.bgcolor groupcolor = LGraphCanvas.node_colors.yellow.groupcolor - properties: { text: string } - serialize_widgets: boolean isVirtualNode: boolean collapsable: boolean title_mode: number - constructor() { + constructor(title?: string) { + super(title) if (!this.properties) { this.properties = { text: '' } } ComfyWidgets.STRING( - // @ts-expect-error - // Should we extends LGraphNode? + // Should we extends LGraphNode? Yesss this, '', + // @ts-expect-error ['', { default: this.properties.text, multiline: true }], app ) @@ -40,7 +40,6 @@ app.registerExtension({ LiteGraph.registerNodeType( 'Note', - // @ts-expect-error Object.assign(NoteNode, { title_mode: LiteGraph.NORMAL_TITLE, title: 'Note', diff --git a/src/extensions/core/rerouteNode.ts b/src/extensions/core/rerouteNode.ts index 0c1c4cbc0..45ce5ce8f 100644 --- a/src/extensions/core/rerouteNode.ts +++ b/src/extensions/core/rerouteNode.ts @@ -11,11 +11,12 @@ app.registerExtension({ __outputType?: string } - class RerouteNode { + class RerouteNode extends LGraphNode { static category: string | undefined static defaultVisibility = false - constructor() { + constructor(title?: string) { + super(title) if (!this.properties) { this.properties = {} } diff --git a/src/extensions/core/widgetInputs.ts b/src/extensions/core/widgetInputs.ts index 089d9765c..43eb77ca1 100644 --- a/src/extensions/core/widgetInputs.ts +++ b/src/extensions/core/widgetInputs.ts @@ -1,8 +1,8 @@ import { ComfyWidgets, addValueControlWidgets } from '../../scripts/widgets' import { app } from '../../scripts/app' import { applyTextReplacements } from '../../scripts/utils' -import { LiteGraph } from '@comfyorg/litegraph' -import type { LGraphNode, INodeInputSlot, IWidget } from '@comfyorg/litegraph' +import { LiteGraph, LGraphNode } from '@comfyorg/litegraph' +import type { INodeInputSlot, IWidget } from '@comfyorg/litegraph' const CONVERTED_TYPE = 'converted-widget' const VALID_TYPES = ['STRING', 'combo', 'number', 'toggle', 'BOOLEAN'] @@ -13,11 +13,12 @@ const TARGET = Symbol() // Used for reroutes to specify the real target widget interface PrimitiveNode extends LGraphNode {} const replacePropertyName = 'Run widget replace on values' -class PrimitiveNode { +class PrimitiveNode extends LGraphNode { controlValues: any[] lastType: string static category: string - constructor() { + constructor(title?: string) { + super(title) this.addOutput('connect to widget input', '*') this.serialize_widgets = true this.isVirtualNode = true diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 3e5ee5625..a647a1348 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -1983,8 +1983,15 @@ export class ComfyApp { async registerNodeDef(nodeId: string, nodeData: ComfyNodeDef) { const self = this - const node = Object.assign( - function ComfyNode() { + const node = class ComfyNode extends LGraphNode { + static comfyClass? = nodeData.name + // TODO: change to "title?" once litegraph.d.ts has been updated + static title = nodeData.display_name || nodeData.name + static nodeData? = nodeData + static category?: string + + constructor(title?: string) { + super(title) var inputs = nodeData['input']['required'] if (nodeData['input']['optional'] != undefined) { inputs = Object.assign( @@ -2050,13 +2057,9 @@ export class ComfyApp { this.serialize_widgets = true app.#invokeExtensionsAsync('nodeCreated', this) - }, - { - title: nodeData.display_name || nodeData.name, - comfyClass: nodeData.name, - nodeData } - ) + } + // @ts-expect-error node.prototype.comfyClass = nodeData.name this.#addNodeContextMenuHandler(node) @@ -2064,9 +2067,7 @@ export class ComfyApp { this.#addNodeKeyHandler(node) await this.#invokeExtensionsAsync('beforeRegisterNodeDef', node, nodeData) - // @ts-expect-error LiteGraph.registerNodeType(nodeId, node) - // @ts-expect-error node.category = nodeData.category } diff --git a/src/scripts/domWidget.ts b/src/scripts/domWidget.ts index a59306a7b..24e23bb42 100644 --- a/src/scripts/domWidget.ts +++ b/src/scripts/domWidget.ts @@ -254,7 +254,6 @@ export function addDomClippingSetting(): void { }) } -//@ts-ignore LGraphNode.prototype.addDOMWidget = function ( name: string, type: string, diff --git a/src/types/litegraph-augmentation.d.ts b/src/types/litegraph-augmentation.d.ts index 359f94870..74b4fe050 100644 --- a/src/types/litegraph-augmentation.d.ts +++ b/src/types/litegraph-augmentation.d.ts @@ -14,6 +14,13 @@ declare module '@comfyorg/litegraph' { * If the node is a frontend only node and should not be serialized into the prompt. */ isVirtualNode?: boolean + + addDOMWidget( + name: string, + type: string, + element: HTMLElement, + options: Record + ): DOMWidget } interface IWidget { @@ -67,4 +74,8 @@ declare module '@comfyorg/litegraph' { slotPos: Vector2 ): number } + + interface ContextMenu { + root?: HTMLDivElement + } }