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

feat(cwb): preserve secondary-only layers

Browse Source 
- build generic and LoRA worklists from the union of inputs
- copy single-source layers exactly and merge available providers
- keep logical LoRA groups canonical, atomic, and collision-safe
- retain mismatch, dtype, quantization, streaming, and OOM contracts
This commit is contained in:
silveroxides
2026-08-02 20:26:41 +02:00
parent aca6147f35
commit 89f4def83e
3 changed files with 498 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/consensus_mergers.md
+26 -12
View File
@@ -10,11 +10,13 @@ similarity-derived weights before producing the output vector.
Two-input and three-input variants are provided for checkpoints, standalone
diffusion models, text encoders, LoRAs, and embeddings.
Checkpoint, diffusion-model, text-encoder, and LoRA outputs are anchored to
Model A. Model A supplies the output key set, metadata, and naming. Keys found
only in later inputs are ignored. Embedding mergers instead use the union of
all keys and use the longest compatible first dimension as the alignment
reference.
Model A supplies output metadata and anchors shared-layer shape, naming, and
mismatch behavior. The output key set is the union of all inputs. A valid
tensor found only in a secondary input is copied unchanged. When Model A lacks
a tensor that is present in multiple secondary inputs, those available sources
are CWB-merged without inserting a zero contribution for Model A. Embedding
mergers additionally use the longest compatible first dimension as their
alignment reference.
## Presets and custom controls
@@ -42,6 +44,11 @@ or receives zero weight, CWB falls back to equal weighting.
- Exclude patterns preserve the anchored source tensor.
- Discard patterns remove the matching output tensor or logical LoRA layer.
The missing-input modes above apply to Model-A-owned layers. Model A being
absent from a secondary-only layer is not itself a mismatch. A shape conflict
between multiple available secondary providers does use the selected mismatch
mode, with the earliest available input as its preservation anchor.
Patterns use regular expressions by default. Enable Glob Patterns to use glob
syntax instead.
@@ -55,19 +62,26 @@ Recognized diffusion-model outputs use the canonical
`diffusion_model.<layer>.lora_A.weight` / `.lora_B.weight` convention.
Secondary formats are matched by normalized logical layer name. Differing
ranks are zero-padded like the DARE/TIES mergers: A/down tensors on dimension 0
and B/up tensors on dimension 1. An existing Model A alpha key is updated to
the resulting maximum rank; no alpha key is invented when Model A has none.
and B/up tensors on dimension 1. An existing anchor-source alpha key is updated
to the resulting maximum rank; no alpha key is invented when the anchor has
none.
Similarity alignment operates on paired latent-rank components. A rows and B
columns share one mapping; fixed input/output feature axes are never reordered.
Input alpha scaling is absorbed into the B factor before blending, and global
scale is applied once to the resulting pair rather than once per factor.
Incomplete pairs and unrecognized Model A tensors are copied unchanged.
Incomplete pairs and unrecognized tensors are copied from the earliest
available source unchanged.
Companion-bearing groups (`lora_mid`, reshape, DoRA, or set-weight) are kept
atomically from Model A rather than partially transformed.
atomically from their earliest available source rather than partially
transformed. Secondary-only logical groups are included automatically: a group
from one source is copied unchanged, while compatible groups from multiple
secondary sources are merged using only those sources. Recognized roles are
written with the same canonical output convention as shared layers.
`include_1d_diffs` is disabled by default. Disabled 1D direct layers preserve
Model A. Enabled 1D direct layers participate as complete tensors and are
always saved in FP32.
`include_1d_diffs` is disabled by default. Disabled shared 1D direct layers
preserve Model A, while secondary-only layers preserve their earliest source.
Enabled 1D direct layers participate as complete tensors and are always saved
in FP32.
## Precision and quantized inputs
nodes/consensus_merger.py
+256 -79
View File
@@ -651,10 +651,22 @@ def _normalized_lora_core(block_name: str) -> str:
return core.replace(".", "_")
def _secondary_lora_map(pairs: dict[str, dict[str, str]]) -> dict[str, str]:
def _secondary_lora_map(
pairs: dict[str, dict[str, str]],
*,
input_index: int | None = None,
) -> dict[str, str]:
result = {}
for block_name in pairs:
result.setdefault(_normalized_lora_core(block_name), block_name)
core = _normalized_lora_core(block_name)
previous = result.get(core)
if previous is not None and previous != block_name:
label = f" in input {input_index + 1}" if input_index is not None else ""
raise ValueError(
f"[CWB LoRA Merge] LoRA keys '{previous}' and '{block_name}' "
f"normalize to the same logical layer '{core}'{label}."
)
result[core] = block_name
return result
@@ -736,10 +748,9 @@ class ConsensusMergerLogic:
embedding_union,
):
primary = handlers[0]
keys = set(primary.keys())
if embedding_union:
for handler in handlers[1:]:
keys.update(handler.keys())
keys = set()
for handler in handlers:
keys.update(handler.keys())
keys = sorted(keys)
output_path = cls._output_path(model_type, params["output_filename"])
requested_dtype = _requested_dtype(params["save_dtype"])
@@ -748,6 +759,8 @@ class ConsensusMergerLogic:
exclude = _compile_patterns(params["exclude_patterns"], glob_mode=glob_mode)
discard = _compile_patterns(params["discard_patterns"], glob_mode=glob_mode)
pbar = comfy.utils.ProgressBar(len(keys))
secondary_only_copied = 0
secondary_only_merged = 0
writer_context = _atomic_output_writer(output_path, primary.metadata())
writer = writer_context.__enter__()
@@ -759,11 +772,21 @@ class ConsensusMergerLogic:
pbar.update(1)
continue
source_indices = [i for i, handler in enumerate(handlers) if key in handler.keys()]
preserve_index = 0 if key in primary.keys() else source_indices[0]
primary_owned = key in primary.keys()
secondary_only = not primary_owned
preserve_index = 0 if primary_owned else source_indices[0]
guarded = any(key in low_bit_sets[i] for i in source_indices)
excluded = _matches_any_pattern(key, exclude, glob_mode=glob_mode)
if guarded or excluded:
write_preserved_tensor(writer, key, handlers[preserve_index])
if secondary_only:
secondary_only_copied += 1
pbar.update(1)
continue
if secondary_only and len(source_indices) == 1:
write_preserved_tensor(writer, key, handlers[preserve_index])
secondary_only_copied += 1
pbar.update(1)
continue
@@ -775,10 +798,12 @@ class ConsensusMergerLogic:
preserve_index + 1,
)
write_preserved_tensor(writer, key, handlers[preserve_index])
if secondary_only:
secondary_only_copied += 1
pbar.update(1)
continue
if not embedding_union and len(source_indices) != len(handlers):
if primary_owned and not embedding_union and len(source_indices) != len(handlers):
if mismatch_mode == "error":
raise ValueError(f"Tensor '{key}' is missing from a CWB input.")
if mismatch_mode == "skip":
@@ -795,13 +820,16 @@ class ConsensusMergerLogic:
)
reference_shape = raw[reference_source].shape
else:
reference_source = 0
reference_shape = raw[0].shape
reference_source = preserve_index
reference_shape = raw[preserve_index].shape
tensors = []
actual_dtypes = []
reference_index = 0
for source_index, handler in enumerate(handlers):
preserve_for_mismatch = False
candidate_indices = range(len(handlers)) if primary_owned else source_indices
for source_index in candidate_indices:
handler = handlers[source_index]
if source_index not in raw:
if mismatch_mode == "error":
raise ValueError(f"Tensor '{key}' is missing from input {source_index + 1}.")
@@ -821,6 +849,9 @@ class ConsensusMergerLogic:
if aligned is None:
if mismatch_mode == "error":
raise ValueError(f"Tensor shape mismatch for '{key}'.")
if mismatch_mode == "skip":
preserve_for_mismatch = True
break
if mismatch_mode == "zeros":
tensors.append(torch.zeros(
reference_shape,
@@ -832,6 +863,12 @@ class ConsensusMergerLogic:
tensors.append(aligned)
actual_dtypes.append(handler.get_dtype(key))
if preserve_for_mismatch:
write_preserved_tensor(writer, key, handlers[preserve_index])
if secondary_only:
secondary_only_copied += 1
pbar.update(1)
continue
if not tensors:
pbar.update(1)
continue
@@ -851,20 +888,38 @@ class ConsensusMergerLogic:
)
writer.write(key, merged)
del merged
if secondary_only:
secondary_only_merged += 1
pbar.update(1)
except BaseException as exc:
writer_context.__exit__(type(exc), exc, exc.__traceback__)
raise
else:
writer_context.__exit__(None, None, None)
if secondary_only_copied or secondary_only_merged:
logging.info(
"[CWB Merge] Secondary-only tensors: %d copied, %d merged.",
secondary_only_copied,
secondary_only_merged,
)
return os.path.basename(output_path)
@classmethod
def _merge_loras(cls, handlers, low_bit_sets, model_type, params, settings):
parsed = [parse_lora_layers(handler.keys()) for handler in handlers]
primary_pairs, primary_passthrough = parsed[0]
validate_canonical_blocks(primary_pairs, "CWB LoRA Merge")
secondary_maps = [_secondary_lora_map(pairs) for pairs, _ in parsed]
logical_maps = []
logical_cores = []
seen_cores = set()
passthrough_keys = set()
for input_index, (pairs, passthrough) in enumerate(parsed):
validate_canonical_blocks(pairs, f"CWB LoRA Merge input {input_index + 1}")
logical_map = _secondary_lora_map(pairs, input_index=input_index)
logical_maps.append(logical_map)
passthrough_keys.update(passthrough)
for core in logical_map:
if core not in seen_cores:
logical_cores.append(core)
seen_cores.add(core)
output_path = cls._output_path(model_type, params["output_filename"])
requested_dtype = _requested_dtype(params["save_dtype"])
mismatch_mode = params["mismatch_mode"]
@@ -874,15 +929,23 @@ class ConsensusMergerLogic:
discard = _compile_patterns(params["discard_patterns"], glob_mode=glob_mode)
written = set()
preserved_companion_groups = 0
pbar = comfy.utils.ProgressBar(len(primary_pairs) + len(primary_passthrough))
secondary_only_copied = 0
secondary_only_merged = 0
pbar = comfy.utils.ProgressBar(len(logical_cores) + len(passthrough_keys))
def preserve_keys(writer, keys: Iterable[str]):
def preserve_keys(writer, keys: Iterable[str], source_index: int = 0):
for key in keys:
if key not in written:
write_preserved_tensor(writer, key, handlers[0])
write_preserved_tensor(writer, key, handlers[source_index])
written.add(key)
def preserve_roles(writer, block_name: str, keys: dict[str, str], roles):
def preserve_roles(
writer,
block_name: str,
keys: dict[str, str],
roles,
source_index: int = 0,
):
recognized_sources = set()
for role, source_key in layer_tensor_keys(keys).items():
if role not in roles:
@@ -890,39 +953,63 @@ class ConsensusMergerLogic:
output_key = canonical_lora_key(block_name, role)
recognized_sources.add(source_key)
if output_key not in written:
write_preserved_tensor(writer, source_key, handlers[0], output_key)
write_preserved_tensor(
writer,
source_key,
handlers[source_index],
output_key,
)
written.add(output_key)
return recognized_sources
def preserve_layer(writer, block_name: str, keys: dict[str, str]):
def preserve_layer(
writer,
block_name: str,
keys: dict[str, str],
source_index: int = 0,
):
recognized_sources = preserve_roles(
writer, block_name, keys, layer_tensor_keys(keys)
writer,
block_name,
keys,
layer_tensor_keys(keys),
source_index,
)
preserve_keys(
writer,
(
key for key in handlers[0].keys()
key for key in handlers[source_index].keys()
if key.startswith(f"{block_name}.") and key not in recognized_sources
),
source_index,
)
writer_context = _atomic_output_writer(output_path, handlers[0].metadata())
writer = writer_context.__enter__()
try:
with torch.no_grad():
for primary_block, primary_keys in tqdm(
primary_pairs.items(), desc="CWB merging LoRA layers", unit="layers"
for core in tqdm(
logical_cores, desc="CWB merging LoRA layers", unit="layers"
):
_clear_previous_layer(params)
core = _normalized_lora_core(primary_block)
matches = []
for index, (pairs, _) in enumerate(parsed):
block = primary_block if index == 0 else secondary_maps[index].get(core)
for index, ((pairs, _), logical_map) in enumerate(zip(parsed, logical_maps)):
block = logical_map.get(core)
matches.append((index, block, pairs.get(block) if block else None))
anchor_index, anchor_block, anchor_keys = next(
(index, block, keys)
for index, block, keys in matches
if block is not None and keys is not None
)
primary_owned = matches[0][2] is not None
secondary_only = not primary_owned
layer_keys = [
key for key in handlers[0].keys()
if key in primary_keys.values() or key.startswith(f"{primary_block}.")
key
for index, block, keys in matches
if block is not None and keys is not None
for key in handlers[index].keys()
if key in keys.values() or key.startswith(f"{block}.")
]
if any(_matches_any_pattern(key, discard, glob_mode=glob_mode) for key in layer_keys):
pbar.update(1)
@@ -940,12 +1027,23 @@ class ConsensusMergerLogic:
for _, _, keys in matches
)
if companion_bearing:
preserve_layer(writer, primary_block, primary_keys)
preserve_layer(writer, anchor_block, anchor_keys, anchor_index)
preserved_companion_groups += 1
if secondary_only:
secondary_only_copied += 1
pbar.update(1)
continue
if guarded or excluded:
preserve_layer(writer, primary_block, primary_keys)
preserve_layer(writer, anchor_block, anchor_keys, anchor_index)
if secondary_only:
secondary_only_copied += 1
pbar.update(1)
continue
available_matches = [match for match in matches if match[2] is not None]
if secondary_only and len(available_matches) == 1:
preserve_layer(writer, anchor_block, anchor_keys, anchor_index)
secondary_only_copied += 1
pbar.update(1)
continue
@@ -958,37 +1056,53 @@ class ConsensusMergerLogic:
if name in {"down", "up", "alpha", "diff", "diff_b"}
)
if "down" in primary_keys and "up" in primary_keys:
group_merged = False
if "down" in anchor_keys and "up" in anchor_keys:
pair_sources = []
pair_failed = False
primary_down_template = handlers[0].get_tensor(primary_keys["down"])
primary_up_template = handlers[0].get_tensor(primary_keys["up"])
anchor_down_template = handlers[anchor_index].get_tensor(anchor_keys["down"])
anchor_up_template = handlers[anchor_index].get_tensor(anchor_keys["up"])
for index, _, keys in matches:
if not keys or "down" not in keys or "up" not in keys:
if not keys:
if secondary_only:
continue
if mismatch_mode == "error":
raise ValueError(f"LoRA pair '{primary_block}' is missing from input {index + 1}.")
raise ValueError(f"LoRA pair '{anchor_block}' is missing from input {index + 1}.")
if mismatch_mode == "skip":
pair_failed = True
break
pair_sources.append((
index,
torch.zeros_like(primary_down_template),
torch.zeros_like(primary_up_template),
torch.zeros_like(anchor_down_template),
torch.zeros_like(anchor_up_template),
1.0,
))
continue
if "down" not in keys or "up" not in keys:
if mismatch_mode == "error":
raise ValueError(f"LoRA pair '{anchor_block}' is incomplete in input {index + 1}.")
if mismatch_mode == "skip":
pair_failed = True
break
pair_sources.append((
index,
torch.zeros_like(anchor_down_template),
torch.zeros_like(anchor_up_template),
1.0,
))
continue
down = handlers[index].get_tensor(keys["down"])
up = handlers[index].get_tensor(keys["up"])
if down.ndim < 2 or up.ndim < 2 or down.shape[0] != up.shape[1]:
if mismatch_mode == "error" or index == 0:
raise ValueError(f"Invalid LoRA rank dimensions for '{primary_block}'.")
if mismatch_mode == "error" or index == anchor_index:
raise ValueError(f"Invalid LoRA rank dimensions for '{anchor_block}'.")
if mismatch_mode == "skip":
pair_failed = True
break
pair_sources.append((
index,
torch.zeros_like(primary_down_template),
torch.zeros_like(primary_up_template),
torch.zeros_like(anchor_down_template),
torch.zeros_like(anchor_up_template),
1.0,
))
continue
@@ -997,38 +1111,50 @@ class ConsensusMergerLogic:
alpha = handlers[index].get_tensor(keys["alpha"])
if alpha.numel() != 1:
raise ValueError(
f"LoRA alpha for '{primary_block}' must be scalar."
f"LoRA alpha for '{anchor_block}' must be scalar."
)
scale = float(alpha.reshape(-1)[0].item()) / down.shape[0]
pair_sources.append((index, down, up, scale))
if pair_failed:
preserve_roles(writer, primary_block, primary_keys, {"down", "up", "alpha"})
preserve_roles(
writer,
anchor_block,
anchor_keys,
{"down", "up", "alpha"},
anchor_index,
)
elif pair_sources:
primary_down = pair_sources[0][1]
primary_up = pair_sources[0][2]
anchor_down = pair_sources[0][1]
anchor_up = pair_sources[0][2]
compatible = []
for index, down, up, scale in pair_sources:
valid = (
tuple(down.shape[1:]) == tuple(primary_down.shape[1:])
and up.shape[0] == primary_up.shape[0]
and tuple(up.shape[2:]) == tuple(primary_up.shape[2:])
tuple(down.shape[1:]) == tuple(anchor_down.shape[1:])
and up.shape[0] == anchor_up.shape[0]
and tuple(up.shape[2:]) == tuple(anchor_up.shape[2:])
)
if not valid:
if mismatch_mode == "error":
raise ValueError(f"LoRA pair shape mismatch for '{primary_block}'.")
raise ValueError(f"LoRA pair shape mismatch for '{anchor_block}'.")
if mismatch_mode == "skip":
compatible = []
break
compatible.append((
index,
torch.zeros_like(primary_down),
torch.zeros_like(primary_up),
torch.zeros_like(anchor_down),
torch.zeros_like(anchor_up),
1.0,
))
continue
compatible.append((index, down, up, scale))
if not compatible:
preserve_roles(writer, primary_block, primary_keys, {"down", "up", "alpha"})
preserve_roles(
writer,
anchor_block,
anchor_keys,
{"down", "up", "alpha"},
anchor_index,
)
else:
target_dtype = select_output_dtype(
logical_dtypes,
@@ -1040,72 +1166,105 @@ class ConsensusMergerLogic:
settings,
params["process_device"],
target_dtype,
operation_label=primary_block,
operation_label=anchor_block,
)
writer.write(
canonical_lora_key(primary_block, "down"),
canonical_lora_key(anchor_block, "down"),
merged_down,
)
writer.write(
canonical_lora_key(primary_block, "up"),
canonical_lora_key(anchor_block, "up"),
merged_up,
)
del merged_down, merged_up
written.update({
canonical_lora_key(primary_block, "down"),
canonical_lora_key(primary_block, "up"),
canonical_lora_key(anchor_block, "down"),
canonical_lora_key(anchor_block, "up"),
})
if "alpha" in primary_keys:
if "alpha" in anchor_keys:
writer.write(
canonical_lora_key(primary_block, "alpha"),
canonical_lora_key(anchor_block, "alpha"),
torch.tensor(float(max_rank), dtype=target_dtype),
)
written.add(canonical_lora_key(primary_block, "alpha"))
written.add(canonical_lora_key(anchor_block, "alpha"))
group_merged = True
for kind in ("diff", "diff_b", "w_norm", "b_norm"):
if kind not in primary_keys:
if kind not in anchor_keys:
continue
primary_key = primary_keys[kind]
output_key = canonical_lora_key(primary_block, kind)
primary_tensor = handlers[0].get_tensor(primary_key)
if primary_tensor.ndim == 1 and not include_1d:
anchor_key = anchor_keys[kind]
output_key = canonical_lora_key(anchor_block, kind)
anchor_tensor = handlers[anchor_index].get_tensor(anchor_key)
if anchor_tensor.ndim == 1 and not include_1d:
if output_key not in written:
write_preserved_tensor(writer, primary_key, handlers[0], output_key)
write_preserved_tensor(
writer,
anchor_key,
handlers[anchor_index],
output_key,
)
written.add(output_key)
continue
direct = []
direct_dtypes = []
failed = False
for index, _, keys in matches:
if not keys or kind not in keys:
if not keys:
if secondary_only:
continue
if mismatch_mode == "error":
raise ValueError(f"Direct LoRA layer '{primary_block}.{kind}' is missing from input {index + 1}.")
raise ValueError(f"Direct LoRA layer '{anchor_block}.{kind}' is missing from input {index + 1}.")
if mismatch_mode == "skip":
failed = True
break
direct.append(torch.zeros_like(primary_tensor, dtype=torch.float32))
direct.append(torch.zeros_like(anchor_tensor, dtype=torch.float32))
continue
if kind not in keys:
if secondary_only:
continue
if mismatch_mode == "error":
raise ValueError(f"Direct LoRA layer '{anchor_block}.{kind}' is missing from input {index + 1}.")
if mismatch_mode == "skip":
failed = True
break
direct.append(torch.zeros_like(anchor_tensor, dtype=torch.float32))
continue
tensor = handlers[index].get_tensor(keys[kind])
if tensor.shape != primary_tensor.shape:
if tensor.shape != anchor_tensor.shape:
if mismatch_mode == "error":
raise ValueError(f"Direct LoRA shape mismatch for '{primary_block}.{kind}'.")
raise ValueError(f"Direct LoRA shape mismatch for '{anchor_block}.{kind}'.")
if mismatch_mode == "skip":
failed = True
break
direct.append(torch.zeros_like(primary_tensor, dtype=torch.float32))
direct.append(torch.zeros_like(anchor_tensor, dtype=torch.float32))
continue
direct.append(tensor)
direct_dtypes.append(handlers[index].get_dtype(keys[kind]))
if failed:
if output_key not in written:
write_preserved_tensor(writer, primary_key, handlers[0], output_key)
write_preserved_tensor(
writer,
anchor_key,
handlers[anchor_index],
output_key,
)
written.add(output_key)
continue
if secondary_only and len(direct) == 1:
if output_key not in written:
write_preserved_tensor(
writer,
anchor_key,
handlers[anchor_index],
output_key,
)
written.add(output_key)
continue
target_dtype = select_output_dtype(
direct_dtypes,
requested_dtype,
force=params["override_dtype"],
is_1d_diff=primary_tensor.ndim == 1,
is_1d_diff=anchor_tensor.ndim == 1,
)
merged = _merge_tensors_to_cpu(
direct,
@@ -1118,18 +1277,30 @@ class ConsensusMergerLogic:
writer.write(output_key, merged)
del merged
written.add(output_key)
group_merged = True
preserve_layer(writer, primary_block, primary_keys)
preserve_layer(writer, anchor_block, anchor_keys, anchor_index)
if secondary_only:
if group_merged:
secondary_only_merged += 1
else:
secondary_only_copied += 1
pbar.update(1)
for key in primary_passthrough:
for key in sorted(passthrough_keys):
_clear_previous_layer(params)
if key in written:
continue
if _matches_any_pattern(key, discard, glob_mode=glob_mode):
pbar.update(1)
continue
preserve_keys(writer, [key])
source_index = next(
index for index, (_, passthrough) in enumerate(parsed)
if key in passthrough
)
preserve_keys(writer, [key], source_index)
if source_index > 0:
secondary_only_copied += 1
pbar.update(1)
except BaseException as exc:
writer_context.__exit__(type(exc), exc, exc.__traceback__)
@@ -1138,9 +1309,15 @@ class ConsensusMergerLogic:
writer_context.__exit__(None, None, None)
if preserved_companion_groups:
logging.warning(
"[CWB LoRA Merge] Preserved %d companion-bearing Model-A group(s)",
"[CWB LoRA Merge] Preserved %d companion-bearing group(s)",
preserved_companion_groups,
)
if secondary_only_copied or secondary_only_merged:
logging.info(
"[CWB LoRA Merge] Secondary-only groups: %d copied, %d merged.",
secondary_only_copied,
secondary_only_merged,
)
return os.path.basename(output_path)
tests/test_consensus_merger.py
+216 -2
View File
@@ -337,7 +337,7 @@ def test_all_ten_schemas_and_lora_switch_contract(cwb):
assert "include_1d_diffs" not in ids
def test_model_a_anchored_streaming_dtype_and_nonfloat_preservation(
def test_streaming_dtype_nonfloat_and_secondary_preservation(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "a.safetensors"
@@ -359,11 +359,12 @@ def test_model_a_anchored_streaming_dtype_and_nonfloat_preservation(
["a", "b"], "diffusion_models", _params("anchored")
)
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
assert set(tensors) == {"shared", "only_a", "metadata_tensor"}
assert set(tensors) == {"shared", "only_a", "extra", "metadata_tensor"}
torch.testing.assert_close(tensors["shared"], torch.tensor([[2.0, 0.0], [0.0, 2.0]]))
assert tensors["shared"].dtype == torch.float32
torch.testing.assert_close(tensors["only_a"], torch.tensor([5.0, 6.0], dtype=torch.float16))
torch.testing.assert_close(tensors["metadata_tensor"], torch.tensor([1, 2]))
torch.testing.assert_close(tensors["extra"], torch.ones((2, 2)))
overridden = cwb.ConsensusMergerLogic.execute(
["a", "b"],
@@ -442,6 +443,97 @@ def test_three_input_streaming_merge(monkeypatch, tmp_path, cwb):
torch.testing.assert_close(tensor, torch.tensor([[3.0, 0.0]]))
def test_two_input_preserves_secondary_only_tensor_exactly(monkeypatch, tmp_path, cwb):
a = tmp_path / "secondary_a.safetensors"
b = tmp_path / "secondary_b.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({
"shared": torch.tensor([1.0]),
"secondary": torch.tensor([2.0, 3.0], dtype=torch.float32),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"diffusion_models",
_params(
"secondary_exact",
global_scale=2.0,
override_dtype=True,
save_dtype="fp16",
),
)
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["secondary"]
torch.testing.assert_close(tensor, torch.tensor([2.0, 3.0]))
assert tensor.dtype == torch.float32
def test_three_input_merges_available_secondary_sources(monkeypatch, tmp_path, cwb):
a = tmp_path / "union_a.safetensors"
b = tmp_path / "union_b.safetensors"
c = tmp_path / "union_c.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({
"shared": torch.tensor([1.0]),
"secondary_shared": torch.tensor([2.0]),
"b_only": torch.tensor([7.0]),
}, str(b))
save_file({
"shared": torch.tensor([1.0]),
"secondary_shared": torch.tensor([4.0]),
"c_only": torch.tensor([9.0]),
}, str(c))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b), "c": str(c)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"], "diffusion_models", _params("secondary_union")
)
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
torch.testing.assert_close(tensors["secondary_shared"], torch.tensor([3.0]))
torch.testing.assert_close(tensors["b_only"], torch.tensor([7.0]))
torch.testing.assert_close(tensors["c_only"], torch.tensor([9.0]))
def test_secondary_shape_conflict_skip_preserves_earliest_source(monkeypatch, tmp_path, cwb):
a = tmp_path / "shape_a.safetensors"
b = tmp_path / "shape_b.safetensors"
c = tmp_path / "shape_c.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({"secondary": torch.tensor([2.0, 3.0])}, str(b))
save_file({"secondary": torch.tensor([[8.0, 9.0]])}, str(c))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b), "c": str(c)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"],
"diffusion_models",
_params("secondary_shape_skip", mismatch_mode="skip"),
)
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["secondary"]
torch.testing.assert_close(tensor, torch.tensor([2.0, 3.0]))
def test_secondary_only_filters_and_low_bit_preservation(monkeypatch, tmp_path, cwb):
a = tmp_path / "filter_a.safetensors"
b = tmp_path / "filter_b.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({
"secondary.keep": torch.tensor([4, 5], dtype=torch.uint8),
"secondary.drop": torch.tensor([6.0]),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"diffusion_models",
_params("secondary_filters", discard_patterns=r"secondary\.drop"),
)
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
assert "secondary.drop" not in tensors
torch.testing.assert_close(
tensors["secondary.keep"], torch.tensor([4, 5], dtype=torch.uint8)
)
def test_lora_cross_format_companion_group_preserves_model_a_canonically(
monkeypatch, tmp_path, cwb
):
@@ -468,6 +560,8 @@ def test_lora_cross_format_companion_group_preserves_model_a_canonically(
"diffusion_model.foo.lora_A.weight",
"diffusion_model.foo.lora_B.weight",
"diffusion_model.foo.dora_scale",
"lora_unet_secondary.lora_A.weight",
"lora_unet_secondary.lora_B.weight",
}
assert tensors["diffusion_model.foo.lora_A.weight"].shape == (1, 2)
assert tensors["diffusion_model.foo.lora_B.weight"].shape == (2, 1)
@@ -475,6 +569,126 @@ def test_lora_cross_format_companion_group_preserves_model_a_canonically(
torch.testing.assert_close(
tensors["diffusion_model.foo.dora_scale"], torch.tensor([1.0, 1.0])
)
torch.testing.assert_close(
tensors["lora_unet_secondary.lora_A.weight"], torch.ones((1, 2))
)
torch.testing.assert_close(
tensors["lora_unet_secondary.lora_B.weight"], torch.ones((2, 1))
)
def test_lora_secondary_only_pair_is_preserved_without_scaling(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "lora_secondary_a.safetensors"
b = tmp_path / "lora_secondary_b.safetensors"
save_file({
"diffusion_model.base.diff": torch.tensor([[1.0]]),
}, str(a))
save_file({
"transformer.extra.lora_down.weight": torch.tensor([[2.0, 3.0]]),
"transformer.extra.lora_up.weight": torch.tensor([[4.0], [5.0]]),
"transformer.extra.alpha": torch.tensor(1.0),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"loras",
_params(
"lora_secondary_exact",
global_scale=2.0,
override_dtype=True,
save_dtype="fp16",
),
lora_mode=True,
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_A.weight"],
torch.tensor([[2.0, 3.0]]),
)
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_B.weight"],
torch.tensor([[4.0], [5.0]]),
)
assert tensors["diffusion_model.extra.lora_A.weight"].dtype == torch.float32
assert tensors["diffusion_model.extra.alpha"].item() == pytest.approx(1.0)
def test_lora_three_input_merges_b_and_c_group_without_a(monkeypatch, tmp_path, cwb):
a = tmp_path / "lora_union_a.safetensors"
b = tmp_path / "lora_union_b.safetensors"
c = tmp_path / "lora_union_c.safetensors"
save_file({"diffusion_model.base.diff": torch.tensor([[1.0]])}, str(a))
save_file({
"diffusion_model.extra.lora_A.weight": torch.full((1, 2), 1.0),
"diffusion_model.extra.lora_B.weight": torch.full((2, 1), 1.0),
}, str(b))
save_file({
"transformer.extra.lora_down.weight": torch.full((1, 2), 3.0),
"transformer.extra.lora_up.weight": torch.full((2, 1), 3.0),
}, str(c))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b), "c": str(c)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"], "loras", _params("lora_secondary_union"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_A.weight"], torch.full((1, 2), 2.0)
)
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_B.weight"], torch.full((2, 1), 2.0)
)
def test_lora_secondary_direct_and_companion_groups_are_atomic(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "lora_atomic_a.safetensors"
b = tmp_path / "lora_atomic_b.safetensors"
save_file({"diffusion_model.base.diff": torch.tensor([[1.0]])}, str(a))
save_file({
"diffusion_model.norm.diff": torch.tensor([2.0, 3.0], dtype=torch.bfloat16),
"diffusion_model.extra.lora_A.weight": torch.tensor([[4.0, 5.0]]),
"diffusion_model.extra.lora_B.weight": torch.tensor([[6.0], [7.0]]),
"diffusion_model.extra.dora_scale": torch.tensor([8.0, 9.0]),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"loras",
_params("lora_secondary_atomic", global_scale=3.0),
lora_mode=True,
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.norm.diff"],
torch.tensor([2.0, 3.0], dtype=torch.bfloat16),
)
assert tensors["diffusion_model.norm.diff"].dtype == torch.bfloat16
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_A.weight"], torch.tensor([[4.0, 5.0]])
)
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_B.weight"], torch.tensor([[6.0], [7.0]])
)
torch.testing.assert_close(
tensors["diffusion_model.extra.dora_scale"], torch.tensor([8.0, 9.0])
)
def test_lora_rejects_duplicate_normalized_cores_in_secondary_input(cwb):
with pytest.raises(ValueError, match="same logical layer 'foo'.*input 2"):
cwb._secondary_lora_map(
{
"diffusion_model.foo": {},
"transformer.foo": {},
},
input_index=1,
)
def test_lora_peft_prefix_matching_and_existing_alpha(monkeypatch, tmp_path, cwb):