From ebb29393aac0c74155898bba734c3080ca91402f Mon Sep 17 00:00:00 2001 From: silveroxides Date: Fri, 4 Sep 2026 15:55:42 +0200 Subject: [PATCH] feat(lora): add Lodestone delta mergers --- __init__.py | 3 + docs/lodestone_mergers.md | 37 ++++ nodes/lodestone_merger.py | 309 +++++++++++++++++++++++++++++++++ tests/test_artifact_paths.py | 1 + tests/test_groups.toml | 8 +- tests/test_lodestone_merger.py | 157 +++++++++++++++++ 6 files changed, 511 insertions(+), 4 deletions(-) create mode 100644 docs/lodestone_mergers.md create mode 100644 nodes/lodestone_merger.py create mode 100644 tests/test_lodestone_merger.py diff --git a/__init__.py b/__init__.py index 25da4c4..2d026aa 100644 --- a/__init__.py +++ b/__init__.py @@ -21,6 +21,7 @@ from .nodes.merger import ( CheckpointThreeMerger, EmbeddingThreeMerger ) from .nodes.consensus_merger import CWB_MERGER_NODES +from .nodes.lodestone_merger import LODESTONE_MERGER_NODES from .nodes.model_analysis import MODEL_ANALYSIS_NODES from .nodes.lora_extract_svd import ( LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile, @@ -81,6 +82,8 @@ class ModelUtilsExtension(ComfyExtension): CheckpointThreeMerger, EmbeddingThreeMerger, # Dedicated CWB Mergers *CWB_MERGER_NODES, + # Dedicated Lodestone LoRA Mergers + *LODESTONE_MERGER_NODES, # Two-Model Analysis *MODEL_ANALYSIS_NODES, # LoRA Extraction diff --git a/docs/lodestone_mergers.md b/docs/lodestone_mergers.md new file mode 100644 index 0000000..e7ab80a --- /dev/null +++ b/docs/lodestone_mergers.md @@ -0,0 +1,37 @@ +# Lodestone LoRA Mergers + +The dedicated Lodestone nodes merge full logical LoRA deltas rather than +blending down and up factors independently. Stored alpha is normalized before +each factor pair is expanded. Results are saved as canonical full direct-difference +tensors and contain no factor or alpha tensors. + +The 2-, 3-, and Multi-Merge nodes treat every selected LoRA as one equal-prior +slot, matching the supplied Lodestone calculation. + +## Calculation Modes + +- `sum`: sum the source deltas. +- `mean`: divide the summed delta by the number of active contributors. +- `slotnorm`: scale each source delta to the median source Frobenius norm, then + average. +- `normmatch`: retain the summed direction and set its Frobenius norm to the + median source norm. +- `slotnorm-normmatch`: equalize source norms, sum, then set the final norm to + the original median source norm. + +A zero target or merged norm remains zero. + +## Memory Behavior + +UEL streams one logical layer at a time and writes each completed full-delta +tensor incrementally through the atomic writer. Dense tensors used for the +merge and Frobenius norms exist only within the active layer. A CUDA +out-of-memory error releases CUDA state and retries that logical layer on CPU. + +## Fixed-Input Controls + +LoRA 1 anchors output layers and filtering behavior. `skip` preserves LoRA 1 +when another input lacks a required layer, `zeros` supplies an explicit zero +slot, and `error` aborts. Source factors may have different ranks when their +expanded dense shapes agree. Exclude and discard patterns apply to complete +logical layers. diff --git a/nodes/lodestone_merger.py b/nodes/lodestone_merger.py new file mode 100644 index 0000000..cf8940c --- /dev/null +++ b/nodes/lodestone_merger.py @@ -0,0 +1,309 @@ +import fnmatch +import gc +import re + +import comfy.utils +import folder_paths +import torch +from comfy_api.latest import io +from tqdm import tqdm +from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned + +from .artifact_paths import canonical_model_artifact_path +from .device_utils import cleanup_after_operation, estimate_model_size, prepare_for_large_operation +from .lora_alpha import normalize_lora_pair +from .lora_resize import build_lora_layer_map, canonical_lora_key, parse_lora_layers, select_output_dtype, validate_canonical_blocks +from .quantization_guard import inspect_low_bit_input +from .uel_io import AsyncTensorCursor, atomic_uel_writer + + +LODESTONE_METHODS = ("sum", "mean", "slotnorm", "normmatch", "slotnorm-normmatch") +_ALLOWED_ROLES = { + "down", "up", "alpha", "diff", "format", "down_suffix", "up_suffix", + "alpha_suffix", +} + + +def _matches(key, patterns, glob_mode): + entries = [entry.strip() for entry in patterns.splitlines() if entry.strip()] + if glob_mode: + return any(fnmatch.fnmatchcase(key, entry) for entry in entries) + return any(re.search(entry, key) is not None for entry in entries) + + +def factor_delta(down, up, layer): + if down.ndim == 2 and up.ndim == 2: + if down.shape[0] != up.shape[1]: + raise ValueError(f"Lodestone rank mismatch for '{layer}'.") + return up.float() @ down.float() + if down.ndim == 4 and up.ndim == 4 and tuple(up.shape[2:]) == (1, 1): + rank, channels, height, width = down.shape + if up.shape[1] != rank: + raise ValueError(f"Lodestone rank mismatch for '{layer}'.") + delta = up.reshape(up.shape[0], rank).float() @ down.reshape(rank, -1).float() + return delta.reshape(up.shape[0], channels, height, width) + raise ValueError( + f"Lodestone layer '{layer}' requires linear factors or convolutional " + "factors with a 1x1 up kernel." + ) + + +def merge_lodestone_deltas(deltas, method): + if method not in LODESTONE_METHODS: + raise ValueError(f"Unknown Lodestone method '{method}'.") + if not deltas: + raise ValueError("Lodestone requires at least one delta.") + shape = tuple(deltas[0].shape) + if any(tuple(delta.shape) != shape for delta in deltas[1:]): + raise ValueError("Lodestone source deltas must have matching shapes.") + norms = torch.stack([torch.linalg.vector_norm(delta.float()) for delta in deltas]) + target = torch.median(norms) + values = deltas + if method in {"slotnorm", "slotnorm-normmatch"}: + values = [ + delta * (target / norm.clamp_min(1e-12)) + for delta, norm in zip(deltas, norms) + ] + merged = torch.stack(values).sum(dim=0) + if method in {"mean", "slotnorm"}: + merged /= len(values) + elif method in {"normmatch", "slotnorm-normmatch"}: + merged_norm = torch.linalg.vector_norm(merged.float()) + if merged_norm > 0: + merged *= target / merged_norm.clamp_min(1e-12) + return merged + + +def _to_device(tensor, device): + if device.type == "cuda": + return transfer_to_gpu_pinned(tensor, device, torch.float32) + return tensor.to(device=device, dtype=torch.float32) + + +def _release_cuda_oom(): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def _is_cuda_oom(error): + return isinstance(error, torch.cuda.OutOfMemoryError) or ( + isinstance(error, RuntimeError) and "out of memory" in str(error).lower() + ) + + +def _build_units(layer_map, infos, handlers, mismatch_mode, exclude_patterns, discard_patterns, glob_patterns): + units = [] + stream_keys = [[] for _ in handlers] + for core, mapped in layer_map.items(): + if not any(index == 0 for index, _ in mapped): + continue + output_key = canonical_lora_key(core, "diff") + if _matches(output_key, discard_patterns, glob_patterns): + continue + preserve = _matches(output_key, exclude_patterns, glob_patterns) + sources = [] + for index, block in sorted(mapped): + roles = infos[index]["pairs"][block] + unsupported = set(roles) - _ALLOWED_ROLES + if unsupported: + raise ValueError(f"Lodestone layer '{block}' has unsupported roles: {', '.join(sorted(unsupported))}.") + if "down" in roles or "up" in roles: + if not all(role in roles for role in ("down", "up")): + raise ValueError(f"Lodestone layer '{block}' has an incomplete factor pair.") + keys = [roles["down"], roles["up"]] + if any(key in infos[index]["low_bit"] for key in keys): + raise ValueError(f"Lodestone cannot expand low-bit factors for '{block}'.") + if "alpha" in roles: + keys.append(roles["alpha"]) + sources.append((index, block, "pair", keys)) + elif "diff" in roles: + key = roles["diff"] + if key in infos[index]["low_bit"]: + raise ValueError(f"Lodestone cannot merge low-bit delta '{block}'.") + sources.append((index, block, "diff", [key])) + if preserve: + sources = [source for source in sources if source[0] == 0] + present = {source[0] for source in sources} + if len(present) < len(handlers) and not preserve: + if mismatch_mode == "error": + raise ValueError(f"Lodestone layer '{core}' is missing from one or more inputs.") + if mismatch_mode == "skip": + sources = [source for source in sources if source[0] == 0] + if not sources: + continue + dtypes = [] + for index, _block, kind, keys in sources: + stream_keys[index].extend(keys) + dtypes.append(handlers[index].get_dtype(keys[0] if kind == "diff" else keys[1])) + units.append({ + "core": core, + "output_key": output_key, + "sources": sources, + "missing": 0 if preserve or mismatch_mode != "zeros" else len(handlers) - len(present), + "dtypes": dtypes, + }) + return units, stream_keys + + +def _process_on_device(unit, loaded, method, save_dtype, device): + deltas = [] + for index, block, kind, keys in unit["sources"]: + if kind == "diff": + delta = _to_device(loaded[(index, keys[0])], device) + else: + down = loaded[(index, keys[0])] + up = loaded[(index, keys[1])] + alpha = loaded[(index, keys[2])] if len(keys) == 3 else None + down, up = normalize_lora_pair(down, up, alpha, layer=block) + delta = factor_delta(_to_device(down, device), _to_device(up, device), block) + deltas.append(delta) + if unit["missing"]: + deltas.extend(torch.zeros_like(deltas[0]) for _ in range(unit["missing"])) + merged = merge_lodestone_deltas(deltas, method) + dtype = select_output_dtype(unit["dtypes"], save_dtype, is_1d_diff=merged.ndim == 1) + return merged.to(dtype).cpu().contiguous() + + +def _process(unit, loaded, method, save_dtype, requested_device): + device = torch.device(requested_device) + try: + return _process_on_device(unit, loaded, method, save_dtype, device), False + except Exception as error: + if device.type != "cuda" or not _is_cuda_oom(error): + raise + _release_cuda_oom() + return _process_on_device(unit, loaded, method, save_dtype, torch.device("cpu")), True + + +def merge_lodestone_loras(lora_paths, method, device, save_dtype, output_filename, + mismatch_mode="skip", exclude_patterns="", discard_patterns="", + glob_patterns=False, verbose=True): + operation = "Lodestone LoRA Merge" + prepare_for_large_operation(sum(estimate_model_size(path) for path in lora_paths) * 2.5, torch.device(device)) + handlers = [MemoryEfficientSafeOpen(path, low_memory=True) for path in lora_paths] + cursors = [] + try: + infos = [] + for index, handler in enumerate(handlers): + low_bit = inspect_low_bit_input(handler, f"LoRA {index + 1}", operation) + pairs, passthrough = parse_lora_layers(handler.keys()) + infos.append({"pairs": pairs, "passthrough": passthrough, "low_bit": low_bit}) + layer_map = build_lora_layer_map(infos, None) + validate_canonical_blocks(layer_map, operation) + units, stream_keys = _build_units( + layer_map, infos, handlers, mismatch_mode, exclude_patterns, discard_patterns, glob_patterns + ) + cursors = [ + AsyncTensorCursor(handler, keys, pin_memory=torch.device(device).type == "cuda") + for handler, keys in zip(handlers, stream_keys) + ] + output_path, _ = canonical_model_artifact_path("loras", output_filename) + metadata = (handlers[0].metadata() or {}).copy() + metadata.update({"merge_method": f"lodestone:{method}", "alpha_normalized": "true", "output_representation": "full_direct_difference"}) + fallback_count = 0 + pbar = comfy.utils.ProgressBar(len(units)) + with atomic_uel_writer(output_path, metadata) as writer, torch.no_grad(): + for unit in tqdm(units, desc=f"Lodestone layers ({method})", unit="layers"): + loaded = {} + try: + for index, _block, _kind, keys in unit["sources"]: + for key in keys: + loaded[(index, key)] = cursors[index].take(key) + output, fallback = _process(unit, loaded, method, save_dtype, device) + writer.write_batch([(unit["output_key"], output)]) + fallback_count += int(fallback) + del output + finally: + for index, key in list(loaded): + del loaded[(index, key)] + cursors[index].release(key) + loaded.clear() + pbar.update(1) + for cursor in cursors: + cursor.finish() + if verbose: + print(f"[{operation}] CUDA OOM CPU fallbacks: {fallback_count}") + print(f"[{operation}] Saved {len(units)} full-delta layers to {output_path}") + return output_path + finally: + for cursor in cursors: + cursor.close() + for handler in handlers: + handler.__exit__(None, None, None) + cleanup_after_operation() + + +def _controls(default_filename): + return [ + io.Combo.Input("calc_mode", options=list(LODESTONE_METHODS), default="sum", tooltip="Dense delta-space Lodestone calculation."), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip", tooltip="Preserve LoRA 1, insert zero deltas, or abort for missing layers."), + io.String.Input("output_filename", default=default_filename, tooltip="Filename without extension under ComfyUI's LoRA directory."), + io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"], default="bf16", tooltip="Saved full-delta tensor dtype."), + io.Combo.Input("process_device", options=["cuda", "cpu"], default="cuda", tooltip="Dense processing device; CUDA OOM retries the layer on CPU."), + io.String.Input("exclude_patterns", default="", multiline=True, tooltip="Preserve matching logical layers from LoRA 1."), + io.String.Input("discard_patterns", default="", multiline=True, tooltip="Omit matching logical layers."), + io.Boolean.Input("glob_patterns", default=False, tooltip="Interpret filters as globs instead of regular expressions."), + ] + + +def _execute(kwargs, count): + paths = [folder_paths.get_full_path_or_raise("loras", kwargs[f"lora_{index}"]) for index in range(1, count + 1)] + dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[kwargs["save_dtype"]] + merge_lodestone_loras(paths, kwargs["calc_mode"], kwargs["process_device"], dtype, + kwargs["output_filename"], kwargs["mismatch_mode"], + kwargs["exclude_patterns"], kwargs["discard_patterns"], kwargs["glob_patterns"]) + return canonical_model_artifact_path("loras", kwargs["output_filename"])[1] + + +class LodestoneLoRATwoMerger(io.ComfyNode): + @classmethod + def define_schema(cls): + options = folder_paths.get_filename_list("loras") + return io.Schema(node_id="LodestoneLoRATwoMerger", display_name="Lodestone Merge LoRAs (2)", + category="ModelUtils/LoRA/Merge/Lodestone", + inputs=[io.Combo.Input("lora_1", options=options, tooltip="First equal-prior LoRA and anchor."), io.Combo.Input("lora_2", options=options, tooltip="Second equal-prior LoRA."), *_controls("lodestone_merged_2_lora")], + outputs=[io.AnyType.Output(display_name="output_filename")], is_output_node=True) + + @classmethod + def execute(cls, **kwargs): + return io.NodeOutput(_execute(kwargs, 2)) + + +class LodestoneLoRAThreeMerger(io.ComfyNode): + @classmethod + def define_schema(cls): + options = folder_paths.get_filename_list("loras") + return io.Schema(node_id="LodestoneLoRAThreeMerger", display_name="Lodestone Merge LoRAs (3)", + category="ModelUtils/LoRA/Merge/Lodestone", + inputs=[io.Combo.Input("lora_1", options=options, tooltip="First equal-prior LoRA and anchor."), io.Combo.Input("lora_2", options=options, tooltip="Second equal-prior LoRA."), io.Combo.Input("lora_3", options=options, tooltip="Third equal-prior LoRA."), *_controls("lodestone_merged_3_lora")], + outputs=[io.AnyType.Output(display_name="output_filename")], is_output_node=True) + + @classmethod + def execute(cls, **kwargs): + return io.NodeOutput(_execute(kwargs, 3)) + + +class LodestoneLoRAMultiMerger(io.ComfyNode): + @classmethod + def define_schema(cls): + options = folder_paths.get_filename_list("loras") + optional = ["None", *options] + inputs = [io.Combo.Input("lora_count", options=[str(index) for index in range(2, 9)], default="2", tooltip="Number of consecutive equal-prior LoRAs."), + io.Combo.Input("lora_1", options=options, tooltip="First equal-prior LoRA and anchor."), io.Combo.Input("lora_2", options=options, tooltip="Second equal-prior LoRA.")] + inputs.extend(io.Combo.Input(f"lora_{index}", options=optional, default="None", tooltip="Optional equal-prior LoRA used when included by LoRA Count.") for index in range(3, 9)) + inputs.extend(_controls("lodestone_merged_multi_lora")) + return io.Schema(node_id="LodestoneLoRAMultiMerger", display_name="Lodestone LoRA Multi-Merge", + category="ModelUtils/LoRA/Merge/Lodestone", inputs=inputs, + outputs=[io.AnyType.Output(display_name="output_filename")], is_output_node=True) + + @classmethod + def execute(cls, **kwargs): + count = int(kwargs["lora_count"]) + missing = [index for index in range(1, count + 1) if not kwargs[f"lora_{index}"] or kwargs[f"lora_{index}"] == "None"] + if missing: + raise ValueError(f"LoRA Count includes unselected input(s): {', '.join(map(str, missing))}.") + return io.NodeOutput(_execute(kwargs, count)) + + +LODESTONE_MERGER_NODES = [LodestoneLoRATwoMerger, LodestoneLoRAThreeMerger, LodestoneLoRAMultiMerger] diff --git a/tests/test_artifact_paths.py b/tests/test_artifact_paths.py index 9a1e957..d275a91 100644 --- a/tests/test_artifact_paths.py +++ b/tests/test_artifact_paths.py @@ -13,6 +13,7 @@ MIGRATED_SAVE_MODULES = ( "consensus_merger.py", "merger.py", "lora_merger.py", + "lodestone_merger.py", "lora_resize.py", "dtype_conversion.py", "lora_extract_svd.py", diff --git a/tests/test_groups.toml b/tests/test_groups.toml index 01f762e..b6e74db 100644 --- a/tests/test_groups.toml +++ b/tests/test_groups.toml @@ -9,7 +9,7 @@ python_tests = [ paths = [ "nodes/artifact_paths.py", "nodes/dora_extract_wd.py", "nodes/dora_learned_wd.py", "nodes/dtype_conversion.py", - "nodes/lora_extract_svd.py", "nodes/lora_merger.py", + "nodes/lodestone_merger.py", "nodes/lora_extract_svd.py", "nodes/lora_merger.py", "nodes/lora_rename.py", "nodes/lora_resize.py", "nodes/merger.py", "nodes/prunekeys.py", "nodes/renamekeys.py", "nodes/text_encoder_extract.py", ] @@ -36,12 +36,12 @@ python_tests = ["tests/test_adaptive_extraction_pipeline.py", "tests/test_knee_p [groups.lora_merge] paths = [ - "nodes/device_utils.py", "nodes/lora_alpha.py", "nodes/lora_merger.py", - "nodes/quantization_guard.py", "nodes/uel_io.py", + "nodes/device_utils.py", "nodes/lodestone_merger.py", "nodes/lora_alpha.py", + "nodes/lora_merger.py", "nodes/quantization_guard.py", "nodes/uel_io.py", ] python_tests = [ "tests/test_lora_diff_dtype_contract.py", "tests/test_lora_merger_async.py", - "tests/test_lora_quantization_guard.py", + "tests/test_lodestone_merger.py", "tests/test_lora_quantization_guard.py", ] [groups.lora_resize] diff --git a/tests/test_lodestone_merger.py b/tests/test_lodestone_merger.py new file mode 100644 index 0000000..bd4b47e --- /dev/null +++ b/tests/test_lodestone_merger.py @@ -0,0 +1,157 @@ +import importlib +import sys +import types +from pathlib import Path + +import pytest +import torch +from unifiedefficientloader import IncrementalSafetensorsWriter, MemoryEfficientSafeOpen + + +REPO_ROOT = Path(__file__).resolve().parents[1] +PACKAGE_NAME = "modelutils_lodestone_tests" + + +def _load_module(name): + if PACKAGE_NAME not in sys.modules: + package = types.ModuleType(PACKAGE_NAME) + package.__path__ = [str(REPO_ROOT)] + sys.modules[PACKAGE_NAME] = package + return importlib.import_module(f"{PACKAGE_NAME}.{name}") + + +@pytest.fixture(scope="module") +def lodestone(): + return _load_module("nodes.lodestone_merger") + + +def _write(path, tensors): + with IncrementalSafetensorsWriter(str(path), max_workers=1) as writer: + writer.write_batch(list(tensors.items())) + + +def _read(path): + result = {} + with MemoryEfficientSafeOpen(str(path), low_memory=True) as handler: + stream = handler.async_stream(handler.keys(), batch_size=1, prefetch_batches=1, pin_memory=False) + for batch in stream: + for key, tensor in batch: + result[key] = tensor.clone() + handler.mark_processed(key) + return result + + +@pytest.mark.parametrize( + ("method", "expected"), + [ + ("sum", torch.diag(torch.tensor([2.0, 4.0]))), + ("mean", torch.diag(torch.tensor([1.0, 2.0]))), + ("slotnorm", torch.eye(2)), + ("normmatch", torch.diag(torch.tensor([2.0, 4.0])) * (2.0 / torch.sqrt(torch.tensor(20.0)))), + ("slotnorm-normmatch", torch.eye(2) * torch.sqrt(torch.tensor(2.0))), + ], +) +def test_methods_match_supplied_dense_math(lodestone, method, expected): + deltas = [torch.diag(torch.tensor([2.0, 0.0])), torch.diag(torch.tensor([0.0, 4.0]))] + torch.testing.assert_close(lodestone.merge_lodestone_deltas(deltas, method), expected) + + +def test_convolutional_factor_expansion(lodestone): + down = torch.arange(18, dtype=torch.float32).reshape(2, 1, 3, 3) + up = torch.tensor([[[[2.0]], [[-1.0]]]]) + expected = (up.reshape(1, 2) @ down.reshape(2, -1)).reshape(1, 1, 3, 3) + torch.testing.assert_close(lodestone.factor_delta(down, up, "conv"), expected) + + +def test_zero_cancellation_remains_zero(lodestone): + delta = torch.tensor([[1.0, -2.0]]) + result = lodestone.merge_lodestone_deltas([delta, -delta], "normmatch") + assert torch.count_nonzero(result) == 0 + + +@pytest.mark.parametrize("count", [3, 8]) +def test_three_and_multi_input_dense_means(lodestone, count): + deltas = [torch.full((2, 2), float(index)) for index in range(1, count + 1)] + expected = torch.full((2, 2), (count + 1) / 2) + torch.testing.assert_close( + lodestone.merge_lodestone_deltas(deltas, "mean"), expected + ) + + +def test_layer_plan_enforces_missing_filters_and_low_bit(lodestone): + class Handler: + def get_dtype(self, _key): + return torch.float32 + + handlers = [Handler(), Handler()] + roles = {"down": "layer.down", "up": "layer.up"} + infos = [ + {"pairs": {"layer": roles}, "low_bit": set()}, + {"pairs": {}, "low_bit": set()}, + ] + layer_map = {"layer": [(0, "layer")]} + with pytest.raises(ValueError, match="missing"): + lodestone._build_units(layer_map, infos, handlers, "error", "", "", False) + units, _ = lodestone._build_units(layer_map, infos, handlers, "zeros", "", "", False) + assert units[0]["missing"] == 1 + excluded, _ = lodestone._build_units(layer_map, infos, handlers, "zeros", "*.diff", "", True) + assert excluded[0]["missing"] == 0 + discarded, _ = lodestone._build_units(layer_map, infos, handlers, "skip", "", "*.diff", True) + assert discarded == [] + + infos[0]["low_bit"] = {"layer.down"} + with pytest.raises(ValueError, match="low-bit"): + lodestone._build_units(layer_map, infos, handlers, "skip", "", "", False) + + +def test_streamed_merge_writes_only_alpha_normalized_full_delta(monkeypatch, tmp_path, lodestone): + first = tmp_path / "first.safetensors" + second = tmp_path / "second.safetensors" + _write(first, { + "diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 0.0]]), + "diffusion_model.layer.lora_B.weight": torch.tensor([[2.0], [0.0]]), + "diffusion_model.layer.alpha": torch.tensor(0.5), + }) + _write(second, { + "diffusion_model.layer.lora_A.weight": torch.tensor([[0.0, 1.0]]), + "diffusion_model.layer.lora_B.weight": torch.tensor([[0.0], [4.0]]), + }) + monkeypatch.setattr(lodestone.folder_paths, "models_dir", str(tmp_path)) + monkeypatch.setattr(lodestone, "prepare_for_large_operation", lambda *args: None) + monkeypatch.setattr(lodestone, "cleanup_after_operation", lambda: None) + + output = lodestone.merge_lodestone_loras( + [str(first), str(second)], "sum", "cpu", torch.float32, "dense", verbose=False + ) + tensors = _read(output) + + assert list(tensors) == ["diffusion_model.layer.diff"] + torch.testing.assert_close(tensors["diffusion_model.layer.diff"], torch.diag(torch.tensor([1.0, 4.0]))) + + +def test_dedicated_nodes_do_not_modify_factor_merger_schema(lodestone): + standard = _load_module("nodes.lora_merger") + assert [node.define_schema().node_id for node in lodestone.LODESTONE_MERGER_NODES] == [ + "LodestoneLoRATwoMerger", "LodestoneLoRAThreeMerger", "LodestoneLoRAMultiMerger" + ] + standard_inputs = {item.id: item for item in standard.LoRAMultiMerge.define_schema().inputs} + assert standard_inputs["merge_mode"].options == ["concatenate", "weighted_sum"] + assert "calc_mode" not in standard_inputs + + +def test_dense_cuda_oom_retries_layer_on_cpu(monkeypatch, lodestone): + unit = {"dtypes": [torch.float32], "missing": 0, "sources": []} + calls = [] + + def process(_unit, _loaded, _method, _dtype, device): + calls.append(device.type) + if device.type == "cuda": + raise torch.cuda.OutOfMemoryError("CUDA out of memory") + return torch.ones(1) + + monkeypatch.setattr(lodestone, "_process_on_device", process) + monkeypatch.setattr(lodestone, "_release_cuda_oom", lambda: None) + output, fallback = lodestone._process(unit, {}, "sum", torch.float32, "cuda") + assert fallback is True + assert calls == ["cuda", "cpu"] + assert torch.equal(output, torch.ones(1))