diff --git a/README.md b/README.md index 8026717..25e9e43 100644 --- a/README.md +++ b/README.md @@ -20,8 +20,7 @@ Multi-stage prompt pipelines reuse the same long system prompts. Pasting that pr |---|---|---|---| | `mode` | combo | `fixed` | How the role is chosen: `fixed`, `increment`, or `random`. See Selection modes below. | | `role` | combo | first role | Role used in `fixed` mode. The list is built from `roles/*.md` at startup. Add a `.md` file and restart ComfyUI to add a role. | -| `value` | INT | `0` | The single stepping input for `increment` and `random` modes (ignored in `fixed`). In `increment` it is the position; set its `control_after_generate` to **increment** to advance one role per generation, wrapping within `start..end`. In `random` it is the seed; set `control_after_generate` to **randomize** for a new pick each generation, and the same value reproduces the same role. | -| `start` | INT | `0` | Lower bound position (0-based) into the alphabetically sorted role list. | +| `start` | INT | `0` | Lower bound position (0-based) into the alphabetically sorted role list. Bounds `increment` and `random`. | | `end` | INT | `-1` | Upper bound position; `-1` means the last role. Out-of-range values are clamped. | ## Outputs @@ -37,19 +36,16 @@ Multi-stage prompt pipelines reuse the same long system prompts. Pasting that pr Roles are sorted alphabetically by display name; `start`/`end` index into that list (0-based, `end = -1` means the last). The range is inclusive. -- **fixed**: output the `role` dropdown. `value`/`start`/`end` ignored. -- **increment**: output the role at `start + (value mod span)`, wrapping back to - `start` after `end`. Set the `value` widget's `control_after_generate` to - **increment** so each generation advances to the next role. -- **random**: output a role picked from `start..end` using `value` as the seed. The - same value always yields the same role; set `control_after_generate` to - **randomize** for a fresh pick per run. +- **fixed**: output the `role` dropdown. `start`/`end` ignored. +- **increment**: output the next role alphabetically, advancing one step per + generation and wrapping from `end` back to `start`. The node steps on its own; + there is no extra widget to set. The counter is per-node and process-local, so + it restarts from `start` when ComfyUI restarts. +- **random**: output a role picked at random within `start..end`, fresh each + generation. -Note: `control_after_generate` is the same per-widget control ComfyUI puts on the -`seed` widget of samplers. `increment` and `random` share the single `value` -widget because only one is ever read at a time. Both still work without -`control_after_generate` by changing `value` yourself or using ComfyUI's batch -count. +`increment` and `random` re-run the node on every generation (so a downstream LLM +node re-runs too); `fixed` only re-runs when the role or its file text changes. ## Adding roles diff --git a/nodes.py b/nodes.py index d651c4b..1e24d97 100644 --- a/nodes.py +++ b/nodes.py @@ -2,8 +2,8 @@ A role picker: a dropdown of system-prompt roles loaded from roles/*.md, output as a STRING to wire into an LLM node's system_prompt input. A mode control -selects the role by fixed dropdown, by an incrementing value, or randomly seeded -by that same value, bounded to a start..end slice of the list. +selects the role: a fixed dropdown pick, an automatic alphabetical step (one role +per generation), or a fresh random pick, bounded to a start..end slice of the list. Dual-API: the shared core (role_core.select) is wrapped by a V1 class (NODE_CLASS_MAPPINGS, authoritative on the 0.25.0 if/elif loader) and an @@ -20,12 +20,10 @@ DESCRIPTION = "Pick an LLM role system prompt from roles/*.md (fixed, increment, _ROLE_TIP = ("Role used in 'fixed' mode. The list is built from roles/*.md at " "startup; add a .md file and restart ComfyUI to add a role.") -_MODE_TIP = ("fixed: use the role dropdown. increment: walk the list by 'value', " - "wrapping within start..end. random: pick within start..end from 'value'.") -_VALUE_TIP = ("The stepping input for increment/random modes (ignored in fixed). " - "increment: set control_after_generate to increment to advance one role " - "per generation, wrapping within start..end. random: set it to randomize " - "for a new pick per generation; the same value reproduces the same role.") +_MODE_TIP = ("fixed: use the role dropdown. increment: step to the next role " + "alphabetically each generation, wrapping within start..end. random: a " + "fresh random role within start..end each generation. increment and random " + "advance on their own; no extra widget to set.") _START_TIP = "Lower bound position (0-based) into the alphabetically sorted role list." _END_TIP = "Upper bound position; -1 means the last role. Out-of-range values are clamped." @@ -51,19 +49,18 @@ class LLMSetRoleNode: "required": { "mode": (list(core.MODES), {"default": "fixed", "tooltip": _MODE_TIP}), "role": (_ROLE_OPTIONS, {"default": _DEFAULT_ROLE, "tooltip": _ROLE_TIP}), - "value": ("INT", {"default": 0, "min": 0, "max": _MAX_INT, - "control_after_generate": True, "tooltip": _VALUE_TIP}), "start": ("INT", {"default": 0, "min": 0, "max": _MAX_INT, "tooltip": _START_TIP}), "end": ("INT", {"default": -1, "min": -1, "max": _MAX_INT, "tooltip": _END_TIP}), - } + }, + "hidden": {"unique_id": "UNIQUE_ID"}, } - def set_role(self, mode, role, value, start, end): - return core.select(mode, role, value, start, end) + def set_role(self, mode, role, start, end, unique_id=None): + return core.select(mode, role, start, end, unique_id) @classmethod - def IS_CHANGED(cls, mode, role, value, start, end): - return core.select_fingerprint(mode, role, value, start, end) + def IS_CHANGED(cls, mode, role, start, end, unique_id=None): + return core.select_fingerprint(mode, role, start, end) NODE_CLASS_MAPPINGS = {NODE_ID: LLMSetRoleNode} @@ -75,15 +72,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {NODE_ID: DISPLAY_NAME} try: from comfy_api.v0_0_2 import io, ComfyExtension - def _int_input(name, *, control=False, **kw): - """io.Int.Input, tolerating comfy_api builds without control_after_generate.""" - if control: - try: - return io.Int.Input(name, control_after_generate=True, **kw) - except TypeError: - pass - return io.Int.Input(name, **kw) - class LLMSetRoleV3(io.ComfyNode): @classmethod def define_schema(cls) -> io.Schema: @@ -97,11 +85,10 @@ try: default="fixed", tooltip=_MODE_TIP), io.Combo.Input("role", options=_ROLE_OPTIONS, default=_DEFAULT_ROLE, tooltip=_ROLE_TIP), - _int_input("value", control=True, default=0, min=0, max=_MAX_INT, - tooltip=_VALUE_TIP), - _int_input("start", default=0, min=0, max=_MAX_INT, tooltip=_START_TIP), - _int_input("end", default=-1, min=-1, max=_MAX_INT, tooltip=_END_TIP), + io.Int.Input("start", default=0, min=0, max=_MAX_INT, tooltip=_START_TIP), + io.Int.Input("end", default=-1, min=-1, max=_MAX_INT, tooltip=_END_TIP), ], + hidden=[io.Hidden.unique_id], outputs=[ io.String.Output(id="system_prompt", display_name="system_prompt"), io.String.Output(id="role_name", display_name="role_name"), @@ -110,12 +97,12 @@ try: ) @classmethod - def fingerprint_inputs(cls, mode, role, value, start, end): - return core.select_fingerprint(mode, role, value, start, end) + def fingerprint_inputs(cls, mode, role, start, end): + return core.select_fingerprint(mode, role, start, end) @classmethod - def execute(cls, mode, role, value, start, end) -> io.NodeOutput: - body, title, pos = core.select(mode, role, value, start, end) + def execute(cls, mode, role, start, end) -> io.NodeOutput: + body, title, pos = core.select(mode, role, start, end, cls.hidden.unique_id) return io.NodeOutput(body, title, pos) class LLMSetRoleExtension(ComfyExtension): diff --git a/pyproject.toml b/pyproject.toml index 32fbfa6..e22a308 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui_llmsetrole" -description = "Pick an LLM role (system prompt) from a dropdown and output it as a STRING for an LLM node's system_prompt. Modes: fixed, increment, or random over a start..end slice, with a single value control_after_generate to step roles across a batch. Roles are plain Markdown files in roles/; add a .md file and restart to add a role. Node name: Set Role. Located under Add Node > LLM." -version = "1.3.0" +description = "Pick an LLM role (system prompt) from a dropdown and output it as a STRING for an LLM node's system_prompt. Modes: fixed (use the dropdown), increment (step to the next role alphabetically each generation), or random (a fresh pick each generation), over a start..end slice. Roles are plain Markdown files in roles/; add a .md file and restart to add a role. Node name: Set Role. Located under Add Node > LLM." +version = "1.3.1" license = { file = "LICENSE" } requires-python = ">=3.10" dependencies = [] diff --git a/role_core.py b/role_core.py index 02ef023..6a3e597 100644 --- a/role_core.py +++ b/role_core.py @@ -185,9 +185,9 @@ def label_index(label): def pick_index(mode, value, start, end, n): - """Resolve a role position for increment/random modes from a single `value`. + """Resolve a role position for increment/random modes from an integer `value`. - `value` is the node's one stepping input (control_after_generate drives it): + Pure helper used by select() and the tests: increment: lo + (value mod span), wrapping at the end back to the start. random: lo + Random(value).randrange(span), deterministic per value. Raises ValueError if there are no roles. @@ -203,6 +203,32 @@ def pick_index(mode, value, start, end, n): return lo + (value % span) +# -- Increment state ----------------------------------------------------------- +# Per-node counters for increment mode, keyed by the node's unique_id. Process- +# local: stepping restarts from the first role when ComfyUI restarts. select() +# advances the counter, and the node forces a re-run each generation +# (select_fingerprint returns NaN) so the walk progresses instead of caching. +_increment_counters = {} + + +def reset_counters(): + """Clear all increment state. Used by tests; harmless at runtime.""" + _increment_counters.clear() + + +def _next_increment_pos(node_key, start, end, n): + """Advance this node's counter and return the next position within bounds.""" + c = _increment_counters.get(node_key, 0) + _increment_counters[node_key] = c + 1 + return pick_index("increment", c, start, end, n) + + +def _random_pos(start, end, n): + """A fresh random position within bounds, non-reproducible by design.""" + lo, _, span = _bounds(start, end, n) + return lo + random.randrange(span) + + def role_at(pos): """(label, filename) at a position. Raises ValueError if no roles.""" roles = list_roles() @@ -213,28 +239,42 @@ def role_at(pos): return roles[pos] -def select(mode, role, value, start, end): +def select(mode, role, start, end, node_key=None): """Resolve (system_prompt, role_name, resolved_index) for any mode. - fixed: use the dropdown `role`. increment/random: compute the position from - `value` within [start, end]. + fixed: use the dropdown `role`. increment: the next role alphabetically, + advancing this node's counter and wrapping within [start, end]. random: a + fresh random role within [start, end], picked anew each call. """ if mode == "fixed": pos = label_index(role) body, title = apply_role(role) return body, title, pos - pos = pick_index(mode, value, start, end, len(list_roles())) + n = len(list_roles()) + if n < 1: + raise ValueError( + f"No role files found in {ROLES_DIR}. Add a .md file and restart ComfyUI." + ) + pos = _random_pos(start, end, n) if mode == "random" \ + else _next_increment_pos(node_key, start, end, n) _, fn = role_at(pos) title, body = load_role_text(fn) return body, title, pos -def select_fingerprint(mode, role, value, start, end): - """Cache identity covering inputs and the resolved file's mtime.""" +def select_fingerprint(mode, role, start, end): + """Cache identity. + + fixed caches on the resolved file's mtime so the node only re-runs when the + role or its text changes. increment/random return NaN, which never equals + itself, so ComfyUI re-runs the node every generation and the walk advances. + """ + if mode != "fixed": + return float("nan") try: - _, _, pos = select(mode, role, value, start, end) + pos = label_index(role) _, fn = role_at(pos) mtime = os.path.getmtime(os.path.join(ROLES_DIR, fn)) - return f"{mode}:{pos}:{fn}:{mtime}" + return f"fixed:{pos}:{fn}:{mtime}" except (ValueError, FileNotFoundError, OSError): - return f"{mode}:{role}:{value}:{start}:{end}" + return f"fixed:{role}" diff --git a/test_set_role.py b/test_set_role.py index 43a37bb..acbf7ae 100644 --- a/test_set_role.py +++ b/test_set_role.py @@ -178,22 +178,29 @@ class TestSelect(unittest.TestCase): orig = core.ROLES_DIR core.ROLES_DIR = d core._load_cached.cache_clear() + core.reset_counters() try: - body, name, pos = core.select("fixed", "B", 0, 0, -1) + body, name, pos = core.select("fixed", "B", 0, -1) self.assertEqual((name, pos, body), ("B", 1, "body-b.md")) - body, name, pos = core.select("increment", "A", 2, 0, -1) - self.assertEqual((name, pos), ("C", 2)) + # increment walks alphabetically, advancing once per call, wrapping + core.reset_counters() + walk = [core.select("increment", "A", 0, -1, "node1")[1] for _ in range(4)] + self.assertEqual(walk, ["A", "B", "C", "A"]) - body, name, pos = core.select("increment", "A", 3, 0, -1) - self.assertEqual((name, pos), ("A", 0)) # wrapped + # distinct node ids keep independent counters + core.reset_counters() + self.assertEqual(core.select("increment", "A", 0, -1, "x")[1], "A") + self.assertEqual(core.select("increment", "A", 0, -1, "y")[1], "A") + self.assertEqual(core.select("increment", "A", 0, -1, "x")[1], "B") - _, name_r, pos_r = core.select("random", "A", 7, 0, -1) + _, name_r, pos_r = core.select("random", "A", 0, -1) self.assertIn(pos_r, (0, 1, 2)) self.assertIn(name_r, ("A", "B", "C")) finally: core.ROLES_DIR = orig core._load_cached.cache_clear() + core.reset_counters() class TestShippedRoles(unittest.TestCase):