feat: add native LoRA stack loader nodes (CORE-186)

comfy_extras/nodes_lora_stack.py

@@ -0,0 +1,106 @@
+from typing_extensions import override
+
+import comfy.sd
+import comfy.utils
+import folder_paths
+from comfy_api.latest import ComfyExtension, io
+from comfy_api.latest._io import _DynamicGroup
+
+
+def _load_lora_file(lora_name: str):
+    lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
+    return comfy.utils.load_torch_file(lora_path, safe_load=True, return_metadata=True)
+
+
+def _lora_template() -> list[io.Input]:
+    return [
+        io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
+                       tooltip="The name of the LoRA file to apply."),
+        io.Float.Input("strength", default=1.0, min=-100.0, max=100.0, step=0.01,
+                       tooltip="How strongly to apply this LoRA. 0 = off, negative inverts the effect."),
+    ]
+
+
+class LoadLoraModel(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LoadLoraModel",
+            display_name="Load LoRA (Model)",
+            search_aliases=["lora", "load lora", "apply lora", "lora model", "lora stack"],
+            category="model/loaders",
+            description="Apply a stack of LoRAs to a diffusion model. Add one row per LoRA; "
+                        "each row picks a LoRA file and its strength.",
+            inputs=[
+                io.Model.Input("model", tooltip="The diffusion model the LoRAs will be applied to."),
+                _DynamicGroup.Input(
+                    "loras",
+                    template=_lora_template(),
+                    min=1,
+                    max=20,
+                    tooltip="Each row applies one LoRA to the model.",
+                    group_name="LoRA",
+                ),
+            ],
+            outputs=[io.Model.Output(tooltip="The modified diffusion model.")],
+        )
+
+    @classmethod
+    def execute(cls, model, loras: list[dict]) -> io.NodeOutput:
+        for row in loras:
+            lora_name = row.get("lora_name")
+            strength = row["strength"]
+            if not lora_name or strength == 0:
+                continue
+            lora, metadata = _load_lora_file(lora_name)
+            model, _ = comfy.sd.load_lora_for_models(model, None, lora, strength, 0, lora_metadata=metadata)
+        return io.NodeOutput(model)
+
+
+class LoadLoraTextEncoder(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LoadLoraTextEncoder",
+            display_name="Load LoRA (Text Encoder)",
+            search_aliases=["lora", "load lora", "apply lora", "clip lora", "lora stack"],
+            category="model/loaders",
+            description="Apply a stack of LoRAs to a CLIP text encoder. Add one row per LoRA; "
+                        "each row picks a LoRA file and its strength.",
+            inputs=[
+                io.Clip.Input("clip", tooltip="The CLIP text encoder the LoRAs will be applied to."),
+                _DynamicGroup.Input(
+                    "loras",
+                    template=_lora_template(),
+                    min=1,
+                    max=20,
+                    tooltip="Each row applies one LoRA to the text encoder.",
+                    group_name="LoRA",
+                ),
+            ],
+            outputs=[io.Clip.Output(tooltip="The modified CLIP text encoder.")],
+        )
+
+    @classmethod
+    def execute(cls, clip, loras: list[dict]) -> io.NodeOutput:
+        for row in loras:
+            lora_name = row.get("lora_name")
+            strength = row["strength"]
+            if not lora_name or strength == 0:
+                continue
+            lora, metadata = _load_lora_file(lora_name)
+            _, clip = comfy.sd.load_lora_for_models(None, clip, lora, 0, strength, lora_metadata=metadata)
+        return io.NodeOutput(clip)
+
+
+class LoraStackExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            LoadLoraModel,
+            LoadLoraTextEncoder,
+        ]
+
+
+async def comfy_entrypoint() -> LoraStackExtension:
+    return LoraStackExtension()

nodes.py

@@ -2459,6 +2459,7 @@ async def init_builtin_extra_nodes():
         "nodes_eps.py",
         "nodes_flux.py",
         "nodes_lora_extract.py",
+        "nodes_lora_stack.py",
         "nodes_torch_compile.py",
         "nodes_mochi.py",
         "nodes_slg.py",

tests-unit/comfy_extras_test/nodes_lora_stack_test.py

