import type { InjectionKey, MaybeRefOrGetter } from 'vue' import { computed, toValue } from 'vue' import { isProxyWidget } from '@/core/graph/subgraph/proxyWidget' import type { Positionable } from '@/lib/litegraph/src/interfaces' import type { LGraphGroup } from '@/lib/litegraph/src/LGraphGroup' import type { LGraphNode } from '@/lib/litegraph/src/LGraphNode' import type { SubgraphNode } from '@/lib/litegraph/src/subgraph/SubgraphNode' import type { IBaseWidget } from '@/lib/litegraph/src/types/widgets' import { isLGraphGroup, isLGraphNode } from '@/utils/litegraphUtil' export const GetNodeParentGroupKey: InjectionKey< (node: LGraphNode) => LGraphGroup | null > = Symbol('getNodeParentGroup') export type NodeWidgetsList = Array<{ node: LGraphNode; widget: IBaseWidget }> export type NodeWidgetsListList = Array<{ node: LGraphNode widgets: NodeWidgetsList }> /** * Searches widgets in a list and returns search results. * Filters by name, localized label, type, and user-input value. * Performs basic tokenization of the query string. */ export function searchWidgets( list: T, query: string ): T { if (query.trim() === '') { return list } const words = query.trim().toLowerCase().split(' ') return list.filter(({ widget }) => { const label = widget.label?.toLowerCase() const name = widget.name.toLowerCase() const type = widget.type.toLowerCase() const value = widget.value?.toString().toLowerCase() return words.every( (word) => name.includes(word) || label?.includes(word) || type?.includes(word) || value?.includes(word) ) }) as T } /** * Searches widgets and nodes in a list and returns search results. * First checks if the node title matches the query (if so, keeps entire node). * Otherwise, filters widgets using searchWidgets. * Performs basic tokenization of the query string. */ export function searchWidgetsAndNodes( list: NodeWidgetsListList, query: string ): NodeWidgetsListList { if (query.trim() === '') { return list } const words = query.trim().toLowerCase().split(' ') return list .map((item) => { const { node } = item const title = node.getTitle().toLowerCase() if (words.every((word) => title.includes(word))) { return { ...item, keep: true } } return { ...item, keep: false, widgets: searchWidgets(item.widgets, query) } }) .filter((item) => item.keep || item.widgets.length > 0) } type MixedSelectionItem = LGraphGroup | LGraphNode type FlatAndCategorizeSelectedItemsResult = { all: MixedSelectionItem[] nodes: LGraphNode[] groups: LGraphGroup[] others: Positionable[] nodeToParentGroup: Map } type FlatItemsContext = { nodeToParentGroup: Map depth: number parentGroup?: LGraphGroup } /** * The selected items may contain "Group" nodes, which can include child nodes. * This function flattens such structures and categorizes items into: * - all: all categorizable nodes (does not include nodes in "others") * - nodes: node items * - groups: group items * - others: items not currently supported * - nodeToParentGroup: a map from each node to its direct parent group (if any) * @param items The selected items to flatten and categorize * @returns An object containing arrays: all, nodes, groups, others, and nodeToParentGroup map */ export function flatAndCategorizeSelectedItems( items: Positionable[] ): FlatAndCategorizeSelectedItemsResult { const ctx: FlatItemsContext = { nodeToParentGroup: new Map(), depth: 0 } const { all, nodes, groups, others } = flatItems(items, ctx) return { all: repeatItems(all), nodes: repeatItems(nodes), groups: repeatItems(groups), others: repeatItems(others), nodeToParentGroup: ctx.nodeToParentGroup } } export function useFlatAndCategorizeSelectedItems( items: MaybeRefOrGetter ) { const result = computed(() => flatAndCategorizeSelectedItems(toValue(items))) return { flattedItems: computed(() => result.value.all), selectedNodes: computed(() => result.value.nodes), selectedGroups: computed(() => result.value.groups), selectedOthers: computed(() => result.value.others), nodeToParentGroup: computed(() => result.value.nodeToParentGroup) } } function flatItems( items: Positionable[], ctx: FlatItemsContext ): Omit { const result: MixedSelectionItem[] = [] const nodes: LGraphNode[] = [] const groups: LGraphGroup[] = [] const others: Positionable[] = [] if (ctx.depth > 1000) { return { all: [], nodes: [], groups: [], others: [] } } for (let i = 0; i < items.length; i++) { const item = items[i] as Positionable if (isLGraphGroup(item)) { result.push(item) groups.push(item) const children = Array.from(item.children) const childCtx: FlatItemsContext = { nodeToParentGroup: ctx.nodeToParentGroup, depth: ctx.depth + 1, parentGroup: item } const { all: childAll, nodes: childNodes, groups: childGroups, others: childOthers } = flatItems(children, childCtx) result.push(...childAll) nodes.push(...childNodes) groups.push(...childGroups) others.push(...childOthers) } else if (isLGraphNode(item)) { result.push(item) nodes.push(item) if (ctx.parentGroup) { ctx.nodeToParentGroup.set(item, ctx.parentGroup) } } else { // Other types of items are not supported yet // Do not add to all others.push(item) } } return { all: result, nodes, groups, others } } function repeatItems(items: T[]): T[] { const itemSet = new Set() const result: T[] = [] for (const item of items) { if (itemSet.has(item)) continue itemSet.add(item) result.push(item) } return result } /** * Renames a widget and its corresponding input. * Handles both regular widgets and proxy widgets in subgraphs. * * @param widget The widget to rename * @param node The node containing the widget * @param newLabel The new label for the widget (empty string or undefined to clear) * @param parents Optional array of parent SubgraphNodes (for proxy widgets) * @returns true if the rename was successful, false otherwise */ export function renameWidget( widget: IBaseWidget, node: LGraphNode, newLabel: string, parents?: SubgraphNode[] ): boolean { // For proxy widgets in subgraphs, we need to rename the original interior widget if (isProxyWidget(widget) && parents?.length) { const subgraph = parents[0].subgraph if (!subgraph) { console.error('Could not find subgraph for proxy widget') return false } const interiorNode = subgraph.getNodeById(parseInt(widget._overlay.nodeId)) if (!interiorNode) { console.error('Could not find interior node for proxy widget') return false } const originalWidget = interiorNode.widgets?.find( (w) => w.name === widget._overlay.widgetName ) if (!originalWidget) { console.error('Could not find original widget for proxy widget') return false } // Rename the original widget originalWidget.label = newLabel || undefined // Also rename the corresponding input on the interior node const interiorInput = interiorNode.inputs?.find( (inp) => inp.widget?.name === widget._overlay.widgetName ) if (interiorInput) { interiorInput.label = newLabel || undefined } } // Always rename the widget on the current node (either regular widget or proxy widget) const input = node.inputs?.find((inp) => inp.widget?.name === widget.name) // Intentionally mutate the widget object here as it's a reference // to the actual widget in the graph widget.label = newLabel || undefined if (input) { input.label = newLabel || undefined } return true }