diff --git a/__init__.py b/__init__.py index a422160..2ec4cd5 100644 --- a/__init__.py +++ b/__init__.py @@ -1,171 +1,13 @@ """ -SimpleVariables - Set and Get nodes for storing/retrieving any data by name in ComfyUI +SimpleVariables - Pure frontend Set/Get Variable nodes for ComfyUI + +These nodes are implemented entirely in JavaScript (web/simplevariables.js). +They work by redirecting graph connections - no Python execution needed. """ -import time - -class AnyType(str): - """A special type that bypasses ComfyUI's type checking - always matches any type.""" - def __eq__(self, _): - return True - - def __ne__(self, _): - return False - - -# Create the any type instance -ANY = AnyType("*") - -# Global storage for variables -_variable_storage = {} - - -def get_preview_text(variable_name, value): - """Generate preview text for a value""" - if value is None: - return f"{variable_name} = None" - - type_name = type(value).__name__ - - if hasattr(value, 'shape'): - return f"{variable_name} = {type_name} {list(value.shape)}" - elif isinstance(value, str): - val_str = value[:50] + "..." if len(value) > 50 else value - return f'{variable_name} = "{val_str}"' - elif isinstance(value, (int, float, bool)): - return f"{variable_name} = {value}" - elif isinstance(value, list): - return f"{variable_name} = list[{len(value)}]" - elif isinstance(value, dict): - return f"{variable_name} = dict[{len(value)} keys]" - else: - return f"{variable_name} = {type_name}" - - -class SetVariable: - """Store any data with a custom variable name - shows preview on node""" - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "variable_name": ("STRING", {"default": "my_var"}), - }, - "optional": { - "value": (ANY,), - } - } - RETURN_TYPES = (ANY,) - RETURN_NAMES = ("value",) - FUNCTION = "set_var" - CATEGORY = "variables" - OUTPUT_NODE = True - - @classmethod - def IS_CHANGED(cls, **kwargs): - return time.time() - - def set_var(self, variable_name, value=None): - _variable_storage[variable_name] = value - preview = get_preview_text(variable_name, value) - return {"ui": {"text": [preview]}, "result": (value,)} - - -class GetVariable: - """Retrieve data by variable name - shows preview on node - - Connect SetVariable's output to 'trigger' to ensure proper execution order. - """ - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "variable_name": ("STRING", {"default": "my_var"}), - }, - "optional": { - "trigger": (ANY,), - } - } - RETURN_TYPES = (ANY,) - RETURN_NAMES = ("value",) - FUNCTION = "get_var" - CATEGORY = "variables" - OUTPUT_NODE = True - - @classmethod - def IS_CHANGED(cls, **kwargs): - return time.time() - - def get_var(self, variable_name, trigger=None): - if variable_name not in _variable_storage: - available = list(_variable_storage.keys()) - preview = f"{variable_name} = NOT SET (available: {available})" - return {"ui": {"text": [preview]}, "result": (None,)} - - value = _variable_storage[variable_name] - preview = get_preview_text(variable_name, value) - return {"ui": {"text": [preview]}, "result": (value,)} - - -class ListVariables: - """Show all stored variables on the node""" - @classmethod - def INPUT_TYPES(cls): - return {"required": {}} - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("variable_list",) - FUNCTION = "list_vars" - CATEGORY = "variables" - OUTPUT_NODE = True - - def list_vars(self): - if not _variable_storage: - text = "(no variables set)" - else: - lines = [get_preview_text(name, value) for name, value in _variable_storage.items()] - text = "\n".join(lines) - - return {"ui": {"text": [text]}, "result": (text,)} - - -class ClearVariables: - """Clear all stored variables""" - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "clear": ("BOOLEAN", {"default": False}), - } - } - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("status",) - FUNCTION = "clear_vars" - CATEGORY = "variables" - OUTPUT_NODE = True - - def clear_vars(self, clear): - if clear: - count = len(_variable_storage) - _variable_storage.clear() - status = f"Cleared {count} variable(s)" - else: - status = "Not cleared (set clear=True)" - - return {"ui": {"text": [status]}, "result": (status,)} - - -NODE_CLASS_MAPPINGS = { - "SetVariable": SetVariable, - "GetVariable": GetVariable, - "ListVariables": ListVariables, - "ClearVariables": ClearVariables, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "SetVariable": "Set Variable", - "GetVariable": "Get Variable", - "ListVariables": "List Variables", - "ClearVariables": "Clear Variables", -} +# No Python nodes - everything is in JavaScript +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} WEB_DIRECTORY = "./web" __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/web/simplevariables.js b/web/simplevariables.js index 425cc2d..1e633b9 100644 --- a/web/simplevariables.js +++ b/web/simplevariables.js @@ -1,48 +1,176 @@ import { app } from "/scripts/app.js"; -import { ComfyWidgets } from "/scripts/widgets.js"; -function displayText(nodeType, nodeData, appInstance) { - const onExecutedOriginal = nodeType.prototype.onExecuted; - nodeType.prototype.onExecuted = function (message) { - onExecutedOriginal?.apply(this, arguments); +// Pure frontend Set/Get Variable nodes - no Python execution needed +// Based on KJNodes SetGet implementation - if (message?.text) { - // Find or create text widget - let textWidget = this.widgets?.find(w => w.name === "display_text"); - - if (!textWidget) { - textWidget = ComfyWidgets.STRING(this, "display_text", [ - "STRING", { multiline: true } - ], appInstance).widget; - textWidget.inputEl.readOnly = true; - textWidget.inputEl.style.opacity = 0.7; - textWidget.inputEl.style.fontFamily = "monospace"; - } - - // Get text from message - let text = message.text; - if (Array.isArray(text)) { - text = text[0] || ""; - } - textWidget.value = text; - - // Resize node to fit - requestAnimationFrame(() => { - const sz = this.computeSize(); - if (sz[0] < this.size[0]) sz[0] = this.size[0]; - if (sz[1] < this.size[1]) sz[1] = this.size[1]; - this.onResize?.(sz); - appInstance.graph.setDirtyCanvas(true, false); - }); - } - }; -} +const LGraphNode = LiteGraph.LGraphNode; app.registerExtension({ - name: "SimpleVariables", - async beforeRegisterNodeDef(nodeType, nodeData, app) { - if (["SetVariable", "GetVariable", "ListVariables", "ClearVariables"].includes(nodeData.name)) { - displayText(nodeType, nodeData, app); + name: "SimpleVariables.SetVariable", + registerCustomNodes() { + class SetVariableNode extends LGraphNode { + constructor(title) { + super(title); + + this.properties = this.properties || {}; + this.properties.variableName = ""; + + // Variable name widget + this.addWidget("text", "name", "", (value) => { + this.properties.variableName = value; + this.title = value ? `Set: ${value}` : "Set Variable"; + this.updateGetters(); + }); + + // Input and passthrough output + this.addInput("value", "*"); + this.addOutput("value", "*"); + + this.onConnectionsChange = (slotType, slot, isConnect, linkInfo) => { + if (slotType === 1 && this.inputs[0]) { // Input changed + if (isConnect && linkInfo) { + const type = this.getInputType(0); + if (type) { + this.inputs[0].type = type; + this.inputs[0].name = type; + this.outputs[0].type = type; + this.outputs[0].name = type; + } + } else if (!isConnect) { + this.inputs[0].type = "*"; + this.inputs[0].name = "value"; + this.outputs[0].type = "*"; + this.outputs[0].name = "value"; + } + this.updateGetters(); + } + }; + + this.isVirtualNode = true; + } + + getInputType(slot) { + const link = this.getInputLink(slot); + if (link) { + const originNode = this.graph._nodes_by_id[link.origin_id]; + if (originNode && originNode.outputs[link.origin_slot]) { + return originNode.outputs[link.origin_slot].type; + } + } + return null; + } + + updateGetters() { + if (!this.graph) return; + const name = this.properties.variableName; + this.graph._nodes.forEach(node => { + if (node.type === "GetVariableNode" && node.properties.variableName === name) { + node.syncWithSetter(); + } + }); + } + + onRemoved() { + this.updateGetters(); + } } - }, + + LiteGraph.registerNodeType("SetVariableNode", Object.assign(SetVariableNode, { + title: "Set Variable", + })); + SetVariableNode.category = "variables"; + } +}); + +app.registerExtension({ + name: "SimpleVariables.GetVariable", + registerCustomNodes() { + class GetVariableNode extends LGraphNode { + constructor(title) { + super(title); + + this.properties = this.properties || {}; + this.properties.variableName = ""; + + // Variable name dropdown (shows available setters) + this.addWidget("combo", "name", "", (value) => { + this.properties.variableName = value; + this.title = value ? `Get: ${value}` : "Get Variable"; + this.syncWithSetter(); + }, { + values: () => this.getAvailableVariables() + }); + + this.addOutput("value", "*"); + + this.onConnectionsChange = () => { + this.validateLinks(); + }; + + this.isVirtualNode = true; + } + + getAvailableVariables() { + if (!this.graph) return [""]; + const setters = this.graph._nodes.filter(n => n.type === "SetVariableNode"); + const names = setters + .map(n => n.properties?.variableName || n.widgets?.[0]?.value || "") + .filter(n => n !== "") + .sort(); + return ["", ...names]; + } + + findSetter() { + if (!this.graph || !this.properties.variableName) return null; + const name = this.properties.variableName; + return this.graph._nodes.find(n => + n.type === "SetVariableNode" && + (n.properties?.variableName === name || n.widgets?.[0]?.value === name) + ); + } + + syncWithSetter() { + const setter = this.findSetter(); + if (setter && setter.inputs[0]) { + this.outputs[0].type = setter.inputs[0].type || "*"; + this.outputs[0].name = setter.inputs[0].name || "value"; + } else { + this.outputs[0].type = "*"; + this.outputs[0].name = "value"; + } + this.validateLinks(); + } + + validateLinks() { + if (!this.outputs[0].links || this.outputs[0].type === "*") return; + // Remove incompatible links + const linksToRemove = []; + this.outputs[0].links.forEach(linkId => { + const link = this.graph.links[linkId]; + if (link && link.type !== "*" && link.type !== this.outputs[0].type) { + linksToRemove.push(linkId); + } + }); + linksToRemove.forEach(id => this.graph.removeLink(id)); + } + + // This is the magic - redirect to the setter's input + getInputLink(slot) { + const setter = this.findSetter(); + if (setter && setter.inputs[0]) { + return this.graph.links[setter.inputs[0].link]; + } + return null; + } + + onAdded() { + this.syncWithSetter(); + } + } + + LiteGraph.registerNodeType("GetVariableNode", Object.assign(GetVariableNode, { + title: "Get Variable", + })); + GetVariableNode.category = "variables"; + } });