@@ -0,0 +1,103 @@
+import asyncio
+from unittest.mock import Mock, call
+
+import pytest
+import torch
+
+from comfy.cli_args import args
+
+if not torch.cuda.is_available():
+    args.cpu = True
+
+from comfy_api.latest._io import build_nested_inputs, create_input_dict_v1, get_finalized_class_inputs
+from comfy_extras import nodes_lora_stack
+
+
+@pytest.mark.parametrize('node_class,target', [
+    (nodes_lora_stack.LoadLoraModel, 'model'),
+    (nodes_lora_stack.LoadLoraTextEncoder, 'clip'),
+])
+def test_applies_sparse_rows_in_order_with_metadata(monkeypatch, node_class, target):
+    monkeypatch.setattr(nodes_lora_stack.folder_paths, 'get_filename_list', lambda _: ['a.safetensors', 'b.safetensors'])
+    schema = create_input_dict_v1(node_class.define_schema().inputs)
+    values = {
+        'loras.0.lora_name': 'a.safetensors', 'loras.0.strength': 0.5,
+        'loras.2.lora_name': 'b.safetensors', 'loras.2.strength': -0.25,
+    }
+    _, _, v3_data = get_finalized_class_inputs(schema, values)
+    rows = build_nested_inputs(values, v3_data)['loras']
+    assert rows[1] == {'lora_name': None, 'strength': None}
+
+    original, first, second = object(), object(), object()
+    lora_a, lora_b = object(), object()
+    metadata_a, metadata_b = {'name': 'a'}, {'name': 'b'}
+    resolve = Mock(side_effect=['/loras/a.safetensors', '/loras/b.safetensors'])
+    load = Mock(side_effect=[(lora_a, metadata_a), (lora_b, metadata_b)])
+    apply = Mock(side_effect=[(first, None), (second, None)] if target == 'model' else [(None, first), (None, second)])
+    monkeypatch.setattr(nodes_lora_stack.folder_paths, 'get_full_path_or_raise', resolve)
+    monkeypatch.setattr(nodes_lora_stack.comfy.utils, 'load_torch_file', load)
+    monkeypatch.setattr(nodes_lora_stack.comfy.sd, 'load_lora_for_models', apply)
+
+    result = node_class.execute(original, rows)
+
+    assert result.result == (second,)
+    assert resolve.call_args_list == [call('loras', 'a.safetensors'), call('loras', 'b.safetensors')]
+    assert load.call_args_list == [
+        call('/loras/a.safetensors', safe_load=True, return_metadata=True),
+        call('/loras/b.safetensors', safe_load=True, return_metadata=True),
+    ]
+    if target == 'model':
+        assert apply.call_args_list == [
+            call(original, None, lora_a, 0.5, 0, lora_metadata=metadata_a),
+            call(first, None, lora_b, -0.25, 0, lora_metadata=metadata_b),
+        ]
+    else:
+        assert apply.call_args_list == [
+            call(None, original, lora_a, 0, 0.5, lora_metadata=metadata_a),
+            call(None, first, lora_b, 0, -0.25, lora_metadata=metadata_b),
+        ]
+
+
+@pytest.mark.parametrize('node_class', [nodes_lora_stack.LoadLoraModel, nodes_lora_stack.LoadLoraTextEncoder])
+def test_accepts_twentieth_position_and_rejects_later_indices(monkeypatch, node_class):
+    monkeypatch.setattr(nodes_lora_stack.folder_paths, 'get_filename_list', lambda _: ['a.safetensors'])
+    schema = create_input_dict_v1(node_class.define_schema().inputs)
+    values = {'loras.19.lora_name': 'a.safetensors', 'loras.19.strength': 0.5}
+    _, _, v3_data = get_finalized_class_inputs(schema, values)
+    rows = build_nested_inputs(values, v3_data)['loras']
+    assert rows == [{'lora_name': None, 'strength': None}] * 19 + [{'lora_name': 'a.safetensors', 'strength': 0.5}]
+
+    with pytest.raises(ValueError, match='exceeds the index limit of 19'):
+        get_finalized_class_inputs(schema, {'loras.20.lora_name': 'a.safetensors', 'loras.20.strength': 0.5})
+
+
+@pytest.mark.parametrize('node_class', [nodes_lora_stack.LoadLoraModel, nodes_lora_stack.LoadLoraTextEncoder])
+def test_skips_empty_positions_and_zero_strength_without_loading(monkeypatch, node_class):
+    load = Mock()
+    monkeypatch.setattr(nodes_lora_stack.folder_paths, 'get_full_path_or_raise', load)
+    original = object()
+    result = node_class.execute(original, [
+        {'lora_name': None, 'strength': None},
+        {'lora_name': 'disabled.safetensors', 'strength': 0},
+    ])
+    assert result.result == (original,)
+    load.assert_not_called()
+
+
+@pytest.mark.parametrize('node_class', [nodes_lora_stack.LoadLoraModel, nodes_lora_stack.LoadLoraTextEncoder])
+def test_missing_lora_fails_at_the_file_resolver(monkeypatch, node_class):
+    resolve = Mock(side_effect=FileNotFoundError('Missing LoRA'))
+    load = Mock()
+    monkeypatch.setattr(nodes_lora_stack.folder_paths, 'get_full_path_or_raise', resolve)
+    monkeypatch.setattr(nodes_lora_stack.comfy.utils, 'load_torch_file', load)
+    with pytest.raises(FileNotFoundError, match='Missing LoRA'):
+        node_class.execute(object(), [{'lora_name': 'missing.safetensors', 'strength': 1}])
+    resolve.assert_called_once_with('loras', 'missing.safetensors')
+    load.assert_not_called()
+
+
+def test_extension_exports_both_loaders():
+    extension = asyncio.run(nodes_lora_stack.comfy_entrypoint())
+    assert asyncio.run(extension.get_node_list()) == [
+        nodes_lora_stack.LoadLoraModel, nodes_lora_stack.LoadLoraTextEncoder,
+    ]