diff --git a/.gitignore b/.gitignore index 68bc17f..35ea65a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,7 @@ +# pre-commit +.flake8 +.pre-commit-config.yaml + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/README.md b/README.md index c98fb75..ec74f45 100644 --- a/README.md +++ b/README.md @@ -5,6 +5,10 @@ Or perhaps you use only one, cause you forgot you have tens of them installed? Try this. +### Note to existing users: + +Version v2 will cause validation error in workflows that used previous version. There's nothing to worry about, simply replace error causing nodes with new ones, or stick to version v1 + # Installation Exactly the same as with other simple custom nodes. @@ -21,13 +25,19 @@ To uninstall: # Use -Simply connect `Embedding Picker` node, that you will find in `utils` menu, in between _CLIP Encoding_ and _Text Multiine_ or other similar nodes. +Right click on the _CLIP Text Encode_ node and select the top option 'Prepend Embedding Picker'. This will create the node itself and copy all your prompts. +A similar option exists on the `Embedding Picker' node itself, use this to quickly chain multiple embeddings. -You can also right click on _CLIP Encoding_ node, and choose `Prepend Embedding Picker` to place and link one automatically. -Similar option exists on `Embedding Picker` node itself, feel free to use both multiple times to chain more nodes. +You can also find the node in `utils` menu. -By default the embedding will be set as the first prompt, set `append` to True if you prefer it to be the last one. +By default the embedding will be the placed as first prompt, and will result in such text -## Example Workflow: +> embedding:EasyNegative.pt, text, watermark + +Set `append` to True if you prefer it to be at the end, i. e. + +> text, watermark, embedding:EasyNegative.pt + +## Example Workflow ![Example](example_workflow.png) diff --git a/__init__.py b/__init__.py index dad1c00..8ec8393 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,7 @@ -from .node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .node import EmbeddingPicker -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] +NODE_CLASS_MAPPINGS = {"EmbeddingPicker": EmbeddingPicker} +NODE_DISPLAY_NAME_MAPPINGS = {"EmbeddingPicker": "Embedding Picker"} # --------------------------- Install web extension ---------------------------- import shutil diff --git a/example_workflow.png b/example_workflow.png index 5f79dc3..ba9f6e5 100644 Binary files a/example_workflow.png and b/example_workflow.png differ diff --git a/js/quickNodes.js b/js/quickNodes.js index 57b8240..aaf4b23 100644 --- a/js/quickNodes.js +++ b/js/quickNodes.js @@ -18,7 +18,9 @@ function hideWidget(node, widget, suffix = "") { if (!node_input || !node_input.link) { return undefined; } - return widget.origSerializeValue ? widget.origSerializeValue() : widget.value; + return widget.origSerializeValue + ? widget.origSerializeValue() + : widget.value; }; // Hide any linked widgets, e.g. seed+seedControl @@ -79,23 +81,30 @@ function addMenuHandler(nodeType, callback) { * @param {object} [options={}] - Optional parameters to adjust the new node's behavior. * @param {boolean} [options.select=true] - If set to true, the new node will be selected. * @param {number} [options.shiftY=0] - The vertical shift from the reference node's position. - * @param {boolean} [options.before=false] - If true, the new node will be positioned + * @param {boolean} [options.before=false] - If true, the new node will be positioned * to the left of the reference node; otherwise, to the right. + * @param {array} [options.size] - The size of the new node. * * @returns {object} The newly created node. */ -function addNode(name, nextTo, options = {}) { +function placeNewNode(name, nextTo, options = {}) { const nodeSeparation = 30; - const { select = true, shiftY = 0, before = false } = options; + const { select = true, shiftY = 0, before = false, size = null } = options; const node = LiteGraph.createNode(name); - node.size = [300, 100]; + + if (size) { + node.size = size; + } + app.graph.add(node); const [nextToX, nextToY] = nextTo.pos; const [nextToWidth] = nextTo.size; - const offsetX = before ? -node.size[0] - nodeSeparation : nextToWidth + nodeSeparation; + const offsetX = before + ? -node.size[0] - nodeSeparation + : nextToWidth + nodeSeparation; node.pos = [nextToX + offsetX, nextToY + shiftY]; if (select) { @@ -105,57 +114,75 @@ function addNode(name, nextTo, options = {}) { return node; } +/** + * Converts every `node` widget that matches `newNodeWidgetNames` + * into an input slots, before linking them with `newNode`. + * + * @param {Object} node - The "right" node. + * @param {Object} nodeData - The "right" node data. + * @param {Object} newNode - The "left" node. + * @param {Array} newNodeWidgetNames - Widget/Slots names that should be connected between nodes. + */ +function prependNode(node, nodeData, newNode, newNodeWidgetNames) { + for (const widget_name of newNodeWidgetNames) { + let slot = node.findInputSlot(widget_name); + if (slot === -1) { + //Convert widget into input + const w = node.widgets.find((obj) => obj.name === widget_name); + const { required, optional } = nodeData?.input; + const config = required[w.name] || + optional?.[w.name] || [w.type, w.options || {}]; + convertToInput(node, w, config); + + slot = node.findInputSlot(widget_name); + } + newNode.connect(newNode.findOutputSlot(widget_name), node, slot); + } +} + app.registerExtension({ name: "trop.EP.QuickNodes", async beforeRegisterNodeDef(nodeType, nodeData, app) { - if (nodeData.name === "EmbeddingPicker" || nodeData.name === "Text Multiline") { + if ( + nodeData.name === "EmbeddingPicker" || + nodeData.name === "CLIPTextEncode" + ) { addMenuHandler(nodeType, function (_, options) { options.unshift({ - content: "Append Embedding Picker", + content: `Prepend Embedding Picker`, callback: () => { - const EPNode = addNode("EmbeddingPicker", this); - const EPNodeLinks = this.outputs[0].links ? this.outputs[0].links.map((l) => ({ ...graph.links[l] })) : []; - - this.disconnectOutput(0); - this.connect(0, EPNode, 0); - - // reconnect all links to new node - for (const link of EPNodeLinks) { - EPNode.connect(0, link.target_id, link.target_slot); - } - }, - }); - }); - } - - if (nodeData.name === "CLIPTextEncode") { - addMenuHandler(nodeType, function (_, options) { - options.unshift({ - content: "Prepend Embedding Picker", - callback: () => { - //TODO: check input for widget - if (this.findInputSlot("text") === -1) { - const w = this.widgets[0]; - const { required, optional } = nodeData?.input || {}; - const config = required[w.name] || optional?.[w.name] || [w.type, w.options || {}]; - convertToInput(this, w, config); - } - - const EPNode = addNode("EmbeddingPicker", this, { + const newNode = placeNewNode("EmbeddingPicker", this, { before: true, - shiftY: 20, + shiftY: nodeData.name === "CLIPTextEncode" ? 20 : 0, + size: [300, 200], }); - // connect new node in between - if (this.getInputLink(1)) { - const previousNode = this.getInputNode(1); - this.disconnectInput(1); - previousNode.connect(0, EPNode, 0); + try { + // Copy colors to new node + ["bgcolor", "color"].forEach((prop) => { + if (typeof this[prop] !== "undefined") { + newNode[prop] = this[prop]; + } + }); + + // copy prompts to new node + const prompts = this.widgets.find((w) => w.name === "text").value; + if (prompts && typeof prompts !== "undefined") { + newNode.widgets[3].value = prompts; + } + } catch (e) { + console.error("Failed to copy prompts", e); + } + + prependNode(this, nodeData, newNode, ["text"]); + + if (this.size[1] > 120) { + // In some cases settings size too low breaks node. Minimum size of CLIPTextEncode is 210x50 and 210x118 for EP + this.size = [this.size[0], 120]; } - EPNode.connect(0, this, 1); }, }); }); } }, -}); \ No newline at end of file +}); diff --git a/node.py b/node.py index 0093437..708817f 100644 --- a/node.py +++ b/node.py @@ -1,54 +1,50 @@ -from pathlib import Path - -import folder_paths - - -class EmbeddingPicker: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(self): - embeddings = folder_paths.get_filename_list("embeddings") - - return { - "required": { - "prompts": ("STRING", {"forceInput": True}), - "embedding": ((embeddings),), - "emphasis": ( - "FLOAT", - { - "default": 1.0, - "min": 0.05, - "max": 3.0, - "step": 0.05, - }, - ), - "append": ( - "BOOLEAN", - {"default": False, "label_on": "true ", "label_off": "false "}, - ), - }, - } - - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("prompts",) - FUNCTION = "concat_embedding" - OUTPUT_NODE = False - - CATEGORY = "utils" - - def concat_embedding(self, prompts, embedding, emphasis, append): - emb = "embedding:" + Path(embedding).stem - - emphasis = f"{emphasis:.3f}" - if emphasis != "1.000": - emb = f"({emb}:{emphasis})" - - output = f"{prompts}, {emb}" if append else f"{emb}, {prompts}" - - return (output,) - - -NODE_CLASS_MAPPINGS = {"EmbeddingPicker": EmbeddingPicker} -NODE_DISPLAY_NAME_MAPPINGS = {"EmbeddingPicker": "Embedding Picker"} +from pathlib import Path + +import folder_paths + + +class EmbeddingPicker: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(self): + embeddings = folder_paths.get_filename_list("embeddings") + + return { + "required": { + "embedding": ((embeddings),), + "emphasis": ( + "FLOAT", + { + "default": 1.0, + "min": 0.05, + "max": 3.0, + "step": 0.05, + }, + ), + "append": ( + "BOOLEAN", + {"default": False, "label_on": "true ", "label_off": "false "}, + ), + "text": ("STRING", {"multiline": True}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("text",) + FUNCTION = "concat_embedding" + OUTPUT_NODE = False + + CATEGORY = "utils" + + def concat_embedding(self, text, embedding, emphasis, append): + emb = "embedding:" + Path(embedding).stem + + emphasis = f"{emphasis:.3f}" + if emphasis != "1.000": + emb = f"({emb}:{emphasis})" + + output = f"{text}, {emb}" if append else f"{emb}, {text}" + + return (output,)