diff --git a/README.md b/README.md index a8a57de..d303168 100644 --- a/README.md +++ b/README.md @@ -169,7 +169,6 @@ Multiple `@combine` blocks are allowed and processed independently. A `---` outs #### Validation The node stops the workflow on error if: -- `options` is empty - A `@combine` block is opened inside another - A `@end` is found without a matching `@combine` - A `@combine` block is never closed @@ -229,6 +228,36 @@ Slot lengths do not need to match. --- +### PromptSwitch + +![PromptSwitch_v1](docs/images/PromptSwitch_v1.png) + +#### Inputs + +- `on_true` (str) — value returned when condition is `True`. Accepts direct input or connected node. +- `on_false` (str) — value returned when condition is `False`. Accepts direct input or connected node. +- `condition` (bool), default `True` — controls which input is routed to the output. + +#### Outputs + +- `text` (str) — the selected value based on condition. + +#### Usage + +Connect or type a value in `on_true` and `on_false`. Toggle `condition` to switch between them. + +If the node is bypassed, an empty string is returned. + +#### Editor + +Input ports and widget borders are colored based on the active condition: +- **Orange**: active input when condition is `True` +- **Blue**: active input when condition is `False` +- **Gray**: inactive input +Colors are only shown when an input is connected or has a value. + +--- + ### Web #### contentEditable editor @@ -268,6 +297,8 @@ The selection is preserved after each keypress, consistent with ComfyUI native b ## TODO +- Centralize colors used + `PromptPresetSelector` - Add a field to easily concatenate text at the end of the output -- Wildcard mode for presets (randomly selected, index to -1 or -2?) \ No newline at end of file +- Wildcard mode for presets (randomly selected, index to -1 or -2?) diff --git a/__init__.py b/__init__.py index 0caff71..f194d9c 100644 --- a/__init__.py +++ b/__init__.py @@ -18,16 +18,24 @@ from .src.nodes.prompt_preset_selector import ( from .src.nodes.prompt_preset_selector import ( NODE_DISPLAY_NAME_MAPPINGS as PromptPresetSelectorDisplayMappings, ) +from .src.nodes.prompt_switch import ( + NODE_CLASS_MAPPINGS as PromptSwitchMappings, +) +from .src.nodes.prompt_switch import ( + NODE_DISPLAY_NAME_MAPPINGS as PromptSwitchDisplayMappings, +) NODE_CLASS_MAPPINGS = {} NODE_CLASS_MAPPINGS.update(PromptPresetSelectorMappings) NODE_CLASS_MAPPINGS.update(PromptOptionPickerMappings) NODE_CLASS_MAPPINGS.update(PromptLayoutFillerMappings) +NODE_CLASS_MAPPINGS.update(PromptSwitchMappings) NODE_DISPLAY_NAME_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS.update(PromptPresetSelectorDisplayMappings) NODE_DISPLAY_NAME_MAPPINGS.update(PromptOptionPickerDisplayMappings) NODE_DISPLAY_NAME_MAPPINGS.update(PromptLayoutFillerDisplayMappings) +NODE_DISPLAY_NAME_MAPPINGS.update(PromptSwitchDisplayMappings) # Node registration __all__ = [ diff --git a/docs/images/PromptPresetSelector_v1.png b/docs/images/PromptPresetSelector_v1.png deleted file mode 100644 index 679c885..0000000 Binary files a/docs/images/PromptPresetSelector_v1.png and /dev/null differ diff --git a/docs/images/PromptPresetSelector_v2.png b/docs/images/PromptPresetSelector_v2.png deleted file mode 100644 index 74ec24c..0000000 Binary files a/docs/images/PromptPresetSelector_v2.png and /dev/null differ diff --git a/docs/images/PromptSwitch_v1.png b/docs/images/PromptSwitch_v1.png new file mode 100644 index 0000000..fb02c72 Binary files /dev/null and b/docs/images/PromptSwitch_v1.png differ diff --git a/pyproject.toml b/pyproject.toml index 9ab1058..df25989 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "ComfyUI-RE-CustomUtils" -version = "0.3.0" +version = "0.4.0" description = "A collection of custom nodes for ComfyUI." license = { file = "LICENSE" } readme = "README.md" diff --git a/src/nodes/prompt_option_picker.py b/src/nodes/prompt_option_picker.py index 0a69b43..d2f5be5 100644 --- a/src/nodes/prompt_option_picker.py +++ b/src/nodes/prompt_option_picker.py @@ -42,14 +42,13 @@ class PromptOptionPicker: @classmethod def VALIDATE_INPUTS(cls, options, **kwargs): - if not options or not options.strip(): - return "options cannot be empty" + try: + parse_options(options) + except ValueError as e: + return str(e) return True def process(self, options, selected): - # Validate only, JS handles the dropdown - parse_options(options) - # -- displays an empty string return ("" if selected == "--" else selected,) diff --git a/src/nodes/prompt_switch.py b/src/nodes/prompt_switch.py new file mode 100644 index 0000000..090f95c --- /dev/null +++ b/src/nodes/prompt_switch.py @@ -0,0 +1,53 @@ +from .. import config + + +class PromptSwitch: + """ + Routes one of two STRING inputs to the output based on a boolean condition. + If condition is True, returns on_true. If False, returns on_false. + If the node is bypassed, ComfyUI returns on_true by convention. + """ + + CATEGORY = config.NODE_CATEGORY + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "on_true": ( + "STRING", + { + "tooltip": "Value returned when condition is True.", + }, + ), + "on_false": ( + "STRING", + { + "tooltip": "Value returned when condition is False.", + }, + ), + "condition": ( + "BOOLEAN", + { + "default": True, + "tooltip": "If True, returns on_true. If False, returns on_false.", + }, + ), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("text",) + FUNCTION = "process" + + def process(self, on_true: str, on_false: str, condition: bool) -> tuple[str]: + return (on_true if condition else on_false,) + + +NODE_CLASS_MAPPINGS = { + "PromptSwitch": PromptSwitch, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PromptSwitch": "Prompt Switch", +} diff --git a/tests/test_prompt_option_picker.py b/tests/test_prompt_option_picker.py index d10e6a4..7b8b8bc 100644 --- a/tests/test_prompt_option_picker.py +++ b/tests/test_prompt_option_picker.py @@ -41,17 +41,22 @@ def test_whitespace_only_lines(node): assert value == "" +def test_empty_options_returns_empty(node): + (value,) = node.process(options="", selected="--") + assert value == "" + + # --- VALIDATE_INPUTS --- def test_validate_empty_options(): result = PromptOptionPicker.VALIDATE_INPUTS(options="") - assert result != True + assert result == True def test_validate_whitespace_only(): result = PromptOptionPicker.VALIDATE_INPUTS(options=" ") - assert result != True + assert result == True def test_validate_valid(): diff --git a/tests/test_prompt_switch.py b/tests/test_prompt_switch.py new file mode 100644 index 0000000..3425804 --- /dev/null +++ b/tests/test_prompt_switch.py @@ -0,0 +1,67 @@ +import pytest +from src.nodes.prompt_switch import PromptSwitch + + +@pytest.fixture +def node(): + return PromptSwitch() + + +def run(node, on_true="", on_false="", condition=True): + return node.process(on_true=on_true, on_false=on_false, condition=condition) + + +# --- Base cases --- + + +def test_condition_true(node): + (text,) = run(node, on_true="hello", on_false="world", condition=True) + assert text == "hello" + + +def test_condition_false(node): + (text,) = run(node, on_true="hello", on_false="world", condition=False) + assert text == "world" + + +def test_empty_on_true(node): + (text,) = run(node, on_true="", on_false="world", condition=True) + assert text == "" + + +def test_empty_on_false(node): + (text,) = run(node, on_true="hello", on_false="", condition=False) + assert text == "" + + +def test_both_empty(node): + (text,) = run(node, on_true="", on_false="", condition=True) + assert text == "" + + +def test_multiline_value(node): + (text,) = run(node, on_true="line1\nline2", on_false="other", condition=True) + assert text == "line1\nline2" + + +# --- With connected nodes (values passed as strings) --- + + +def test_connected_node_true(node): + (text,) = run( + node, + on_true="masterpiece, detailed", + on_false="simple", + condition=True, + ) + assert text == "masterpiece, detailed" + + +def test_connected_node_false(node): + (text,) = run( + node, + on_true="detailed", + on_false="{% preset_a | preset_b %}", + condition=False, + ) + assert text == "{% preset_a | preset_b %}" diff --git a/web/prompt_layout_filler.js b/web/prompt_layout_filler.js index 0a19a4d..f6fc906 100644 --- a/web/prompt_layout_filler.js +++ b/web/prompt_layout_filler.js @@ -118,7 +118,7 @@ const attachEditor = (node) => { renderColored(); }; - // --- Initial render --- + // Initial render updateSlotVisibility(node); renderColored(); }; diff --git a/web/prompt_switch.js b/web/prompt_switch.js new file mode 100644 index 0000000..0d6943b --- /dev/null +++ b/web/prompt_switch.js @@ -0,0 +1,127 @@ +import { hookWidget } from "./utils.js"; +import { app } from "/scripts/app.js"; + +// --- Constants --- + +const NODE_NAME = "PromptSwitch"; + +const COLORS = { + true: "#ff9800", + false: "#64b5f6", + inactive: "#888888", +}; + +// --- Highlight --- + +const updateInputColors = (node, condition) => { + const trueInput = node.inputs?.find((inp) => inp.name === "on_true"); + const falseInput = node.inputs?.find((inp) => inp.name === "on_false"); + const trueWidget = node.widgets?.find((w) => w.name === "on_true"); + const falseWidget = node.widgets?.find((w) => w.name === "on_false"); + if (!trueInput || !falseInput) return; + + const trueActive = + trueInput.link != null || (trueWidget?.value ?? "").trim() !== ""; + const falseActive = + falseInput.link != null || (falseWidget?.value ?? "").trim() !== ""; + + trueInput.color_on = trueActive + ? condition + ? COLORS.true + : COLORS.inactive + : undefined; + trueInput.color_off = trueInput.color_on; + + falseInput.color_on = falseActive + ? condition + ? COLORS.inactive + : COLORS.false + : undefined; + falseInput.color_off = falseInput.color_on; + + if (node.graph) node.graph.setDirtyCanvas(true, true); +}; + +// --- Attach --- + +const attachSwitch = (node) => { + if (node.type !== NODE_NAME) return; + if (node._switchAttached) return; + node._switchAttached = true; + + const trueWidget = node.widgets?.find((w) => w.name === "on_true"); + const falseWidget = node.widgets?.find((w) => w.name === "on_false"); + const conditionWidget = node.widgets?.find((w) => w.name === "condition"); + if (!conditionWidget) return; + + const refresh = () => { + const condition = conditionWidget.value ?? true; + updateInputColors(node, condition); + if (node.graph) node.graph.setDirtyCanvas(true, true); + }; + + // Input port colors + hookWidget(conditionWidget, refresh); + hookWidget(trueWidget, refresh); + hookWidget(falseWidget, refresh); + + // Widget border via canvas + const original = node.onDrawForeground; + node.onDrawForeground = function (ctx) { + if (original) original.call(this, ctx); + + const trueWidget = node.widgets?.find((w) => w.name === "on_true"); + const falseWidget = node.widgets?.find((w) => w.name === "on_false"); + const trueInput = node.inputs?.find((inp) => inp.name === "on_true"); + const falseInput = node.inputs?.find((inp) => inp.name === "on_false"); + if (!trueWidget || !falseWidget) return; + + const condition = + node.widgets?.find((w) => w.name === "condition")?.value ?? true; + + const drawBorder = (widget, color, input) => { + // Skip if not connected AND widget value is empty + const isConnected = input?.link != null; + const hasValue = widget.value != null && widget.value.trim() !== ""; + if (!isConnected && !hasValue) return; + + const x = 0; + const y = widget.last_y - 2; + const w = node.size[0]; + const h = widget.computedHeight ?? 20; + + ctx.strokeStyle = color; + ctx.lineWidth = 2; + ctx.strokeRect(x, y, w, h); + }; + + // Draw inactive first, active on top + if (condition) { + drawBorder(falseWidget, COLORS.inactive, falseInput); + drawBorder(trueWidget, COLORS.true, trueInput); + } else { + drawBorder(trueWidget, COLORS.inactive, trueInput); + drawBorder(falseWidget, COLORS.false, falseInput); + } + }; + + // Initial render + refresh(); +}; + +// --- Registration --- + +app.registerExtension({ + name: NODE_NAME, + beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== NODE_NAME) return; + + const original = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + if (original) original.call(this); + + const node = this; + requestAnimationFrame(() => attachSwitch(node)); + }; + }, +});