diff --git a/py/nodes/lora_stack_combiner.py b/py/nodes/lora_stack_combiner.py index 9b5c4fbc..eda9238e 100644 --- a/py/nodes/lora_stack_combiner.py +++ b/py/nodes/lora_stack_combiner.py @@ -5,7 +5,8 @@ class LoraStackCombinerLM: @classmethod def INPUT_TYPES(cls): return { - "required": { + "required": {}, + "optional": { "lora_stack_a": ("LORA_STACK",), "lora_stack_b": ("LORA_STACK",), }, @@ -15,7 +16,7 @@ class LoraStackCombinerLM: RETURN_NAMES = ("LORA_STACK",) FUNCTION = "combine_stacks" - def combine_stacks(self, lora_stack_a, lora_stack_b): + def combine_stacks(self, lora_stack_a=None, lora_stack_b=None): combined_stack = [] if lora_stack_a: diff --git a/tests/nodes/test_lora_stack_combiner.py b/tests/nodes/test_lora_stack_combiner.py index b5b83954..609f7385 100644 --- a/tests/nodes/test_lora_stack_combiner.py +++ b/tests/nodes/test_lora_stack_combiner.py @@ -49,3 +49,22 @@ def test_combine_stacks_allows_duplicate_entries(): (combined_stack,) = node.combine_stacks([duplicate_entry], [duplicate_entry]) assert combined_stack == [duplicate_entry, duplicate_entry] + + +def test_combine_stacks_returns_empty_when_both_unconnected(): + node = LoraStackCombinerLM() + + (combined_stack,) = node.combine_stacks() + + assert combined_stack == [] + + +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) + + assert combined_stack_a == stack_a + assert combined_stack_b == stack_a