diff --git a/__init__.py b/__init__.py index 1cea570..5764c46 100644 --- a/__init__.py +++ b/__init__.py @@ -13,6 +13,7 @@ from .py.lora_stack import RgthreeLoraLoaderStack from .py.seed import RgthreeSeed from .py.sdxl_empty_latent_image import RgthreeSDXLEmptyLatentImage from .py.power_prompt import RgthreePowerPrompt +from .py.power_prompt_simple import RgthreePowerPromptSimple NODE_CLASS_MAPPINGS = { RgthreeContext.NAME: RgthreeContext, @@ -22,6 +23,7 @@ NODE_CLASS_MAPPINGS = { RgthreeSeed.NAME: RgthreeSeed, RgthreeSDXLEmptyLatentImage.NAME: RgthreeSDXLEmptyLatentImage, RgthreePowerPrompt.NAME: RgthreePowerPrompt, + RgthreePowerPromptSimple.NAME: RgthreePowerPromptSimple, } THIS_DIR=os.path.dirname(os.path.abspath(__file__)) diff --git a/js/power_prompt.js b/js/power_prompt.js index 4c48708..040df77 100644 --- a/js/power_prompt.js +++ b/js/power_prompt.js @@ -1,70 +1,138 @@ import { app } from '../../scripts/app.js'; import { api } from '../../scripts/api.js'; +import { addConnectionLayoutSupport } from './utils.js'; class PowerPrompt { constructor(node, nodeData) { + this.combos = {}; + this.combosValues = {}; this.node = node; this.node.properties = this.node.properties || {}; this.nodeData = nodeData; + this.isSimple = this.nodeData.name.includes('Simple'); this.promptEl = node.widgets[0].inputEl; - this.refreshCombos(nodeData); - this.boundOnFreshNodeDefs = this.onFreshNodeDefs.bind(this); - api.addEventListener('fresh-node-defs', this.boundOnFreshNodeDefs); - const oldNodeRemoved = this.node.onRemoved; - this.node.onRemoved = () => { - oldNodeRemoved === null || oldNodeRemoved === void 0 ? void 0 : oldNodeRemoved.call(this.node); - api.removeEventListener('fresh-node-defs', this.boundOnFreshNodeDefs); + this.addAndHandleKeyboardLoraEditWeight(); + this.patchNodeRefresh(); + const oldOnConnectionsChange = this.node.onConnectionsChange; + this.node.onConnectionsChange = (type, slotIndex, isConnected, link_info, _ioSlot) => { + oldOnConnectionsChange === null || oldOnConnectionsChange === void 0 ? void 0 : oldOnConnectionsChange.apply(this.node, [type, slotIndex, isConnected, link_info, _ioSlot]); + this.onNodeConnectionsChange(type, slotIndex, isConnected, link_info, _ioSlot); }; + const oldOnConnectInput = this.node.onConnectInput; + this.node.onConnectInput = (inputIndex, outputType, outputSlot, outputNode, outputIndex) => { + let canConnect = true; + if (oldOnConnectInput) { + canConnect = oldOnConnectInput.apply(this.node, [inputIndex, outputType, outputSlot, outputNode, outputIndex]); + } + return canConnect && !this.node.inputs[inputIndex].disabled; + }; + const oldOnConnectOutput = this.node.onConnectOutput; + this.node.onConnectOutput = (outputIndex, inputType, inputSlot, inputNode, inputIndex) => { + let canConnect = true; + if (oldOnConnectOutput) { + canConnect = oldOnConnectOutput === null || oldOnConnectOutput === void 0 ? void 0 : oldOnConnectOutput.apply(this.node, [outputIndex, inputType, inputSlot, inputNode, inputIndex]); + } + return canConnect && !this.node.outputs[outputIndex].disabled; + }; + this.node.widgets.splice(1); + this.refreshCombos(nodeData); + setTimeout(() => { + this.stabilizeInputsOutputs(); + }, 32); + } + onNodeConnectionsChange(_type, _slotIndex, _isConnected, _linkInfo, _ioSlot) { + this.stabilizeInputsOutputs(); + } + stabilizeInputsOutputs() { + const clipLinked = this.node.inputs.some(i => i.name.includes('clip') && !!i.link); + const modelLinked = this.node.inputs.some(i => i.name.includes('model') && !!i.link); + for (const [index, output] of this.node.outputs.entries()) { + const type = output.type.toLowerCase(); + if (type.includes('model')) { + output.disabled = !modelLinked; + } + else if (type.includes('conditioning')) { + output.disabled = !clipLinked; + } + else if (type.includes('clip')) { + output.disabled = !clipLinked; + } + else if (type.includes('string')) { + output.color_off = '#7F7'; + output.color_on = '#7F7'; + } + if (output.disabled) { + } + } } onFreshNodeDefs(event) { this.refreshCombos(event.detail[this.nodeData.name]); } - refreshCombos(nodeData) { - var _a; - this.nodeData = nodeData; - for (const [key, value] of Object.entries(((_a = this.nodeData.input) === null || _a === void 0 ? void 0 : _a.hidden) || {})) { - if (key.includes('embedding') && Array.isArray(value[0])) { - const values = value[0]; - if (!this.embeddingWidget) { - this.embeddingWidget = this.node.addWidget('combo', key.replace(/_/g, ' '), values[0], (selected) => { - if (selected !== values[0]) { - this.insertText(`embedding:${selected}`); - } - this.embeddingWidget.value = values[0]; - }, { - values, - serialize: false, - }); - } - this.embeddingWidget.options.values = values; - this.embeddingWidget.value = values[0]; + findAndPatchCombos() { + } + onPromptComboCallback(widget, selected) { + const values = widget.options.values; + if (selected !== values[0] && !selected.match(/^disable\s[a-z]/i)) { + if (widget.name.includes('embedding')) { + this.insertText(`embedding:${selected}`); } - else if (key.includes('saved') && Array.isArray(value[0])) { + else if (widget.name.includes('saved')) { + this.insertText(this.combosValues[`saved_${widget.name}`][values.indexOf(selected)]); + } + else if (widget.name.includes('lora')) { + this.insertText(``); + } + } + } + refreshCombos(nodeData) { + var _a, _b; + this.nodeData = nodeData; + let data = ((_a = this.nodeData.input) === null || _a === void 0 ? void 0 : _a.optional) || {}; + data = Object.assign(data, ((_b = this.nodeData.input) === null || _b === void 0 ? void 0 : _b.hidden) || {}); + for (const [key, value] of Object.entries(data)) { + if (Array.isArray(value[0])) { const values = value[0]; - if (values.length <= 1) { - this.savedValues = []; - if (this.savedWidget) { - this.node.widgets.splice(this.node.widgets.indexOf(this.savedWidget), 1); - this.savedWidget = undefined; + if (key.startsWith('insert')) { + const shouldShow = values.length > 2 || (values.length > 1 && !values[1].match(/^disable\s[a-z]/i)); + if (shouldShow) { + if (!this.combos[key]) { + this.combos[key] = this.node.addWidget('combo', key, values, (selected) => { + if (selected !== values[0] && !selected.match(/^disable\s[a-z]/i)) { + if (key.includes('embedding')) { + this.insertText(`embedding:${selected}`); + } + else if (key.includes('saved')) { + this.insertText(this.combosValues[`values_${key}`][values.indexOf(selected)]); + } + else if (key.includes('lora')) { + this.insertText(``); + } + this.combos[key].value = values[0]; + } + }, { + values, + serialize: true, + }); + this.combos[key].oldComputeSize = this.combos[key].computeSize; + let node = this.node; + this.combos[key].computeSize = function (width) { + var _a, _b; + const size = ((_b = (_a = this).oldComputeSize) === null || _b === void 0 ? void 0 : _b.call(_a, width)) || [width, LiteGraph.NODE_WIDGET_HEIGHT]; + if (this === node.widgets[node.widgets.length - 1]) { + size[1] += 10; + } + return size; + }; + } + this.combos[key].options.values = values; + this.combos[key].value = values[0]; } - continue; - } - if (key.startsWith('values')) { - this.savedValues = values; - } - else { - if (!this.savedWidget) { - this.savedWidget = this.node.addWidget('combo', key.replace(/_/g, ' '), values[0], (selected) => { - if (selected !== values[0]) { - this.insertText(this.savedValues[values.indexOf(selected)]); - } - this.savedWidget.value = values[0]; - }, { - values, - serialize: false, - }); + else if (!shouldShow && this.combos[key]) { + this.node.widgets.splice(this.node.widgets.indexOf(this.combos[key]), 1); + delete this.combos[key]; } - this.savedWidget.options.values = values; - this.savedWidget.value = values[0]; + } + else if (key.startsWith('values')) { + this.combosValues[key] = values; } } } @@ -82,16 +150,89 @@ class PowerPrompt { this.promptEl.selectionEnd = first.length + text.length; } } + addAndHandleKeyboardLoraEditWeight() { + this.promptEl.addEventListener('keydown', (event) => { + var _a, _b; + if (!(event.key === "ArrowUp" || event.key === "ArrowDown")) + return; + if (!event.ctrlKey && !event.metaKey) + return; + const delta = event.shiftKey ? .01 : .1; + let start = this.promptEl.selectionStart; + let end = this.promptEl.selectionEnd; + let fullText = this.promptEl.value; + let selectedText = fullText.substring(start, end); + if (!selectedText) { + const stopOn = "<>() \r\n\t"; + if (fullText[start] == '>') { + start -= 2; + end -= 2; + } + if (fullText[end - 1] == '<') { + start += 2; + end += 2; + } + while (!stopOn.includes(fullText[start]) && start > 0) { + start--; + } + while (!stopOn.includes(fullText[end - 1]) && end < fullText.length) { + end++; + } + selectedText = fullText.substring(start, end); + } + if (!selectedText.startsWith('')) { + return; + } + let weight = (_b = Number((_a = selectedText.match(/:(-?\d*(\.\d*)?)>$/)) === null || _a === void 0 ? void 0 : _a[1])) !== null && _b !== void 0 ? _b : 1; + weight += event.key === "ArrowUp" ? delta : -delta; + const updatedText = selectedText.replace(/(:-?\d*(\.\d*)?)?>$/, `:${weight.toFixed(2)}>`); + this.promptEl.setRangeText(updatedText, start, end, 'select'); + event.preventDefault(); + event.stopPropagation(); + }); + } + patchNodeRefresh() { + this.boundOnFreshNodeDefs = this.onFreshNodeDefs.bind(this); + api.addEventListener('fresh-node-defs', this.boundOnFreshNodeDefs); + const oldNodeRemoved = this.node.onRemoved; + this.node.onRemoved = () => { + oldNodeRemoved === null || oldNodeRemoved === void 0 ? void 0 : oldNodeRemoved.call(this.node); + api.removeEventListener('fresh-node-defs', this.boundOnFreshNodeDefs); + }; + } } +let nodeData = null; app.registerExtension({ name: 'rgthree.PowerPrompt', - async beforeRegisterNodeDef(nodeType, nodeData, _app) { - if (nodeData.name === 'Power Prompt (rgthree)') { + async beforeRegisterNodeDef(nodeType, passedNodeData, _app) { + if (passedNodeData.name.startsWith('Power Prompt') && passedNodeData.name.includes('rgthree')) { + nodeData = passedNodeData; const onNodeCreated = nodeType.prototype.onNodeCreated; nodeType.prototype.onNodeCreated = function () { onNodeCreated ? onNodeCreated.apply(this, []) : undefined; - this.powerPrompt = new PowerPrompt(this, nodeData); + this.powerPrompt = new PowerPrompt(this, passedNodeData); }; + addConnectionLayoutSupport(nodeType, app, [['Left', 'Right'], ['Right', 'Left']]); } }, + async loadedGraphNode(node) { + if (node.type === 'Power Prompt (rgthree)') { + setTimeout(() => { + if (node.outputs[0].type === 'STRING') { + if (node.outputs[0].links) { + node.outputs[3].links = node.outputs[3].links || []; + for (const link of node.outputs[0].links) { + node.outputs[3].links.push(link); + app.graph.links[link].origin_slot = 3; + } + node.outputs[0].links = null; + } + node.outputs[0].type = nodeData.output[0]; + node.outputs[0].name = nodeData.output_name[0] || node.outputs[0].type; + node.outputs[0].color_on = undefined; + node.outputs[0].color_off = undefined; + } + }, 50); + } + } }); diff --git a/js/utils.js b/js/utils.js index 5bacbdb..85491e8 100644 --- a/js/utils.js +++ b/js/utils.js @@ -105,11 +105,28 @@ export function getConnectionPosForLayout(node, isInput, slotNumber, out) { const offset = (_a = node.constructor.layout_slot_offset) !== null && _a !== void 0 ? _a : (LiteGraph.NODE_SLOT_HEIGHT * 0.5); const side = isInput ? layout[0] : layout[1]; const data = LAYOUT_LABEL_TO_DATA[side]; - const cxn = node[isInput ? 'inputs' : 'outputs'][slotNumber]; + const slotList = node[isInput ? 'inputs' : 'outputs']; + const cxn = slotList[slotNumber]; if (!cxn) { console.log('No connection found.. weird', isInput, slotNumber); return out; } + if (cxn.disabled) { + if (cxn.color_on !== '#666665') { + cxn._color_on_org = cxn._color_on_org || cxn.color_on; + cxn._color_off_org = cxn._color_off_org || cxn.color_off; + } + cxn.color_on = '#666665'; + cxn.color_off = '#666665'; + } + else if (cxn.color_on === '#666665') { + cxn.color_on = cxn._color_on_org || undefined; + cxn.color_off = cxn._color_off_org || undefined; + } + const displaySlot = collapseConnections ? 0 : (slotNumber - slotList.reduce((count, ioput, index) => { + count += index < slotNumber && ioput.hidden ? 1 : 0; + return count; + }, 0)); cxn.dir = data[0]; if (side === 'Left') { if (node.flags.collapsed) { @@ -133,7 +150,6 @@ export function getConnectionPosForLayout(node, isInput, slotNumber, out) { out[1] = node.pos[1] + (node.size[1] * .5); } else { - const displaySlot = collapseConnections ? 0 : slotNumber; out[1] = node.pos[1] + (displaySlot + 0.7) * LiteGraph.NODE_SLOT_HEIGHT + @@ -163,7 +179,6 @@ export function getConnectionPosForLayout(node, isInput, slotNumber, out) { out[1] = node.pos[1] + (node.size[1] * .5); } else { - const displaySlot = collapseConnections ? 0 : slotNumber; out[1] = node.pos[1] + (displaySlot + 0.7) * LiteGraph.NODE_SLOT_HEIGHT + diff --git a/py/log.py b/py/log.py index 3ad6d74..fc9e95f 100644 --- a/py/log.py +++ b/py/log.py @@ -1,22 +1,22 @@ # https://stackoverflow.com/questions/4842424/list-of-ansi-color-escape-sequences +# https://en.wikipedia.org/wiki/ANSI_escape_code#3-bit_and_4-bit COLORS_FG = { 'BLACK': '\33[30m', 'RED': '\33[31m', 'GREEN': '\33[32m', 'YELLOW': '\33[33m', 'BLUE': '\33[34m', - 'VIOLET': '\33[35m', - 'BEIGE': '\33[36m', + 'MAGENTA': '\33[35m', + 'CYAN': '\33[36m', 'WHITE': '\33[37m', - 'GREY': '\33[90m', - 'LIGHTRED': '\33[91m', - 'LIGHTGREEN': '\33[92m', - 'LIGHTYELLOW': '\33[93m', - 'LIGHTBLUE': '\33[94m', - 'LIGHTVIOLET': '\33[95m', - 'LIGHTBEIGE': '\33[96m', - 'LIGHTWHITE': '\33[97m', + 'BRIGHT_RED': '\33[91m', + 'BRIGHT_GREEN': '\33[92m', + 'BRIGHT_YELLOW': '\33[93m', + 'BRIGHT_BLUE': '\33[94m', + 'BRIGHT_MAGENTA': '\33[95m', + 'BRIGHT_CYAN': '\33[96m', + 'BRIGHT_WHITE': '\33[97m', } COLORS_STYLE = { 'RESET': '\33[0m', @@ -29,23 +29,22 @@ COLORS_STYLE = { 'SELECTED': '\33[7m', } COLORS_BG = { - 'BLACKBG': '\33[40m', - 'REDBG': '\33[41m', - 'GREENBG': '\33[42m', - 'YELLOWBG': '\33[43m', - 'BLUEBG': '\33[44m', - 'VIOLETBG': '\33[45m', - 'BEIGEBG': '\33[46m', - 'WHITEBG': '\33[47m', - - 'GREYBG': '\33[100m', - 'LIGHTREDBG': '\33[101m', - 'LIGHTGREENBG': '\33[102m', - 'LIGHTYELLOWBG': '\33[103m', - 'LIGHTBLUEBG': '\33[104m', - 'LIGHTVIOLETBG': '\33[105m', - 'LIGHTBEIGEBG': '\33[106m', - 'LIGHTWHITEBG': '\33[107m', + 'BLACK': '\33[40m', + 'RED': '\33[41m', + 'GREEN': '\33[42m', + 'YELLOW': '\33[43m', + 'BLUE': '\33[44m', + 'MAGENTA': '\33[45m', + 'CYAN': '\33[46m', + 'WHITE': '\33[47m', + 'GREY': '\33[100m', + 'BRIGHT_RED': '\33[101m', + 'BRIGHT_GREEN': '\33[102m', + 'BRIGHT_YELLOW': '\33[103m', + 'BRIGHT_BLUE': '\33[104m', + 'BRIGHT_MAGENTA': '\33[105m', + 'BRIGHT_CYAN': '\33[106m', + 'BRIGHT_WHITE': '\33[107m', } def log_welcome(num_nodes=None): @@ -54,3 +53,18 @@ def log_welcome(num_nodes=None): print('{} {} exciting nodes.{}'.format(msg, num_nodes, COLORS_STYLE['RESET'])) else: print('{}.{}'.format(msg, COLORS_STYLE['RESET'])) + +def log_node_success(node_name, message): + _log_node(COLORS_FG["GREEN"], node_name, message, prefix='✓ ') + +def log_node_info(node_name, message): + _log_node(COLORS_FG["CYAN"], node_name, message, prefix='🛈 ') + +def log_node(node_name, message): + _log_node(COLORS_FG["CYAN"], node_name, message, prefix=' ') + +def log_node_warn(node_name, message): + _log_node(COLORS_FG["YELLOW"], node_name, message, prefix='⚠ ') + +def _log_node(color, node_name, message, prefix=''): + print(f'{COLORS_STYLE["BOLD"]}{color}{prefix}rgthree {node_name.replace(" (rgthree)", "")}:{COLORS_STYLE["RESET"]} {message}') diff --git a/py/power_prompt.py b/py/power_prompt.py index 47cf1bb..514ba97 100644 --- a/py/power_prompt.py +++ b/py/power_prompt.py @@ -1,14 +1,54 @@ import os -import inspect +import re + +from .log import log_node_warn, log_node_info, log_node_success from .constants import get_category, get_name -from nodes import LoraLoader +from nodes import LoraLoader, CLIPTextEncode import folder_paths +NODE_NAME=get_name('Power Prompt') + +def get_and_strip_loras(prompt, silent=False): + pattern=']*?)(?::(-?\d*(?:\.\d*)?))?>' + lora_paths=folder_paths.get_filename_list('loras') + lora_filenames_no_ext=[os.path.splitext(os.path.basename(x))[0] for x in lora_paths] + + matches = re.findall(pattern, prompt) + + loras=[] + for match in matches: + tag_filename=match[0] + strength=float(match[1] if len(match) > 1 and len(match[1]) else 1.0) + if strength == 0 and not silent: + log_node_info(NODE_NAME, f'Skipping "{tag_filename}" with strength of zero') + + # Let's be flexible. If the lora filename in the tag doesn't have the extension or + # path prefix, let's still find and load it. + if tag_filename not in lora_paths: + found_tag_filename=None + for index, value in enumerate(lora_filenames_no_ext): + if value in tag_filename: + found_tag_filename=lora_paths[index] + break + if found_tag_filename: + # if not silent: + # log_node_info(NODE_NAME, f'Found "{found_tag_filename}" for "{tag_filename}" in prompt') + tag_filename=found_tag_filename + else: + if not silent: + log_node_warn(NODE_NAME, f'Lora "{tag_filename}" not found, skipping.') + continue + + loras.append({'lora':tag_filename, 'strength':strength}) + + return (re.sub(pattern, '', prompt), loras) + + class RgthreePowerPrompt: - NAME = get_name('Power Prompt') + NAME = NODE_NAME CATEGORY = get_category() @classmethod @@ -22,17 +62,43 @@ class RgthreePowerPrompt: 'required': { 'prompt': ('STRING', {'multiline': True}), }, + 'optional': { + "opt_model": ("MODEL",), + "opt_clip": ("CLIP", ), + 'insert_lora': (['CHOOSE', 'DISABLE LORAS'] + [os.path.splitext(x)[0] for x in folder_paths.get_filename_list('loras')],), + 'insert_embedding': (['CHOOSE',] + [os.path.splitext(x)[0] for x in folder_paths.get_filename_list('embeddings')],), + 'insert_saved': (['CHOOSE',] + SAVED_PROMPTS_FILES,), + }, 'hidden': { - 'insert_embedding': (['Choose'] + folder_paths.get_filename_list('embeddings'),), - 'insert_saved': (['Choose'] + SAVED_PROMPTS_FILES,), - 'values_insert_saved': (['Choose'] + SAVED_PROMPTS_CONTENT,), + 'values_insert_saved': (['CHOOSE'] + SAVED_PROMPTS_CONTENT,), } } - RETURN_TYPES = ('STRING',) - RETURN_NAMES = ('TEXT',) + RETURN_TYPES = ('CONDITIONING', 'MODEL', 'CLIP', 'STRING',) + RETURN_NAMES = ('CONDITIONING', 'MODEL', 'CLIP', 'TEXT',) FUNCTION = 'main' - def main(self, prompt, insert_embedding=None, insert_saved=None, values_insert_saved=None): - return (prompt,) + def main(self, prompt, opt_model=None, opt_clip=None, insert_lora=None, insert_embedding=None, insert_saved=None, values_insert_saved=None): + if insert_lora == 'DISABLE LORAS': + prompt, loras = get_and_strip_loras(prompt, True) + log_node_info(NODE_NAME, f'Disabling all found loras ({len(loras)}) and stripping lora tags for TEXT output.') + elif opt_model != None and opt_clip != None: + prompt, loras = get_and_strip_loras(prompt) + if len(loras): + for lora in loras: + opt_model, opt_clip = LoraLoader().load_lora(opt_model, opt_clip, lora['lora'], lora['strength'], lora['strength']) + log_node_success(NODE_NAME, f'Loaded "{lora["lora"]}" from prompt') + log_node_info(NODE_NAME, f'{len(loras)} Loras processed; stripping tags for TEXT output.') + elif ' void; + readonly combos: {[key:string]: IComboWidget} = {}; + readonly combosValues: {[key:string]: string[]} = {}; + boundOnFreshNodeDefs!: (event: CustomEvent) => void; constructor(node: ComfyGraphNode, nodeData: ComfyObjectInfo) { this.node = node; this.node.properties = this.node.properties || {}; this.nodeData = nodeData; + this.isSimple = this.nodeData.name.includes('Simple'); this.promptEl = (node.widgets[0]! as any).inputEl; + this.addAndHandleKeyboardLoraEditWeight(); + + // this.findAndPatchCombos(); + this.patchNodeRefresh(); + + const oldOnConnectionsChange = this.node.onConnectionsChange; + this.node.onConnectionsChange = (type: number, slotIndex: number, isConnected: boolean, link_info: LLink, _ioSlot: (INodeOutputSlot | INodeInputSlot)) => { + oldOnConnectionsChange?.apply(this.node, [type, slotIndex, isConnected, link_info,_ioSlot]); + this.onNodeConnectionsChange(type, slotIndex, isConnected, link_info,_ioSlot); + } + + const oldOnConnectInput = this.node.onConnectInput; + this.node.onConnectInput = (inputIndex: number, outputType: INodeOutputSlot["type"], outputSlot: INodeOutputSlot, outputNode: TLGraphNode, outputIndex: number) => { + let canConnect = true; + if (oldOnConnectInput) { + canConnect = oldOnConnectInput.apply(this.node, [inputIndex, outputType, outputSlot, outputNode,outputIndex]); + } + return canConnect && !this.node.inputs[inputIndex]!.disabled; + } + + const oldOnConnectOutput = this.node.onConnectOutput; + this.node.onConnectOutput = (outputIndex: number, inputType: INodeInputSlot["type"], inputSlot: INodeInputSlot, inputNode: TLGraphNode, inputIndex: number) => { + let canConnect = true; + if (oldOnConnectOutput) { + canConnect = oldOnConnectOutput?.apply(this.node, [outputIndex, inputType, inputSlot, inputNode, inputIndex]); + } + return canConnect && !this.node.outputs[outputIndex]!.disabled; + } + + // Strip all widgets but prompt (we'll re-add them in refreshCombos) + this.node.widgets.splice(1); this.refreshCombos(nodeData); + setTimeout(()=> { + this.stabilizeInputsOutputs(); + }, 32); + } - // We patched over api.getNodeDefs in utils.js to fire a custom event that we can not - // listen to here to manually refresh our combos when a request comes in to fetch the - // node data; which only happens one at startup (but before custom nodes js runs), and - // then after clicking the "Refresh" button in the floating menu. - this.boundOnFreshNodeDefs = this.onFreshNodeDefs.bind(this); - api.addEventListener('fresh-node-defs', this.boundOnFreshNodeDefs); - const oldNodeRemoved = this.node.onRemoved; - this.node.onRemoved = () => { - oldNodeRemoved?.call(this.node); - api.removeEventListener('fresh-node-defs', this.boundOnFreshNodeDefs); + /** + * Cleans up optional out puts when we don't have the optional input. Purely a vanity function. + */ + onNodeConnectionsChange(_type: number, _slotIndex: number, _isConnected: boolean, _linkInfo: LLink, _ioSlot: (INodeOutputSlot | INodeInputSlot)) { + this.stabilizeInputsOutputs(); + } + + private stabilizeInputsOutputs() { + // If our first input is connected, then we can show the proper output. + const clipLinked = this.node.inputs.some(i=>i.name.includes('clip') && !!i.link); + const modelLinked = this.node.inputs.some(i=>i.name.includes('model') && !!i.link); + for (const [index, output] of this.node.outputs.entries()) { + const type = (output.type as string).toLowerCase(); + if (type.includes('model')) { + output.disabled = !modelLinked; + } else if (type.includes('conditioning')) { + output.disabled = !clipLinked; + } else if (type.includes('clip')) { + output.disabled = !clipLinked; + } else if (type.includes('string')) { + // Our text prompt is always enabled, but let's color it so it stands out + // if the others are disabled. #7F7 is Litegraph's default. + output.color_off = '#7F7'; + output.color_on = '#7F7'; + } + if (output.disabled) { + // this.node.disconnectOutput(index); + } } } @@ -49,53 +104,91 @@ class PowerPrompt { this.refreshCombos(event.detail[this.nodeData.name]); } - refreshCombos(nodeData: ComfyObjectInfo) { - this.nodeData = nodeData; - // Add the combo for embeddings, in hidden inputs of nodeData - for (const [key, value] of Object.entries(this.nodeData.input?.hidden || {})) { - if (key.includes('embedding') && Array.isArray(value[0])) { - const values = value[0]; - if (!this.embeddingWidget) { - this.embeddingWidget = this.node.addWidget('combo', key.replace(/_/g, ' '), values[0], (selected) => { - if (selected !== values[0]) { - this.insertText(`embedding:${selected}`); - } - this.embeddingWidget!.value = values[0]; - }, { - values, - serialize: false, // Don't include this in prompt. - }); - } - this.embeddingWidget.options.values = values; - this.embeddingWidget.value = values[0]; + findAndPatchCombos() { + // for (const widget of this.node.widgets) { + // if (widget.type === 'combo' && widget.name!.startsWith('insert_')) { + // widget.callback = (selected) => this.onPromptComboCallback(widget as IComboWidget, selected); + // if (widget.options.values.length === 1) { + // widget.disabled = true; + // } + // // Override comput size so we can add some padding after the last widget. Not sure why it's + // // funky, perhaps the multiline text area. + // (widget as any).oldComputeSize = widget.computeSize; + // let node = this.node; + // widget.computeSize = function(width: number) { + // const size = (this as any).oldComputeSize?.(width) || [width, LiteGraph.NODE_WIDGET_HEIGHT]; + // if (this === node.widgets[node.widgets.length- 1]) { + // size[1] += 10; + // } + // return size; + // }; + // } + // } + } - } else if (key.includes('saved') && Array.isArray(value[0])) { - const values = value[0]; - // If all we have is the "Choose" option and no prompts configured, then skip and don't show. - if (values.length <= 1) { - this.savedValues = []; - if (this.savedWidget) { - this.node.widgets.splice(this.node.widgets.indexOf(this.savedWidget as IWidget), 1); - this.savedWidget = undefined; + onPromptComboCallback(widget: IComboWidget, selected: string) { + const values = widget.options.values as string[]; + if (selected !== values[0] && !selected.match(/^disable\s[a-z]/i)) { + if (widget.name!.includes('embedding')) { + this.insertText(`embedding:${selected}`); + } else if (widget.name!.includes('saved')) { + this.insertText(this.combosValues[`saved_${widget.name!}`]![values.indexOf(selected)]!); + } else if (widget.name!.includes('lora')) { + this.insertText(``); + } + } + } + + + + refreshCombos(nodeData: ComfyObjectInfo) { + + this.nodeData = nodeData; + // Add the combo for hidden inputs of nodeData + let data = this.nodeData.input?.optional || {}; + data = Object.assign(data, this.nodeData.input?.hidden || {}); + + for (const [key, value] of Object.entries(data)) {//Object.entries(this.nodeData.input?.hidden || {})) { + if (Array.isArray(value[0])) { + const values = value[0] as string[]; + if (key.startsWith('insert')) { + const shouldShow = values.length > 2 || (values.length > 1 && !values[1]!.match(/^disable\s[a-z]/i)) + if (shouldShow) { + if (!this.combos[key]) { + this.combos[key] = this.node.addWidget('combo', key, values, (selected) => { + if (selected !== values[0] && !selected.match(/^disable\s[a-z]/i)) { + if (key.includes('embedding')) { + this.insertText(`embedding:${selected}`); + } else if (key.includes('saved')) { + this.insertText(this.combosValues[`values_${key}`]![values.indexOf(selected)]!); + } else if (key.includes('lora')) { + this.insertText(``); + } + this.combos[key]!.value = values[0]; + } + }, { + values, + serialize: true, // Don't include this in prompt. + }); + (this.combos[key]! as any).oldComputeSize = this.combos[key]!.computeSize; + let node = this.node; + this.combos[key]!.computeSize = function(width: number) { + const size = (this as any).oldComputeSize?.(width) || [width, LiteGraph.NODE_WIDGET_HEIGHT]; + if (this === node.widgets[node.widgets.length- 1]) { + size[1] += 10; + } + return size; + }; + } + this.combos[key]!.options.values = values; + this.combos[key]!.value = values[0]; + } else if (!shouldShow && this.combos[key]) { + this.node.widgets.splice(this.node.widgets.indexOf(this.combos[key]!), 1); + delete this.combos[key]; } - continue; - } - if (key.startsWith('values')) { - this.savedValues = values; - } else { - if (!this.savedWidget) { - this.savedWidget = this.node.addWidget('combo', key.replace(/_/g, ' '), values[0], (selected) => { - if (selected !== values[0]) { - this.insertText(this.savedValues![values.indexOf(selected)]!); - } - this.savedWidget!.value = values[0]; - }, { - values, - serialize: false, // Don't include this in prompt. - }); - } - this.savedWidget.options.values = values; - this.savedWidget.value = values[0]; + + } else if (key.startsWith('values')) { + this.combosValues[key] = values; } } } @@ -114,17 +207,93 @@ class PowerPrompt { this.promptEl.selectionEnd = first.length + text.length; } } + + /** + * Adds a keydown event listener to our prompt so we can see if we're using the + * ctrl/cmd + up/down arrows shortcut. This kind of competes with the core extension + * "Comfy.EditAttention" but since that only handles parenthesis and listens on window, we should + * be able to intercept and cancel the bubble if we're doing the same action within the lora tag. + */ + addAndHandleKeyboardLoraEditWeight() { + this.promptEl.addEventListener('keydown', (event: KeyboardEvent)=> { + // If we're not doing a ctrl/cmd + arrow key, then bail. + if (!(event.key === "ArrowUp" || event.key === "ArrowDown")) return; + if (!event.ctrlKey && !event.metaKey) return; + // Unfortunately, we can't see Comfy.EditAttention delta in settings, so we hardcode to 0.01. + // We can acutally do better too, let's make it .1 by default, and .01 if also holding shift. + const delta = event.shiftKey ? .01 : .1; + + let start = this.promptEl.selectionStart; + let end = this.promptEl.selectionEnd; + let fullText = this.promptEl.value; + let selectedText = fullText.substring(start, end); + + // We don't care about fully rewriting Comfy.EditAttention, we just want to see if our + // selected text is a lora, which will always start with "') { + start-=2; + end-=2; + } + if (fullText[end-1] == '<') { + start+=2; + end+=2; + } + while (!stopOn.includes(fullText[start]!) && start > 0) { + start--; + } + while (!stopOn.includes(fullText[end-1]!) && end < fullText.length) { + end++; + } + selectedText = fullText.substring(start, end); + } + + // Bail if this isn't a lora. + if (!selectedText.startsWith('')) { + return; + } + + let weight = Number(selectedText.match(/:(-?\d*(\.\d*)?)>$/)?.[1]) ?? 1; + weight += event.key === "ArrowUp" ? delta : -delta; + const updatedText = selectedText.replace(/(:-?\d*(\.\d*)?)?>$/, `:${weight.toFixed(2)}>`); + + // Handle the new value and cancel the bubble so Comfy.EditAttention doesn't also try. + this.promptEl.setRangeText(updatedText, start, end, 'select'); + event.preventDefault(); + event.stopPropagation(); + }); + } + + /** + * Patches over api.getNodeDefs in comfy's api.js to fire a custom event that we can listen to + * here and manually refresh our combos when a request comes in to fetch the node data; which + * only happens once at startup (but before custom nodes js runs), and then after clicking + * the "Refresh" button in the floating menu, which is what we care about. + */ + patchNodeRefresh() { + this.boundOnFreshNodeDefs = this.onFreshNodeDefs.bind(this); + api.addEventListener('fresh-node-defs', this.boundOnFreshNodeDefs); + const oldNodeRemoved = this.node.onRemoved; + this.node.onRemoved = () => { + oldNodeRemoved?.call(this.node); + api.removeEventListener('fresh-node-defs', this.boundOnFreshNodeDefs); + } + } } +let nodeData: ComfyObjectInfo | null = null; app.registerExtension({ name: 'rgthree.PowerPrompt', - async beforeRegisterNodeDef(nodeType: typeof LGraphNode, nodeData: ComfyObjectInfo, _app: ComfyApp) { - if (nodeData.name === 'Power Prompt (rgthree)') { + async beforeRegisterNodeDef(nodeType: typeof LGraphNode, passedNodeData: ComfyObjectInfo, _app: ComfyApp) { + if (passedNodeData.name.startsWith('Power Prompt') && passedNodeData.name.includes('rgthree')) { + nodeData = passedNodeData; const onNodeCreated = nodeType.prototype.onNodeCreated; nodeType.prototype.onNodeCreated = function () { onNodeCreated ? onNodeCreated.apply(this, []) : undefined; - (this as any).powerPrompt = new PowerPrompt(this as ComfyGraphNode, nodeData); + (this as any).powerPrompt = new PowerPrompt(this as ComfyGraphNode, passedNodeData); } // This won't actually work until such a thing exists in app.js#refreshComboInNodes @@ -132,6 +301,32 @@ app.registerExtension({ // nodeType.prototype.onRefreshCombos = function (newNodeData: any) { // (this as any).powerPrompt.refreshCombos(newNodeData); // } + + // This isn't super useful, because R->L removes the names in order to work with + // litegraph's hardcoded L->R math.. but, ¯\_(ツ)_/¯ + addConnectionLayoutSupport(nodeType, app, [['Left', 'Right'], ['Right', 'Left']]); } }, + async loadedGraphNode(node: TLGraphNode) { + if (node.type === 'Power Prompt (rgthree)') { + setTimeout(() => { + // If the first output is STRING, then it's the text output from the initial launch. + // Let's port it to the new + if (node.outputs[0]!.type === 'STRING') { + if (node.outputs[0]!.links) { + node.outputs[3]!.links = node.outputs[3]!.links || []; + for (const link of node.outputs[0]!.links) { + node.outputs[3]!.links.push(link); + app.graph.links[link].origin_slot = 3; + } + node.outputs[0]!.links = null; + } + node.outputs[0]!.type = nodeData!.output![0] as string; + node.outputs[0]!.name = nodeData!.output_name![0] || node.outputs[0]!.type as string; + node.outputs[0]!.color_on = undefined; + node.outputs[0]!.color_off = undefined; + } + }, 50) + } + } }); \ No newline at end of file diff --git a/ts/typings/litegraph.d.ts b/ts/typings/litegraph.d.ts index 4158d39..5e37407 100644 --- a/ts/typings/litegraph.d.ts +++ b/ts/typings/litegraph.d.ts @@ -34,6 +34,10 @@ export interface INodeSlot { locked?: boolean; nameLocked?: boolean; pos?: Vector2; + // @rgthree + hidden?: boolean; + // @rgthree + disabled?: boolean; } export interface INodeInputSlot extends INodeSlot { @@ -52,6 +56,16 @@ export type WidgetCallback = ( event?: MouseEvent ) => void; +// #rgthree +export type WidgetComboCallback = ( + this: T, + value: T["value"][0], + graphCanvas: LGraphCanvas, + node: LGraphNode, + pos: Vector2, + event?: MouseEvent +) => void; + export interface IWidget { name: string | null; value: TValue; @@ -107,7 +121,9 @@ export interface IComboWidget | ((widget: IComboWidget, node: LGraphNode) => string[]); } | TOptions > { + value: T[0]; type: "combo"; + callback?: WidgetComboCallback; } export interface ITextWidget extends IWidget { @@ -817,7 +833,8 @@ export declare class LGraphNode { type: T["type"], name: string, value: T["value"], - callback?: WidgetCallback | string, + // @rgthree + callback?: T["callback"] | string,//WidgetCallback | string, options?: T["options"] ): T; diff --git a/ts/utils.ts b/ts/utils.ts index dc2cea0..829c0b0 100644 --- a/ts/utils.ts +++ b/ts/utils.ts @@ -144,11 +144,36 @@ export function getConnectionPosForLayout(node: TLGraphNode, isInput: boolean, s const offset = (node.constructor as any).layout_slot_offset ?? (LiteGraph.NODE_SLOT_HEIGHT * 0.5); const side = isInput ? layout[0] : layout[1]; const data = LAYOUT_LABEL_TO_DATA[side]!; - const cxn = node[isInput ? 'inputs' : 'outputs'][slotNumber]; + const slotList = node[isInput ? 'inputs' : 'outputs']; + const cxn = slotList[slotNumber]; if (!cxn) { console.log('No connection found.. weird', isInput, slotNumber); return out; } + // Experimental; doesn't work without node.clip_area set (so it won't draw outside), + // but litegraph.core inexplicably clips the title off which we want... so, no go. + // if (cxn.hidden) { + // out[0] = node.pos[0] - 100000 + // out[1] = node.pos[1] - 100000 + // return out + // } + if (cxn.disabled) { + // Let's store the original colors if have them and haven't yet overridden + if (cxn.color_on !== '#666665') { + (cxn as any)._color_on_org = (cxn as any)._color_on_org || cxn.color_on; + (cxn as any)._color_off_org = (cxn as any)._color_off_org || cxn.color_off; + } + cxn.color_on = '#666665'; + cxn.color_off = '#666665'; + } else if (cxn.color_on === '#666665') { + cxn.color_on = (cxn as any)._color_on_org || undefined; + cxn.color_off = (cxn as any)._color_off_org || undefined; + } + // @ts-ignore + const displaySlot = collapseConnections ? 0 : (slotNumber - slotList.reduce((count, ioput, index) => { + count += index < slotNumber && ioput.hidden ? 1 : 0; + return count + }, 0)); cxn.dir = data[0]; if (side === 'Left') { if (node.flags.collapsed) { @@ -171,7 +196,6 @@ export function getConnectionPosForLayout(node: TLGraphNode, isInput: boolean, s if ((node.constructor as any)?.type.includes('Reroute')) { out[1] = node.pos[1] + (node.size[1] * .5); } else { - const displaySlot = collapseConnections ? 0 : slotNumber; out[1] = node.pos[1] + (displaySlot + 0.7) * LiteGraph.NODE_SLOT_HEIGHT + @@ -200,7 +224,6 @@ export function getConnectionPosForLayout(node: TLGraphNode, isInput: boolean, s if ((node.constructor as any)?.type.includes('Reroute')) { out[1] = node.pos[1] + (node.size[1] * .5); } else { - const displaySlot = collapseConnections ? 0 : slotNumber; out[1] = node.pos[1] + (displaySlot + 0.7) * LiteGraph.NODE_SLOT_HEIGHT +