Add lora support to Power Prompt. Also, add a simple Power Prompt w/o lora (for negative, really)

This commit is contained in:
rgthree
2023-08-21 00:33:10 -04:00
parent 48fd07874a
commit 751de06084
9 changed files with 676 additions and 160 deletions
+2
View File
@@ -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__))
+194 -53
View File
@@ -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(`<lora:${selected}:1.0>`);
}
}
}
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(`<lora:${selected}:1.0>`);
}
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('<lora:') || !selectedText.endsWith('>')) {
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);
}
}
});
+18 -3
View File
@@ -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 +
+41 -27
View File
@@ -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}')
+76 -10
View File
@@ -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='<lora:([^:>]*?)(?::(-?\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 '<lora:' in prompt:
_prompt_stripped, loras = get_and_strip_loras(prompt, True)
if len(loras):
log_node_warn(NODE_NAME, f'Found {len(loras)} lora tags in prompt but model & clip were not supplied!')
log_node_info(NODE_NAME, f'Loras not processed, keeping for TEXT output.')
conditioning=None
if opt_clip != None:
conditioning = CLIPTextEncode().encode(opt_clip, prompt)[0]
return (conditioning, opt_model, opt_clip, prompt)
+43
View File
@@ -0,0 +1,43 @@
import os
import folder_paths
from nodes import CLIPTextEncode
from .constants import get_category, get_name
from .power_prompt import RgthreePowerPrompt
class RgthreePowerPromptSimple(RgthreePowerPrompt):
NAME=get_name('Power Prompt - Simple')
CATEGORY = get_category()
@classmethod
def INPUT_TYPES(s):
SAVED_PROMPTS_FILES=folder_paths.get_filename_list('saved_prompts')
SAVED_PROMPTS_CONTENT=[]
for filename in SAVED_PROMPTS_FILES:
with open(folder_paths.get_full_path('saved_prompts', filename), 'r') as f:
SAVED_PROMPTS_CONTENT.append(f.read())
return {
'required': {
'prompt': ('STRING', {'multiline': True}),
},
'optional': {
"opt_clip": ("CLIP", ),
'insert_embedding': (['CHOOSE',] + [os.path.splitext(x)[0] for x in folder_paths.get_filename_list('embeddings')],),
'insert_saved': (['CHOOSE',] + SAVED_PROMPTS_FILES,),
},
'hidden': {
'values_insert_saved': (['CHOOSE'] + SAVED_PROMPTS_CONTENT,),
}
}
RETURN_TYPES = ('CONDITIONING', 'STRING',)
RETURN_NAMES = ('CONDITIONING', 'TEXT',)
FUNCTION = 'main'
def main(self, prompt, opt_clip=None, insert_embedding=None, insert_saved=None, values_insert_saved=None):
conditioning=None
if opt_clip != None:
conditioning = CLIPTextEncode().encode(opt_clip, prompt)[0]
return (conditioning, prompt)
+258 -63
View File
@@ -5,8 +5,9 @@ import {app} from '../../scripts/app.js';
import {api} from '../../scripts/api.js';
// @ts-ignore
import { ComfyWidgets } from '../../scripts/widgets.js';
import type {IWidget, IComboWidget, LGraphNode as TLGraphNode, LiteGraph as TLiteGraph} from './typings/litegraph.js';
import type {LLink, IComboWidget, LGraphNode as TLGraphNode, LiteGraph as TLiteGraph, INodeOutputSlot, INodeInputSlot, Vector2} from './typings/litegraph.js';
import type {ComfyApp, ComfyObjectInfo, ComfyGraphNode} from './typings/comfy.js'
import {addConnectionLayoutSupport} from './utils.js';
declare const LiteGraph: typeof TLiteGraph;
declare const LGraphNode: typeof TLGraphNode;
@@ -14,34 +15,88 @@ declare const LGraphNode: typeof TLGraphNode;
/** Wraps a node instance keeping closure without mucking the finicky types. */
class PowerPrompt {
readonly isSimple: boolean;
readonly node: ComfyGraphNode;
readonly promptEl: HTMLTextAreaElement;
nodeData: ComfyObjectInfo;
embeddingWidget?: IComboWidget;
savedWidget?: IComboWidget;
savedValues?: string[];
readonly boundOnFreshNodeDefs: (event: CustomEvent) => 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(`<lora:${selected}:1.0>`);
}
}
}
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(`<lora:${selected}:1.0>`);
}
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 "<lora:". So work backwards until we
// find something that we know can't be a lora, or a "<".
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);
}
// Bail if this isn't a lora.
if (!selectedText.startsWith('<lora:') || !selectedText.endsWith('>')) {
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)
}
}
});
+18 -1
View File
@@ -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<T extends IWidget = IWidget> = (
event?: MouseEvent
) => void;
// #rgthree
export type WidgetComboCallback<T extends IWidget = IWidget> = (
this: T,
value: T["value"][0],
graphCanvas: LGraphCanvas,
node: LGraphNode,
pos: Vector2,
event?: MouseEvent
) => void;
export interface IWidget<TValue = any, TOptions = any> {
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<string, {}> {
@@ -817,7 +833,8 @@ export declare class LGraphNode {
type: T["type"],
name: string,
value: T["value"],
callback?: WidgetCallback<T> | string,
// @rgthree
callback?: T["callback"] | string,//WidgetCallback<T> | string,
options?: T["options"]
): T;
+26 -3
View File
@@ -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<Number>((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 +