From 823f71f269312407407dde10022a01436b948e57 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Sun, 2 Aug 2026 22:04:40 +0800 Subject: [PATCH] feat(nodes): make Lora Stack Combiner inputs dynamic --- py/nodes/lora_stack_combiner.py | 95 ++++++++- .../core/loraStackDynamicInputs.test.js | 195 ++++++++++++++++++ tests/nodes/test_lora_stack_combiner.py | 100 ++++++++- web/comfyui/lora_stack_dynamic_inputs.js | 127 ++++++++++++ 4 files changed, 504 insertions(+), 13 deletions(-) create mode 100644 tests/frontend/core/loraStackDynamicInputs.test.js create mode 100644 web/comfyui/lora_stack_dynamic_inputs.js diff --git a/py/nodes/lora_stack_combiner.py b/py/nodes/lora_stack_combiner.py index eda9238e..4c5d8e5d 100644 --- a/py/nodes/lora_stack_combiner.py +++ b/py/nodes/lora_stack_combiner.py @@ -1,27 +1,102 @@ +from __future__ import annotations + +import inspect +import re +from typing import Any + +_STACK_INPUT_PATTERN = re.compile(r"^lora_stack(?:_([ab])|(\d+))$") + + +def _is_stack_input(name: str) -> bool: + return bool(_STACK_INPUT_PATTERN.match(name)) + + +def _stack_slot_number(name: str) -> int: + """Numeric slot used to order stack inputs; legacy a/b map to 1/2.""" + match = _STACK_INPUT_PATTERN.match(name) + if not match: + return -1 + letter, digits = match.group(1), match.group(2) + if digits is not None: + return int(digits) + return 1 if letter == "a" else 2 + + +class _LoraStackOptionalInputs: + """Lookup that preserves explicit optional inputs and dynamic lora_stack slots.""" + + def __init__(self, explicit_inputs: dict[str, tuple[str, dict[str, Any]]]) -> None: + self._explicit_inputs = explicit_inputs + + def __contains__(self, item: object) -> bool: + if not isinstance(item, str): + return False + return item in self._explicit_inputs or _is_stack_input(item) + + def __getitem__(self, key: str) -> tuple[str, dict[str, Any]]: + if key in self._explicit_inputs: + return self._explicit_inputs[key] + if _is_stack_input(key): + return ( + "LORA_STACK", + { + "tooltip": "A LoRA stack to combine. Connect to add more inputs.", + }, + ) + raise KeyError(key) + + class LoraStackCombinerLM: NAME = "Lora Stack Combiner (LoraManager)" CATEGORY = "Lora Manager/stackers" + DESCRIPTION = ( + "Combines multiple LoRA stacks into a single stack. " + "Supports dynamic inputs: connect a stack to add more inputs." + ) @classmethod def INPUT_TYPES(cls): + optional_inputs: dict[str, tuple[str, dict[str, Any]]] = { + "lora_stack1": ( + "LORA_STACK", + { + "tooltip": "A LoRA stack to combine. Connect to add more inputs.", + }, + ), + "lora_stack2": ( + "LORA_STACK", + { + "tooltip": "A LoRA stack to combine. Connect to add more inputs.", + }, + ), + } + + stack = inspect.stack() + if len(stack) > 2 and stack[2].function == "get_input_info": + optional_inputs = _LoraStackOptionalInputs(optional_inputs) # type: ignore[assignment] + return { "required": {}, - "optional": { - "lora_stack_a": ("LORA_STACK",), - "lora_stack_b": ("LORA_STACK",), - }, + "optional": optional_inputs, } RETURN_TYPES = ("LORA_STACK",) RETURN_NAMES = ("LORA_STACK",) FUNCTION = "combine_stacks" - def combine_stacks(self, lora_stack_a=None, lora_stack_b=None): - combined_stack = [] + def combine_stacks(self, lora_stack1=None, lora_stack2=None, **kwargs): + stacks = { + "lora_stack1": lora_stack1, + "lora_stack2": lora_stack2, + } + for key, value in kwargs.items(): + if _is_stack_input(key) and value is not None: + stacks[key] = value - if lora_stack_a: - combined_stack.extend(lora_stack_a) - if lora_stack_b: - combined_stack.extend(lora_stack_b) + combined_stack = [] + for key in sorted(stacks, key=_stack_slot_number): + stack = stacks[key] + if stack: + combined_stack.extend(stack) return (combined_stack,) diff --git a/tests/frontend/core/loraStackDynamicInputs.test.js b/tests/frontend/core/loraStackDynamicInputs.test.js new file mode 100644 index 00000000..9e92607b --- /dev/null +++ b/tests/frontend/core/loraStackDynamicInputs.test.js @@ -0,0 +1,195 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const { APP_MODULE, EXTENSION_MODULE, appMock, registeredExtensions } = + vi.hoisted(() => { + const registeredExtensions = []; + const appMock = { + configuringGraph: false, + registerExtension: (ext) => registeredExtensions.push(ext), + }; + return { + APP_MODULE: new URL("../../../scripts/app.js", import.meta.url).pathname, + EXTENSION_MODULE: new URL( + "../../../web/comfyui/lora_stack_dynamic_inputs.js", + import.meta.url + ).pathname, + appMock, + registeredExtensions, + }; + }); + +vi.mock(APP_MODULE, () => ({ + app: appMock, +})); + +describe("Lora Stack Combiner dynamic inputs", () => { + let extension; + + beforeEach(async () => { + vi.resetModules(); + registeredExtensions.length = 0; + appMock.configuringGraph = false; + await import(EXTENSION_MODULE); + extension = registeredExtensions.find( + (ext) => ext.name === "Comfy.LoraManager.LoraStackCombiner" + ); + expect(extension).toBeDefined(); + }); + + function createNodeType() { + const nodeType = { prototype: {} }; + extension.beforeRegisterNodeDef( + nodeType, + { name: "Lora Stack Combiner (LoraManager)" }, + appMock + ); + return nodeType; + } + + function createNode(inputs = []) { + const node = { + comfyClass: "Lora Stack Combiner (LoraManager)", + inputs: inputs.map((name) => ({ name, type: "LORA_STACK" })), + addInput: vi.fn(function (name, type, opts) { + this.inputs.push({ name, type, ...opts }); + }), + removeInput: vi.fn(function (index) { + this.inputs.splice(index, 1); + }), + }; + return node; + } + + function makeLinkInfo() { + return { id: 999, origin_id: 1, target_id: 2 }; + } + + it("adds a third input when the last slot gets connected", () => { + const nodeType = createNodeType(); + const node = createNode(["lora_stack1", "lora_stack2"]); + node.onConnectionsChange = nodeType.prototype.onConnectionsChange; + + node.onConnectionsChange(1, 1, true, makeLinkInfo()); + + expect(node.inputs.map((input) => input.name)).toEqual([ + "lora_stack1", + "lora_stack2", + "lora_stack3", + ]); + }); + + it("does not add an input when a non-last slot gets connected", () => { + const nodeType = createNodeType(); + const node = createNode(["lora_stack1", "lora_stack2", "lora_stack3"]); + node.onConnectionsChange = nodeType.prototype.onConnectionsChange; + + node.onConnectionsChange(1, 0, true, makeLinkInfo()); + + expect(node.inputs.map((input) => input.name)).toEqual([ + "lora_stack1", + "lora_stack2", + "lora_stack3", + ]); + }); + + it("removes a disconnected middle slot and renumbers", () => { + // Simulates a real LiteGraph disconnect event: it fires only for slots that + // had a link, and input.link has already been cleared before the event fires. + const nodeType = createNodeType(); + const node = createNode(["lora_stack1", "lora_stack2", "lora_stack3"]); + node.inputs[0].link = 11; + node.inputs[1].link = null; // slot 2 was just disconnected + node.inputs[2].link = 13; + node.onConnectionsChange = nodeType.prototype.onConnectionsChange; + + node.onConnectionsChange(1, 1, false, makeLinkInfo()); + + expect(node.inputs.map((input) => input.name)).toEqual([ + "lora_stack1", + "lora_stack2", + ]); + }); + + it("keeps the last slot when it is disconnected", () => { + const nodeType = createNodeType(); + const node = createNode(["lora_stack1", "lora_stack2", "lora_stack3"]); + node.inputs[0].link = 11; + node.inputs[1].link = 12; + node.inputs[2].link = null; // last slot was just disconnected + node.onConnectionsChange = nodeType.prototype.onConnectionsChange; + + node.onConnectionsChange(1, 2, false, makeLinkInfo()); + + expect(node.inputs.map((input) => input.name)).toEqual([ + "lora_stack1", + "lora_stack2", + "lora_stack3", + ]); + expect(node.removeInput).not.toHaveBeenCalled(); + }); + + it("keeps at least two inputs when disconnecting", () => { + const nodeType = createNodeType(); + const node = createNode(["lora_stack1", "lora_stack2"]); + node.inputs[0].link = 11; + node.inputs[1].link = null; // slot 2 was just disconnected + node.onConnectionsChange = nodeType.prototype.onConnectionsChange; + + node.onConnectionsChange(1, 1, false, makeLinkInfo()); + + expect(node.inputs.map((input) => input.name)).toEqual([ + "lora_stack1", + "lora_stack2", + ]); + expect(node.removeInput).not.toHaveBeenCalled(); + }); + + it("does nothing while the graph is being configured", () => { + appMock.configuringGraph = true; + const nodeType = createNodeType(); + const node = createNode(["lora_stack1", "lora_stack2"]); + node.onConnectionsChange = nodeType.prototype.onConnectionsChange; + + node.onConnectionsChange(1, 1, true, makeLinkInfo()); + + expect(node.inputs.map((input) => input.name)).toEqual([ + "lora_stack1", + "lora_stack2", + ]); + expect(node.addInput).not.toHaveBeenCalled(); + }); + + it("leaves legacy lora_stack_a/b inputs untouched", () => { + const nodeType = createNodeType(); + const node = createNode(["lora_stack_a", "lora_stack_b"]); + node.onConnectionsChange = nodeType.prototype.onConnectionsChange; + + node.onConnectionsChange(1, 0, true, makeLinkInfo()); + + expect(node.inputs.map((input) => input.name)).toEqual([ + "lora_stack_a", + "lora_stack_b", + ]); + expect(node.addInput).not.toHaveBeenCalled(); + }); + + it("ensures two numbered inputs exist on creation", () => { + const node = createNode([]); + extension.nodeCreated(node, {}); + + expect(node.inputs.map((input) => input.name)).toEqual([ + "lora_stack1", + "lora_stack2", + ]); + }); + + it("does not add numbered inputs to legacy workflows", () => { + const node = createNode(["lora_stack_a", "lora_stack_b"]); + extension.nodeCreated(node, {}); + + expect(node.inputs.map((input) => input.name)).toEqual([ + "lora_stack_a", + "lora_stack_b", + ]); + }); +}); diff --git a/tests/nodes/test_lora_stack_combiner.py b/tests/nodes/test_lora_stack_combiner.py index 609f7385..79e52ef9 100644 --- a/tests/nodes/test_lora_stack_combiner.py +++ b/tests/nodes/test_lora_stack_combiner.py @@ -1,4 +1,11 @@ -from py.nodes.lora_stack_combiner import LoraStackCombinerLM +import types + +import pytest + +from py.nodes.lora_stack_combiner import ( + LoraStackCombinerLM, + _LoraStackOptionalInputs, +) def test_combine_stacks_preserves_order(): @@ -63,8 +70,95 @@ def test_combine_stacks_returns_other_when_one_unconnected(): node = LoraStackCombinerLM() stack_a = [("folder/a.safetensors", 0.7, 0.6)] - (combined_stack_a,) = node.combine_stacks(lora_stack_a=stack_a) - (combined_stack_b,) = node.combine_stacks(lora_stack_b=stack_a) + (combined_stack_a,) = node.combine_stacks(lora_stack1=stack_a) + (combined_stack_b,) = node.combine_stacks(lora_stack2=stack_a) assert combined_stack_a == stack_a assert combined_stack_b == stack_a + + +def test_combine_stacks_with_dynamic_third_slot(): + node = LoraStackCombinerLM() + stack_a = [("folder/a.safetensors", 0.7, 0.6)] + stack_b = [("folder/b.safetensors", 0.8, 0.8)] + stack_c = [("folder/c.safetensors", 1.0, 0.9)] + + (combined_stack,) = node.combine_stacks( + lora_stack1=stack_a, lora_stack2=stack_b, lora_stack3=stack_c + ) + + assert combined_stack == stack_a + stack_b + stack_c + + +def test_combine_stacks_orders_by_slot_number_not_call_order(): + node = LoraStackCombinerLM() + stack_a = [("folder/a.safetensors", 0.7, 0.6)] + stack_b = [("folder/b.safetensors", 0.8, 0.8)] + stack_c = [("folder/c.safetensors", 1.0, 0.9)] + + (combined_stack,) = node.combine_stacks( + lora_stack3=stack_c, lora_stack2=stack_b, lora_stack1=stack_a + ) + + assert combined_stack == stack_a + stack_b + stack_c + + +def test_combine_stacks_accepts_only_dynamic_slot(): + node = LoraStackCombinerLM() + stack_c = [("folder/c.safetensors", 1.0, 0.9)] + + (combined_stack,) = node.combine_stacks(lora_stack3=stack_c) + + assert combined_stack == stack_c + + +def test_combine_stacks_handles_legacy_input_names(): + node = LoraStackCombinerLM() + stack_a = [("folder/a.safetensors", 0.7, 0.6)] + stack_b = [("folder/b.safetensors", 0.8, 0.8)] + + (combined_stack,) = node.combine_stacks(lora_stack_a=stack_a, lora_stack_b=stack_b) + + assert combined_stack == stack_a + stack_b + + +def test_input_types_exposes_two_default_slots(): + input_types = LoraStackCombinerLM.INPUT_TYPES() + + assert set(input_types["optional"]) == {"lora_stack1", "lora_stack2"} + assert input_types["optional"]["lora_stack1"][0] == "LORA_STACK" + assert input_types["optional"]["lora_stack2"][0] == "LORA_STACK" + + +def test_input_types_recognizes_dynamic_slots_from_get_input_info(monkeypatch): + frames = [None, None, types.SimpleNamespace(function="get_input_info")] + monkeypatch.setattr( + "py.nodes.lora_stack_combiner.inspect.stack", lambda: frames + ) + + input_types = LoraStackCombinerLM.INPUT_TYPES() + optional = input_types["optional"] + + assert "lora_stack3" in optional + assert optional["lora_stack3"][0] == "LORA_STACK" + assert "lora_stack25" in optional + assert optional["lora_stack25"][0] == "LORA_STACK" + + +def test_lora_stack_optional_inputs_proxy(): + proxy = _LoraStackOptionalInputs({"lora_stack1": ("LORA_STACK", {})}) + + assert "lora_stack1" in proxy + assert "lora_stack2" in proxy + assert "lora_stack10" in proxy + assert "lora_stack_a" in proxy + assert "lora_stack" not in proxy + assert "lora_stacka" not in proxy + assert "lora_stack_1" not in proxy + assert "text" not in proxy + + assert proxy["lora_stack1"][0] == "LORA_STACK" + assert proxy["lora_stack5"][0] == "LORA_STACK" + + with pytest.raises(KeyError): + proxy["not_a_stack"] diff --git a/web/comfyui/lora_stack_dynamic_inputs.js b/web/comfyui/lora_stack_dynamic_inputs.js new file mode 100644 index 00000000..50bf9e6c --- /dev/null +++ b/web/comfyui/lora_stack_dynamic_inputs.js @@ -0,0 +1,127 @@ +import { app } from "../../scripts/app.js"; + +/** + * Extension for LoraStackCombinerLM node to support dynamic lora_stack inputs. + * Defaults to two inputs; connecting the last slot adds a new empty one, and + * disconnecting a non-last slot removes it (at least two are always kept). + * Based on the dynamic input pattern from Impact Pack's Switch (Any) node. + */ +const STACK_INPUT_PATTERN = /^lora_stack\d+$/; + +app.registerExtension({ + name: "Comfy.LoraManager.LoraStackCombiner", + + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name !== "Lora Stack Combiner (LoraManager)") { + return; + } + + const onConnectionsChange = nodeType.prototype.onConnectionsChange; + + nodeType.prototype.onConnectionsChange = function(type, index, connected, link_info) { + // Skip while the graph is being (re)configured (load, paste, subgraph ops) + if (app.configuringGraph) { + return onConnectionsChange?.apply?.(this, arguments); + } + + const stackTrace = new Error().stack; + + // Skip during graph loading/pasting to avoid interference + if (stackTrace.includes('loadGraphData') || stackTrace.includes('pasteFromClipboard')) { + return onConnectionsChange?.apply?.(this, arguments); + } + + // Skip subgraph operations + if (stackTrace.includes('convertToSubgraph') || stackTrace.includes('Subgraph.configure')) { + return onConnectionsChange?.apply?.(this, arguments); + } + + if (!link_info) { + return onConnectionsChange?.apply?.(this, arguments); + } + + // Handle input connections (type === 1) + if (type === 1) { + const input = this.inputs[index]; + + // Only process numbered lora_stack inputs (legacy a/b slots are left untouched) + if (!input || !STACK_INPUT_PATTERN.test(input.name)) { + return onConnectionsChange?.apply?.(this, arguments); + } + + // Count existing numbered lora_stack inputs + let stackInputCount = 0; + for (const inp of this.inputs) { + if (STACK_INPUT_PATTERN.test(inp.name)) { + stackInputCount++; + } + } + + // Renumber all numbered lora_stack inputs sequentially + let slotIndex = 1; + for (const inp of this.inputs) { + if (STACK_INPUT_PATTERN.test(inp.name)) { + inp.name = `lora_stack${slotIndex}`; + slotIndex++; + } + } + + // Add new input slot if connected and this was the last one + if (connected) { + const lastStackIndex = stackInputCount; + if (index === lastStackIndex || index === this.inputs.findIndex(i => i.name === `lora_stack${lastStackIndex}`)) { + this.addInput(`lora_stack${slotIndex}`, "LORA_STACK", { + tooltip: "A LoRA stack to combine. Connect to add more inputs." + }); + } + } + + // Remove disconnected input slots (but keep at least two). + // LiteGraph fires this event only for slots that had a link, and + // it has already cleared input.link by the time the event fires, + // so the disconnected slot is always empty at this point. + if (!connected && stackInputCount > 2) { + const disconnectedInput = this.inputs[index]; + if (disconnectedInput && STACK_INPUT_PATTERN.test(disconnectedInput.name)) { + // Keep the last slot so there is always an empty slot to reconnect into + const isLastStackSlot = index === this.inputs.findLastIndex(i => STACK_INPUT_PATTERN.test(i.name)); + if (!isLastStackSlot) { + this.removeInput(index); + + // Renumber again after removal + let newSlotIndex = 1; + for (const inp of this.inputs) { + if (STACK_INPUT_PATTERN.test(inp.name)) { + inp.name = `lora_stack${newSlotIndex}`; + newSlotIndex++; + } + } + } + } + } + } + + return onConnectionsChange?.apply?.(this, arguments); + }; + }, + + nodeCreated(node, app) { + if (node.comfyClass !== "Lora Stack Combiner (LoraManager)") { + return; + } + + // Leave legacy (a/b) workflows untouched + const hasLegacyInputs = node.inputs.some(inp => inp.name === "lora_stack_a" || inp.name === "lora_stack_b"); + if (hasLegacyInputs) { + return; + } + + // Ensure at least two numbered lora_stack inputs exist on creation + const stackInputCount = node.inputs.filter(inp => STACK_INPUT_PATTERN.test(inp.name)).length; + for (let i = stackInputCount + 1; i <= 2; i++) { + node.addInput(`lora_stack${i}`, "LORA_STACK", { + tooltip: "A LoRA stack to combine. Connect to add more inputs." + }); + } + } +});