nodes/silveroxides-ComfyUI-ModelUtils
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(lora): add factorized Delta CWB output

Browse Source
This commit is contained in:
silveroxides
2026-09-04 23:20:46 +02:00
parent af32ebea45
commit ae7f742e08
3 changed files with 125 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/cwb_delta_lora_mergers.md
+5
View File
@@ -13,6 +13,11 @@ Missing-layer handling can preserve LoRA 1, include explicit zero contributors,
or abort. Excluded layers preserve LoRA 1 and discarded layers are omitted.
Outputs contain only canonical `.diff` tensors with alpha already normalized.
`factorized_lora` output factorizes each completed merged layer immediately and
writes canonical `lora_A` and `lora_B` tensors at the selected maximum rank. It
never writes or retains a complete full-difference model. One-dimensional
differences remain `.diff` tensors.
Processing streams one logical layer from each input at a time and writes each
completed layer incrementally. CUDA out-of-memory retries only the failed layer
on CPU. The destination is replaced only after the complete output is written.
nodes/cwb_delta_lora_merger.py
+68 -9
View File
@@ -147,6 +147,29 @@ def _factor_delta(down, up, layer):
)
def _factorize_delta(delta, rank):
if delta.ndim not in (2, 4):
return None
shape = delta.shape
matrix = delta if delta.ndim == 2 else delta.reshape(shape[0], -1)
target_rank = max(1, min(int(rank), min(matrix.shape)))
if target_rank == min(matrix.shape):
if matrix.shape[0] <= matrix.shape[1]:
up = torch.eye(matrix.shape[0], device=matrix.device, dtype=matrix.dtype)
down = matrix
else:
up = matrix
down = torch.eye(matrix.shape[1], device=matrix.device, dtype=matrix.dtype)
else:
up, singular_values, right = torch.svd_lowrank(matrix, q=target_rank, niter=2)
up = up * singular_values.unsqueeze(0)
down = right.t()
if delta.ndim == 4:
down = down.reshape(down.shape[0], *shape[1:])
up = up.unsqueeze(-1).unsqueeze(-1)
return down, up
def _inspect_inputs(handlers):
infos = []
for index, handler in enumerate(handlers):
@@ -281,7 +304,8 @@ def _to_processing_device(tensor, device):
def _process_layer_on_device(
unit, loaded, settings: CWBSettings, save_dtype, device, diagnostics
unit, loaded, settings: CWBSettings, save_dtype, device, diagnostics,
output_representation, output_rank,
):
deltas = []
try:
@@ -321,7 +345,15 @@ def _process_layer_on_device(
save_dtype,
is_1d_diff=merged.ndim == 1,
)
return merged.to(dtype=dtype, device="cpu").contiguous(), dtype
if output_representation == "factorized_lora":
factors = _factorize_delta(merged, output_rank)
if factors is not None:
down, up = factors
return [
(canonical_lora_key(unit.core, "down"), down.to(dtype=dtype, device="cpu").contiguous()),
(canonical_lora_key(unit.core, "up"), up.to(dtype=dtype, device="cpu").contiguous()),
], dtype
return [(unit.output_key, merged.to(dtype=dtype, device="cpu").contiguous())], dtype
finally:
deltas.clear()
@@ -338,12 +370,16 @@ def _release_cuda_oom():
torch.cuda.empty_cache()
def _process_layer(unit, loaded, settings, save_dtype, requested_device, diagnostics):
def _process_layer(
unit, loaded, settings, save_dtype, requested_device, diagnostics,
output_representation="full_difference", output_rank=384,
):
device = torch.device(requested_device)
attempt_diagnostics = CWBDiagnostics()
try:
output, dtype = _process_layer_on_device(
unit, loaded, settings, save_dtype, device, attempt_diagnostics
unit, loaded, settings, save_dtype, device, attempt_diagnostics,
output_representation, output_rank,
)
diagnostics.absorb_runtime(attempt_diagnostics)
return output, dtype, False
@@ -353,7 +389,8 @@ def _process_layer(unit, loaded, settings, save_dtype, requested_device, diagnos
_release_cuda_oom()
retry_diagnostics = CWBDiagnostics()
output, dtype = _process_layer_on_device(
unit, loaded, settings, save_dtype, torch.device("cpu"), retry_diagnostics
unit, loaded, settings, save_dtype, torch.device("cpu"), retry_diagnostics,
output_representation, output_rank,
)
diagnostics.absorb_runtime(retry_diagnostics)
return output, dtype, True
@@ -370,6 +407,8 @@ class DeltaCWBLoRAMergerLogic:
mismatch_mode,
output_filename,
save_dtype,
output_representation,
output_rank,
process_device,
exclude_patterns,
discard_patterns,
@@ -416,7 +455,10 @@ class DeltaCWBLoRAMergerLogic:
"merge_method": "delta_cwb",
"source_models": ",".join(source_names),
"alpha_normalized": "true",
"output_representation": "full_direct_difference",
"output_representation": (
"factorized_lora" if output_representation == "factorized_lora"
else "full_direct_difference"
),
"cwb_consensus_type": settings.consensus_type,
"cwb_alignment_method": settings.alignment_method,
"cwb_alignment_threshold": str(settings.alignment_threshold),
@@ -442,21 +484,23 @@ class DeltaCWBLoRAMergerLogic:
loaded[(source.source_index, key)] = cursors[
source.source_index
].take(key)
output, dtype, used_cpu = _process_layer(
output_items, dtype, used_cpu = _process_layer(
unit,
loaded,
settings,
save_dtype,
process_device,
diagnostics,
output_representation,
output_rank,
)
writer.write_batch([(unit.output_key, output)])
writer.write_batch(output_items)
report.merged_layers += int(not unit.preserve_anchor)
report.anchor_preserved_layers += int(unit.preserve_anchor)
report.zero_filled_contributors += unit.zero_contributors
report.cuda_cpu_fallbacks += int(used_cpu)
report.record_dtype(dtype)
del output
del output_items
finally:
for source_index, key in list(loaded):
del loaded[(source_index, key)]
@@ -496,6 +540,8 @@ class DeltaCWBLoRAMergerLogic:
kwargs["mismatch_mode"],
kwargs["output_filename"],
requested_dtype,
kwargs["output_representation"],
kwargs["output_rank"],
kwargs["process_device"],
kwargs["exclude_patterns"],
kwargs["discard_patterns"],
@@ -551,6 +597,19 @@ def _fixed_inputs(count, default_filename):
default="bf16",
tooltip="Requested saved precision. Participating FP32 inputs preserve FP32.",
),
io.Combo.Input(
"output_representation",
options=["full_difference", "factorized_lora"],
default="full_difference",
tooltip="Save exact full differences or factorize each merged layer directly into compact LoRA factors.",
),
io.Int.Input(
"output_rank",
default=384,
min=1,
max=16384,
tooltip="Maximum rank used when Output Representation is factorized_lora.",
),
io.Combo.Input(
"process_device",
options=["cuda", "cpu"],
tests/test_cwb_delta_lora_merger.py
+52 -6
View File
@@ -160,6 +160,8 @@ def test_streamed_merge_normalizes_alpha_once_and_supports_mixed_rank(
"skip",
"dense/mixed_rank",
torch.float32,
"full_difference",
384,
"cpu",
"",
"",
@@ -196,7 +198,8 @@ def test_existing_diff_and_1d_output_remain_fp32(monkeypatch, tmp_path, delta_cw
filename, _ = delta_cwb.DeltaCWBLoRAMergerLogic.execute(
[str(first), str(second)], [first.name, second.name], _settings(delta_cwb),
"balanced_mean", False, "skip", "one_dim", torch.float16, "cpu",
"balanced_mean", False, "skip", "one_dim", torch.float16,
"full_difference", 384, "cpu",
"", "", False, False,
)
output = _read(tmp_path / "loras" / filename)[key]
@@ -204,6 +207,42 @@ def test_existing_diff_and_1d_output_remain_fp32(monkeypatch, tmp_path, delta_cw
torch.testing.assert_close(output, torch.tensor([2.0, 4.0]))
def test_factorized_output_writes_compact_factors_without_dense_tensor(
monkeypatch, tmp_path, delta_cwb
):
first = tmp_path / "factor_first.safetensors"
second = tmp_path / "factor_second.safetensors"
tensors = {
"diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 2.0]]),
"diffusion_model.layer.lora_B.weight": torch.tensor([[3.0], [6.0]]),
}
_write(first, tensors)
_write(second, tensors)
monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None)
filename, _ = delta_cwb.DeltaCWBLoRAMergerLogic.execute(
[str(first), str(second)], [first.name, second.name], _settings(delta_cwb),
"balanced_mean", False, "skip", "factorized", torch.float32,
"factorized_lora", 1, "cpu", "", "", False, False,
)
output_path = tmp_path / "loras" / filename
output = _read(output_path)
assert set(output) == {
"diffusion_model.layer.lora_A.weight",
"diffusion_model.layer.lora_B.weight",
}
assert "diffusion_model.layer.diff" not in output
reconstructed = (
output["diffusion_model.layer.lora_B.weight"]
@ output["diffusion_model.layer.lora_A.weight"]
)
torch.testing.assert_close(reconstructed, torch.tensor([[3.0, 6.0], [6.0, 12.0]]))
with MemoryEfficientSafeOpen(str(output_path), low_memory=True) as handler:
assert handler.metadata()["output_representation"] == "factorized_lora"
def test_atomic_writer_preserves_destination_on_cwb_failure(
monkeypatch, tmp_path, delta_cwb
):
@@ -250,7 +289,8 @@ def test_atomic_writer_preserves_destination_on_cwb_failure(
with pytest.raises(RuntimeError, match="planned failure"):
delta_cwb.DeltaCWBLoRAMergerLogic.execute(
[str(first), str(second)], [first.name, second.name], _settings(delta_cwb),
"balanced_mean", False, "skip", "atomic", torch.float32, "cpu",
"balanced_mean", False, "skip", "atomic", torch.float32,
"full_difference", 384, "cpu",
"", "", False, False,
)
assert destination.read_bytes() == b"existing destination"
@@ -268,7 +308,10 @@ def test_complete_layer_cuda_oom_retries_on_cpu(monkeypatch, delta_cwb):
)
calls = []
def process(_unit, _loaded, _settings, _dtype, device, _diagnostics):
def process(
_unit, _loaded, _settings, _dtype, device, _diagnostics,
_output_representation, _output_rank,
):
calls.append(device.type)
if device.type == "cuda":
raise torch.cuda.OutOfMemoryError("CUDA out of memory")
@@ -305,7 +348,8 @@ def test_exact_node_schema_contract(delta_cwb):
fixed_ids = [value.id for value in schemas[0].inputs]
assert fixed_ids == [
"lora_1", "lora_2", "cwb_preset", "cwb_config", "mismatch_mode",
"output_filename", "save_dtype", "process_device", "exclude_patterns",
"output_filename", "save_dtype", "output_representation", "output_rank",
"process_device", "exclude_patterns",
"discard_patterns", "glob_patterns", "force_clear_cache",
]
multi_ids = [value.id for value in schemas[2].inputs]
@@ -313,11 +357,13 @@ def test_exact_node_schema_contract(delta_cwb):
assert schemas[0].inputs[2].default == "balanced_mean"
assert schemas[0].inputs[4].default == "skip"
assert schemas[0].inputs[6].default == "bf16"
assert schemas[0].inputs[7].default == "cuda"
assert schemas[0].inputs[7].default == "full_difference"
assert schemas[0].inputs[8].default == 384
assert schemas[0].inputs[9].default == "cuda"
assert schemas[0].inputs[-2].default is False
assert schemas[0].inputs[-1].default is True
assert "agreement" in schemas[0].inputs[2].tooltip
assert "Only a failed CUDA layer retries on CPU" in schemas[0].inputs[7].tooltip
assert "Only a failed CUDA layer retries on CPU" in schemas[0].inputs[9].tooltip
def test_invalid_regex_and_unselected_multi_input_fail(delta_cwb):