feat(nodes): make Lora Stack Combiner inputs dynamic

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,)

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",
+    ]);
+  });
+});

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"]

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."
+            });
+        }
+    }
+});