1805 lines
67 KiB
Python
1805 lines
67 KiB
Python
import importlib
|
|
import json
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
from safetensors.torch import load_file, save_file
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
PACKAGE_NAME = "modelutils_cwb_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 cwb():
|
|
return _load_module("nodes.consensus_merger")
|
|
|
|
|
|
def _patch_io(monkeypatch, module, tmp_path, paths):
|
|
monkeypatch.setattr(module.folder_paths, "get_full_path", lambda _, name: paths.get(name))
|
|
monkeypatch.setattr(module.folder_paths, "get_folder_paths", lambda _: [str(tmp_path)])
|
|
monkeypatch.setattr(module.folder_paths, "models_dir", str(tmp_path))
|
|
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
|
|
|
|
|
|
def _result_path(tmp_path, category, result):
|
|
return tmp_path / category / result
|
|
|
|
|
|
def _params(output_filename, **overrides):
|
|
settings = {
|
|
"consensus_type": "mean",
|
|
"alignment_method": "index",
|
|
"alignment_threshold": 0.4,
|
|
"similarity_threshold": 0.0,
|
|
"power_alpha": 2.0,
|
|
"diversity_beta": 0.0,
|
|
"rescale_norm": False,
|
|
"global_scale": 1.0,
|
|
"dynamic_similarity_contrast": False,
|
|
"soft_comfort_bandpass": False,
|
|
"position_weight": 0.0,
|
|
"preserve_common_prefix": False,
|
|
}
|
|
for key in list(settings):
|
|
if key in overrides:
|
|
settings[key] = overrides.pop(key)
|
|
params = {
|
|
"cwb_preset": "broad_sim_medn_rn_softcb",
|
|
"cwb_config": _load_module("nodes.consensus_merger").build_custom_cwb_settings(
|
|
**settings
|
|
),
|
|
"mismatch_mode": "skip",
|
|
"output_filename": output_filename,
|
|
"save_dtype": "fp16",
|
|
"process_device": "cpu",
|
|
"exclude_patterns": "",
|
|
"discard_patterns": "",
|
|
"glob_patterns": False,
|
|
"lazy_load": True,
|
|
"force_clear_cache": False,
|
|
"override_dtype": False,
|
|
"include_1d_diffs": False,
|
|
}
|
|
params.update(overrides)
|
|
return params
|
|
|
|
|
|
def _settings(cwb, **overrides):
|
|
values = {
|
|
"consensus_type": "mean",
|
|
"alignment_method": "index",
|
|
"alignment_threshold": 0.4,
|
|
"similarity_threshold": 0.0,
|
|
"power_alpha": 2.0,
|
|
"diversity_beta": 0.0,
|
|
"rescale_norm": False,
|
|
"global_scale": 1.0,
|
|
"dynamic_similarity_contrast": False,
|
|
"soft_comfort_bandpass": False,
|
|
"position_weight": 0.0,
|
|
"preserve_common_prefix": False,
|
|
}
|
|
values.update(overrides)
|
|
return cwb.build_custom_cwb_settings(**values)
|
|
|
|
|
|
def test_preset_resolution_matches_extended_contract(cwb):
|
|
default = cwb.resolve_cwb_settings(
|
|
"broad_sim_medn_rn_softcb", cwb.LORA_CWB_PRESETS
|
|
)
|
|
assert default.consensus_type == "median"
|
|
assert default.alignment_method == "similarity"
|
|
assert default.alignment_threshold == pytest.approx(0.0)
|
|
assert default.similarity_threshold == pytest.approx(0.0)
|
|
assert default.power_alpha == pytest.approx(2.0)
|
|
assert default.diversity_beta == pytest.approx(4.0)
|
|
assert default.rescale_norm is True
|
|
assert default.dynamic_similarity_contrast is False
|
|
assert default.soft_comfort_bandpass is True
|
|
assert default.position_weight == pytest.approx(0.05)
|
|
|
|
with pytest.raises(ValueError, match="Unsupported CWB preset"):
|
|
cwb.resolve_cwb_settings("baseline", cwb.LORA_CWB_PRESETS)
|
|
|
|
|
|
def test_preset_construction_rejects_omitted_fields(cwb):
|
|
with pytest.raises(TypeError):
|
|
cwb._preset(
|
|
consensus_type="median",
|
|
alignment_method="similarity",
|
|
alignment_threshold=0.0,
|
|
)
|
|
|
|
|
|
def test_lora_alignment_presets_use_sweep_verified_thresholds(cwb):
|
|
assert cwb.LORA_CWB_PRESETS["broad_sim_medn_rn_softcb"][
|
|
"alignment_threshold"
|
|
] == pytest.approx(0.0)
|
|
assert cwb.LORA_CWB_PRESETS["moderate_sim_medn_rn_softcb"][
|
|
"alignment_threshold"
|
|
] == pytest.approx(0.0005)
|
|
assert cwb.LORA_CWB_PRESETS["conservative_sim_medn_rn_softcb"][
|
|
"alignment_threshold"
|
|
] == pytest.approx(0.0025)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("registry_name", "preset", "alignment_threshold", "similarity_threshold"),
|
|
[
|
|
(registry_name, preset, alignment_threshold, similarity_threshold)
|
|
for registry_name in ("EMBEDDING_CWB_PRESETS", "LORA_CWB_PRESETS")
|
|
for preset, alignment_threshold, similarity_threshold in (
|
|
("focused_strong_sim_medn", 0.85, 0.60),
|
|
("focused_balance_sim_medn", 0.75, 0.55),
|
|
("focused_soft_sim_medn", 0.55, 0.50),
|
|
("focused_weak_sim_medn", 0.25, 0.35),
|
|
)
|
|
],
|
|
)
|
|
def test_focused_median_similarity_presets(
|
|
cwb, registry_name, preset, alignment_threshold, similarity_threshold
|
|
):
|
|
settings = getattr(cwb, registry_name)[preset]
|
|
assert settings == {
|
|
"consensus_type": "median",
|
|
"alignment_method": "similarity",
|
|
"alignment_threshold": alignment_threshold,
|
|
"similarity_threshold": similarity_threshold,
|
|
"power_alpha": 1.25,
|
|
"diversity_beta": 0.0,
|
|
"rescale_norm": False,
|
|
"global_scale": 1.0,
|
|
"dynamic_similarity_contrast": False,
|
|
"soft_comfort_bandpass": False,
|
|
"position_weight": 0.20,
|
|
"preserve_common_prefix": False,
|
|
}
|
|
|
|
|
|
def test_cwb_math_index_prefix_and_scale(cwb):
|
|
settings = _settings(
|
|
cwb,
|
|
consensus_type="mean",
|
|
alignment_method="index",
|
|
global_scale=1.0,
|
|
)
|
|
first = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
|
|
second = torch.tensor([[3.0, 0.0], [0.0, 3.0]])
|
|
torch.testing.assert_close(
|
|
cwb.merge_cwb_tensors([first, second], settings),
|
|
torch.tensor([[2.0, 0.0], [0.0, 2.0]]),
|
|
)
|
|
|
|
prefix_settings = _settings(
|
|
cwb,
|
|
consensus_type="mean",
|
|
alignment_method="index",
|
|
global_scale=2.0,
|
|
preserve_common_prefix=True,
|
|
)
|
|
prefixed = cwb.merge_cwb_tensors(
|
|
[torch.tensor([[9.0, 9.0], [1.0, 0.0]]),
|
|
torch.tensor([[9.0, 9.0], [3.0, 0.0]])],
|
|
prefix_settings,
|
|
)
|
|
torch.testing.assert_close(prefixed[0], torch.tensor([9.0, 9.0]))
|
|
torch.testing.assert_close(prefixed[1], torch.tensor([4.0, 0.0]))
|
|
|
|
|
|
def test_similarity_alignment_reorders_matching_rows(cwb):
|
|
settings = _settings(
|
|
cwb,
|
|
consensus_type="mean",
|
|
alignment_method="similarity",
|
|
alignment_threshold=0.5,
|
|
)
|
|
reference = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
|
|
reversed_source = torch.tensor([[0.0, 2.0], [2.0, 0.0]])
|
|
torch.testing.assert_close(
|
|
cwb.merge_cwb_tensors([reference, reversed_source], settings),
|
|
torch.tensor([[1.5, 0.0], [0.0, 1.5]]),
|
|
)
|
|
|
|
scores = torch.ones((2, 2))
|
|
biased = cwb._position_biased_scores(scores, 1.0)
|
|
assert biased[0, 0] > biased[0, 1]
|
|
|
|
|
|
def test_fixed_coordinate_merge_never_reorders_model_rows(cwb):
|
|
settings = _settings(
|
|
cwb,
|
|
consensus_type="mean",
|
|
alignment_method="similarity",
|
|
alignment_threshold=0.5,
|
|
)
|
|
reference = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
|
|
reversed_source = torch.tensor([[0.0, 2.0], [2.0, 0.0]])
|
|
expected = torch.stack([
|
|
cwb.merge_consensus_group(
|
|
torch.stack([reference[row], reversed_source[row]]), settings
|
|
)
|
|
for row in range(reference.shape[0])
|
|
])
|
|
torch.testing.assert_close(
|
|
cwb.merge_cwb_tensors(
|
|
[reference, reversed_source],
|
|
settings,
|
|
allow_similarity_alignment=False,
|
|
),
|
|
expected,
|
|
)
|
|
|
|
|
|
def test_lora_similarity_alignment_pairs_a_rows_with_b_columns(cwb):
|
|
settings = _settings(
|
|
cwb,
|
|
consensus_type="mean",
|
|
alignment_method="similarity",
|
|
alignment_threshold=0.5,
|
|
)
|
|
reference_down = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
|
|
reference_up = torch.tensor([[2.0, 0.0], [0.0, 3.0]])
|
|
source_down = torch.tensor([[0.0, 1.0], [-1.0, 0.0]])
|
|
source_up = torch.tensor([[0.0, -2.0], [3.0, 0.0]])
|
|
|
|
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
|
|
[reference_down, source_down],
|
|
[reference_up, source_up],
|
|
settings,
|
|
reference_index=0,
|
|
)
|
|
torch.testing.assert_close(merged_down, reference_down)
|
|
torch.testing.assert_close(merged_up, reference_up)
|
|
|
|
|
|
def test_lora_krea_shape_scores_only_rank_components(monkeypatch, cwb):
|
|
rank = 256
|
|
output_features = 36864
|
|
down = torch.ones((rank, 2))
|
|
up = torch.ones((output_features, rank))
|
|
matrix_shapes = []
|
|
|
|
def bounded_mm(left, right):
|
|
matrix_shapes.append((left.shape, right.shape))
|
|
return torch.eye(left.shape[0], right.shape[1], dtype=left.dtype)
|
|
|
|
monkeypatch.setattr(cwb.torch, "mm", bounded_mm)
|
|
monkeypatch.setattr(
|
|
cwb,
|
|
"merge_consensus_group",
|
|
lambda stacked, settings, **kwargs: stacked[0].clone(),
|
|
)
|
|
settings = _settings(
|
|
cwb,
|
|
alignment_method="similarity",
|
|
alignment_threshold=0.5,
|
|
)
|
|
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
|
|
[down, down], [up, up], settings, reference_index=0
|
|
)
|
|
|
|
assert merged_down.shape == down.shape
|
|
assert merged_up.shape == up.shape
|
|
assert matrix_shapes
|
|
assert all(left[0] == rank and right[1] == rank for left, right in matrix_shapes)
|
|
|
|
|
|
def test_mixed_rank_similarity_searches_full_reference_without_padding(
|
|
monkeypatch, cwb
|
|
):
|
|
reference_down = torch.eye(3)
|
|
reference_up = torch.eye(3)
|
|
source_down = torch.tensor([[0.0, 0.0, 2.0]])
|
|
source_up = torch.tensor([[0.0], [0.0], [4.0]])
|
|
settings = _settings(
|
|
cwb,
|
|
consensus_type="mean",
|
|
alignment_method="similarity",
|
|
alignment_threshold=0.5,
|
|
rescale_norm=True,
|
|
)
|
|
diagnostics = cwb.CWBDiagnostics()
|
|
matrix_shapes = []
|
|
original_mm = cwb.torch.mm
|
|
|
|
def record_mm(left, right):
|
|
matrix_shapes.append((tuple(left.shape), tuple(right.shape)))
|
|
return original_mm(left, right)
|
|
|
|
monkeypatch.setattr(cwb.torch, "mm", record_mm)
|
|
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
|
|
[reference_down, source_down],
|
|
[reference_up, source_up],
|
|
settings,
|
|
reference_index=0,
|
|
diagnostics=diagnostics,
|
|
)
|
|
|
|
assert matrix_shapes == [((3, 3), (3, 1)), ((3, 3), (3, 1))]
|
|
assert diagnostics.alignment_matches == 1
|
|
assert diagnostics.anchor_only_groups == 2
|
|
torch.testing.assert_close(merged_down[:2], reference_down[:2])
|
|
torch.testing.assert_close(merged_up[:, :2], reference_up[:, :2])
|
|
torch.testing.assert_close(merged_down[2], torch.tensor([0.0, 0.0, 1.5]))
|
|
torch.testing.assert_close(merged_up[:, 2], torch.tensor([0.0, 0.0, 2.5]))
|
|
|
|
|
|
def test_mixed_rank_index_does_not_rescale_absent_components(cwb):
|
|
reference_down = torch.eye(3)
|
|
reference_up = torch.eye(3)
|
|
source_down = torch.tensor([[2.0, 0.0, 0.0]])
|
|
source_up = torch.tensor([[2.0], [0.0], [0.0]])
|
|
settings = _settings(
|
|
cwb,
|
|
consensus_type="mean",
|
|
alignment_method="index",
|
|
rescale_norm=True,
|
|
)
|
|
|
|
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
|
|
[reference_down, source_down],
|
|
[reference_up, source_up],
|
|
settings,
|
|
reference_index=0,
|
|
)
|
|
|
|
torch.testing.assert_close(merged_down[1:], reference_down[1:])
|
|
torch.testing.assert_close(merged_up[:, 1:], reference_up[:, 1:])
|
|
assert torch.linalg.vector_norm(merged_down[0]).item() == pytest.approx(1.5)
|
|
assert torch.linalg.vector_norm(merged_up[:, 0]).item() == pytest.approx(1.5)
|
|
|
|
|
|
def test_genuine_zero_component_is_not_treated_as_structural_padding(cwb):
|
|
reference_down = torch.eye(2)
|
|
reference_up = torch.eye(2)
|
|
source_down = torch.zeros((1, 2))
|
|
source_up = torch.zeros((2, 1))
|
|
settings = _settings(
|
|
cwb,
|
|
alignment_method="similarity",
|
|
alignment_threshold=0.0,
|
|
)
|
|
diagnostics = cwb.CWBDiagnostics()
|
|
|
|
cwb.merge_cwb_lora_pairs(
|
|
[reference_down, source_down],
|
|
[reference_up, source_up],
|
|
settings,
|
|
reference_index=0,
|
|
diagnostics=diagnostics,
|
|
)
|
|
|
|
assert diagnostics.alignment_matches == 1
|
|
assert diagnostics.anchor_only_groups == 1
|
|
|
|
|
|
def test_explicit_mismatch_zero_applies_to_every_reference_component(cwb):
|
|
down = torch.eye(2)
|
|
up = torch.eye(2)
|
|
settings = _settings(
|
|
cwb,
|
|
consensus_type="mean",
|
|
alignment_method="index",
|
|
rescale_norm=True,
|
|
)
|
|
diagnostics = cwb.CWBDiagnostics()
|
|
|
|
merged_down, merged_up, rank = cwb._merge_lora_pair_to_cpu(
|
|
[cwb.LoRAPairSource(0, down, up), cwb.LoRAPairSource.zero(1)],
|
|
settings,
|
|
"cpu",
|
|
torch.float32,
|
|
diagnostics=diagnostics,
|
|
)
|
|
|
|
assert rank == 2
|
|
torch.testing.assert_close(merged_down, down * 0.5)
|
|
torch.testing.assert_close(merged_up, up * 0.5)
|
|
assert diagnostics.explicit_zero_contributors == 1
|
|
assert diagnostics.structural_rank_slots_excluded == 0
|
|
|
|
|
|
def test_mixed_rank_diagnostics_count_only_genuine_components(cwb):
|
|
settings = _settings(cwb, alignment_method="index")
|
|
diagnostics = cwb.CWBDiagnostics()
|
|
sources = [
|
|
cwb.LoRAPairSource(0, torch.ones((4, 2)), torch.ones((2, 4))),
|
|
cwb.LoRAPairSource(1, torch.ones((2, 2)), torch.ones((2, 2))),
|
|
cwb.LoRAPairSource(2, torch.ones((3, 2)), torch.ones((2, 3))),
|
|
]
|
|
|
|
_, _, rank = cwb._merge_lora_pair_to_cpu(
|
|
sources,
|
|
settings,
|
|
"cpu",
|
|
torch.float32,
|
|
diagnostics=diagnostics,
|
|
)
|
|
|
|
assert rank == 4
|
|
assert diagnostics.lora_groups == 1
|
|
assert diagnostics.mixed_rank_lora_groups == 1
|
|
assert diagnostics.genuine_lora_components == 9
|
|
assert diagnostics.structural_rank_slots_excluded == 3
|
|
report = diagnostics.render("custom", 3, True)
|
|
assert "Mixed-rank LoRA groups: 1" in report
|
|
assert "Genuine LoRA components: 9" in report
|
|
assert "Structural rank slots excluded: 3" in report
|
|
assert "CWB LORA DIMENSION GROUPS" in report
|
|
assert "Input ranks: (4, 2, 3)" in report
|
|
assert "Matched source components: 5/5 (100.00%)" in report
|
|
|
|
|
|
def test_lora_report_groups_dimensions_and_summarizes_rectangular_similarity(cwb):
|
|
settings = _settings(
|
|
cwb,
|
|
consensus_type="mean",
|
|
alignment_method="similarity",
|
|
alignment_threshold=0.5,
|
|
)
|
|
diagnostics = cwb.CWBDiagnostics()
|
|
sources = [
|
|
cwb.LoRAPairSource(0, torch.eye(3), torch.eye(3)),
|
|
cwb.LoRAPairSource(
|
|
1,
|
|
torch.tensor([[0.0, 0.0, 2.0]]),
|
|
torch.tensor([[0.0], [0.0], [2.0]]),
|
|
),
|
|
]
|
|
|
|
for layer in ("blocks.1.attn", "blocks.2.attn"):
|
|
cwb._merge_lora_pair_to_cpu(
|
|
sources,
|
|
settings,
|
|
"cpu",
|
|
torch.float32,
|
|
operation_label=layer,
|
|
diagnostics=diagnostics,
|
|
)
|
|
|
|
report = diagnostics.render("custom", 2, True)
|
|
assert "Group 1: 2 layer(s)" in report
|
|
assert "Alignment: similarity; matrices: 3x1" in report
|
|
assert "Matched source components: 2/2 (100.00%)" in report
|
|
assert "Reference-only components: 4/6" in report
|
|
assert "Structural rank slots excluded: 4" in report
|
|
assert "Candidate similarity count: 6" in report
|
|
assert "Matched similarity p05/median/p95: 1 / 1 / 1" in report
|
|
assert "blocks.1.attn" in report
|
|
assert "blocks.2.attn" in report
|
|
assert all(
|
|
not isinstance(value, torch.Tensor)
|
|
for layer_report in diagnostics.lora_layer_reports
|
|
for value in vars(layer_report).values()
|
|
)
|
|
|
|
|
|
def test_lora_alpha_and_global_scale_apply_once_to_pair(cwb):
|
|
settings = _settings(
|
|
cwb,
|
|
consensus_type="mean",
|
|
alignment_method="index",
|
|
global_scale=3.0,
|
|
)
|
|
down = torch.ones((1, 1))
|
|
up = torch.ones((1, 1))
|
|
merged_down, merged_up, rank = cwb._merge_lora_pair_to_cpu(
|
|
[
|
|
cwb.LoRAPairSource(0, down, up, 2.0),
|
|
cwb.LoRAPairSource(1, down, up, 2.0),
|
|
],
|
|
settings,
|
|
"cpu",
|
|
torch.float32,
|
|
)
|
|
assert rank == 1
|
|
torch.testing.assert_close(merged_down, torch.ones((1, 1)))
|
|
torch.testing.assert_close(merged_up, torch.full((1, 1), 6.0))
|
|
|
|
|
|
def test_lora_pair_alignment_supports_convolution_factors(cwb):
|
|
settings = _settings(
|
|
cwb,
|
|
consensus_type="mean",
|
|
alignment_method="index",
|
|
)
|
|
down = torch.arange(6, dtype=torch.float32).reshape(2, 3, 1, 1)
|
|
up = torch.arange(8, dtype=torch.float32).reshape(4, 2, 1, 1)
|
|
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
|
|
[down, down], [up, up], settings, reference_index=0
|
|
)
|
|
torch.testing.assert_close(merged_down, down)
|
|
torch.testing.assert_close(merged_up, up)
|
|
|
|
|
|
def test_lora_cuda_oom_retries_current_pair_on_cpu(monkeypatch, caplog, cwb):
|
|
settings = _settings(
|
|
cwb,
|
|
consensus_type="mean",
|
|
alignment_method="index",
|
|
)
|
|
original_to_compute = cwb._to_compute
|
|
attempted_devices = []
|
|
|
|
def fail_cuda(tensor, device):
|
|
attempted_devices.append(device)
|
|
if str(device).startswith("cuda"):
|
|
raise torch.OutOfMemoryError("injected allocation failure")
|
|
return original_to_compute(tensor, device)
|
|
|
|
monkeypatch.setattr(cwb, "_to_compute", fail_cuda)
|
|
monkeypatch.setattr(cwb, "_release_failed_cuda_operation", lambda: None)
|
|
reference_down = torch.eye(2)
|
|
reference_up = torch.eye(2)
|
|
source_down = torch.ones((1, 2))
|
|
source_up = torch.ones((2, 1))
|
|
diagnostics = cwb.CWBDiagnostics()
|
|
with caplog.at_level("WARNING"):
|
|
merged_down, merged_up, rank = cwb._merge_lora_pair_to_cpu(
|
|
[
|
|
cwb.LoRAPairSource(0, reference_down, reference_up),
|
|
cwb.LoRAPairSource(1, source_down, source_up),
|
|
],
|
|
settings,
|
|
"cuda",
|
|
torch.float32,
|
|
operation_label="diffusion_model.foo",
|
|
diagnostics=diagnostics,
|
|
)
|
|
|
|
assert rank == 2
|
|
assert merged_down.shape == reference_down.shape
|
|
assert merged_up.shape == reference_up.shape
|
|
assert attempted_devices[0] == "cuda"
|
|
assert "cpu" in attempted_devices
|
|
assert "retrying this layer on CPU" in caplog.text
|
|
assert "diffusion_model.foo" in caplog.text
|
|
assert diagnostics.cpu_fallbacks == 1
|
|
assert diagnostics.lora_groups == 1
|
|
assert diagnostics.structural_rank_slots_excluded == 1
|
|
|
|
|
|
def test_mixed_rank_common_prefix_is_limited_to_smallest_rank(cwb):
|
|
reference_down = torch.eye(3)
|
|
reference_up = torch.eye(3)
|
|
source_down = torch.tensor([[1.0, 0.0, 0.0], [0.0, 2.0, 0.0]])
|
|
source_up = torch.tensor([[1.0, 0.0], [0.0, 2.0], [0.0, 0.0]])
|
|
settings = _settings(
|
|
cwb,
|
|
consensus_type="mean",
|
|
alignment_method="index",
|
|
preserve_common_prefix=True,
|
|
)
|
|
|
|
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
|
|
[reference_down, source_down],
|
|
[reference_up, source_up],
|
|
settings,
|
|
reference_index=0,
|
|
)
|
|
|
|
torch.testing.assert_close(merged_down[0], reference_down[0])
|
|
torch.testing.assert_close(merged_up[:, 0], reference_up[:, 0])
|
|
torch.testing.assert_close(merged_down[2], reference_down[2])
|
|
torch.testing.assert_close(merged_up[:, 2], reference_up[:, 2])
|
|
|
|
|
|
def test_norm_rescale_and_dsc_bandpass_are_finite(cwb):
|
|
rescale = _settings(
|
|
cwb,
|
|
consensus_type="mean",
|
|
rescale_norm=True,
|
|
)
|
|
merged = cwb.merge_consensus_group(
|
|
torch.tensor([[2.0, 0.0], [0.0, 2.0]]), rescale
|
|
)
|
|
assert torch.linalg.vector_norm(merged).item() == pytest.approx(2.0)
|
|
|
|
dsc = _settings(
|
|
cwb,
|
|
consensus_type="median",
|
|
diversity_beta=1.5,
|
|
dynamic_similarity_contrast=True,
|
|
soft_comfort_bandpass=True,
|
|
)
|
|
result = cwb.merge_consensus_group(
|
|
torch.tensor([[1.0, 0.0], [0.8, 0.2], [0.0, 1.0]]), dsc
|
|
)
|
|
assert torch.isfinite(result).all()
|
|
|
|
|
|
def test_all_cwb_schemas_and_compact_control_contract(cwb):
|
|
assert len(cwb.CWB_MERGER_NODES) == 16
|
|
schemas = [node.define_schema() for node in cwb.CWB_MERGER_NODES]
|
|
assert len({schema.node_id for schema in schemas}) == 16
|
|
advanced = {
|
|
"consensus_type", "alignment_method", "alignment_threshold",
|
|
"similarity_threshold", "power_alpha", "diversity_beta",
|
|
"rescale_norm", "global_scale", "dynamic_similarity_contrast",
|
|
"soft_comfort_bandpass", "position_weight", "preserve_common_prefix",
|
|
}
|
|
config_schema = cwb.CWBCustomConfiguration.define_schema()
|
|
assert {value.id for value in config_schema.inputs} == advanced
|
|
|
|
for node in cwb.CWB_MERGER_NODES:
|
|
schema = node.define_schema()
|
|
ids = [value.id for value in schema.inputs]
|
|
assert schema.description
|
|
assert all(value.tooltip for value in schema.inputs)
|
|
if node in (cwb.CWBCustomConfiguration, cwb.CWBEmbeddingStack):
|
|
continue
|
|
assert [output.display_name for output in schema.outputs] == [
|
|
"output_filename", "documentation", "cwb_report"
|
|
]
|
|
assert schema.outputs[0].get_io_type() == "*"
|
|
assert schema.outputs[1].get_io_type() == "STRING"
|
|
assert schema.outputs[2].get_io_type() == "STRING"
|
|
assert advanced.isdisjoint(ids)
|
|
assert "cwb_config" in ids
|
|
if node is cwb.CWBLoRAMultiMerger:
|
|
assert ids[:4] == ["execution_mode", "lora_count", "lora_1", "lora_2"]
|
|
assert schema.inputs[1].default == "2"
|
|
assert ids[-3:] == [
|
|
"include_1d_diffs",
|
|
"counterfactual_weight_sweep",
|
|
"include_mode",
|
|
]
|
|
continue
|
|
if node in (
|
|
cwb.CWBEmbeddingSelfCoalesce,
|
|
cwb.CWBEmbeddingMultiMerger,
|
|
cwb.CWBEmbeddingStackMerger,
|
|
):
|
|
assert "vision_boundary_embeddings" in ids
|
|
assert "legacy_boundary_search" in ids
|
|
assert "boundary_reference_embedding" in ids
|
|
assert "boundary_similarity_threshold" in ids
|
|
assert "target_vector_count" in ids
|
|
continue
|
|
assert ids[:3] == ["execution_mode", "model_a", "model_b"]
|
|
if node.INPUT_COUNT == 3:
|
|
assert ids[3] == "model_c"
|
|
if node.LORA_MODE:
|
|
assert schema.inputs[-2].id == "include_1d_diffs"
|
|
else:
|
|
assert "include_1d_diffs" not in ids
|
|
assert schema.inputs[-1].id == "include_mode"
|
|
assert schema.inputs[-1].default is False
|
|
|
|
|
|
def test_use_case_preset_names_match_complete_settings(cwb):
|
|
fields = set(cwb.CWBSettings.__dataclass_fields__)
|
|
for registry in (
|
|
cwb.DENSE_CWB_PRESETS,
|
|
cwb.EMBEDDING_CWB_PRESETS,
|
|
cwb.LORA_CWB_PRESETS,
|
|
):
|
|
assert len(registry) >= 6
|
|
for name, values in registry.items():
|
|
assert set(values) == fields
|
|
assert ("_rn" in name) is values["rescale_norm"]
|
|
assert ("_dsc" in name) is values["dynamic_similarity_contrast"]
|
|
assert ("_softcb" in name) is values["soft_comfort_bandpass"]
|
|
assert ("_pcp" in name) is values["preserve_common_prefix"]
|
|
if registry is not cwb.DENSE_CWB_PRESETS:
|
|
expected_alignment = "similarity" if "_sim_" in name else "index"
|
|
assert values["alignment_method"] == expected_alignment
|
|
expected_consensus = "median" if "_medn" in name else "mean"
|
|
assert values["consensus_type"] == expected_consensus
|
|
|
|
|
|
def test_connected_custom_config_completely_overrides_preset(cwb):
|
|
custom = _settings(
|
|
cwb,
|
|
consensus_type="median",
|
|
alignment_method="similarity",
|
|
alignment_threshold=0.17,
|
|
similarity_threshold=-0.25,
|
|
diversity_beta=7.0,
|
|
dynamic_similarity_contrast=True,
|
|
)
|
|
resolved = cwb.resolve_cwb_settings(
|
|
"direct_idx_medn_rn_softcb", cwb.LORA_CWB_PRESETS, custom
|
|
)
|
|
assert resolved is custom
|
|
|
|
|
|
def test_diagnostics_distinguish_alignment_and_weight_fallbacks(cwb):
|
|
alignment_diagnostics = cwb.CWBDiagnostics()
|
|
settings = _settings(
|
|
cwb, alignment_method="similarity", alignment_threshold=0.9
|
|
)
|
|
cwb.merge_cwb_tensors(
|
|
[torch.eye(2), -torch.eye(2)],
|
|
settings,
|
|
diagnostics=alignment_diagnostics,
|
|
)
|
|
assert alignment_diagnostics.alignment_matches == 0
|
|
assert alignment_diagnostics.anchor_only_groups == 2
|
|
|
|
weighting_diagnostics = cwb.CWBDiagnostics()
|
|
cwb.merge_consensus_group(
|
|
torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
|
|
_settings(cwb, similarity_threshold=0.9),
|
|
diagnostics=weighting_diagnostics,
|
|
)
|
|
assert weighting_diagnostics.all_rejected_fallbacks == 1
|
|
assert weighting_diagnostics.weighting_contributors == 2
|
|
assert weighting_diagnostics.accepted_weighting_contributors == 0
|
|
assert weighting_diagnostics.dominant_weight_min == pytest.approx(0.5)
|
|
assert weighting_diagnostics.dominant_weight_max == pytest.approx(0.5)
|
|
assert weighting_diagnostics.effective_contributor_min == pytest.approx(2.0)
|
|
assert weighting_diagnostics.effective_contributor_max == pytest.approx(2.0)
|
|
|
|
|
|
def test_weight_diagnostics_measure_parameter_influence(cwb):
|
|
diagnostics = cwb.CWBDiagnostics()
|
|
cwb.merge_consensus_group(
|
|
torch.tensor([[1.0, 0.0], [0.8, 0.2], [-1.0, 0.0]]),
|
|
_settings(
|
|
cwb,
|
|
similarity_threshold=0.0,
|
|
power_alpha=2.0,
|
|
diversity_beta=0.0,
|
|
),
|
|
diagnostics=diagnostics,
|
|
)
|
|
|
|
assert diagnostics.weighting_groups == 1
|
|
assert diagnostics.weighting_contributors == 3
|
|
assert diagnostics.accepted_weighting_contributors == 2
|
|
assert diagnostics.dominant_weight_max > 0.5
|
|
assert 1.0 < diagnostics.effective_contributor_min < 2.0
|
|
report = diagnostics.render("test", 3, custom=False)
|
|
assert "Accepted weighting contributors: 2/3" in report
|
|
assert "Dominant normalized weight min/mean/max:" in report
|
|
assert "Effective contributors min/mean/max:" in report
|
|
|
|
|
|
def test_counterfactual_weight_sweep_reuses_similarity_vectors(cwb):
|
|
sweep = cwb.CWBWeightSweepDiagnostics()
|
|
similarities = torch.tensor([0.25, 0.5, 0.75])
|
|
sweep._record_similarities("median", similarities)
|
|
|
|
baseline = sweep.stats[("median", 0.0, 2.0, 0.0, False, False)]
|
|
expected = similarities.square()
|
|
expected /= expected.sum()
|
|
assert baseline.groups == 1
|
|
assert baseline.accepted == 3
|
|
assert baseline.dominant_sum == pytest.approx(float(expected.max()))
|
|
assert baseline.effective_sum == pytest.approx(
|
|
float(1.0 / expected.square().sum())
|
|
)
|
|
assert len(sweep.stats) == 180
|
|
assert all(not isinstance(value, torch.Tensor) for value in sweep.__dict__.values())
|
|
|
|
|
|
def test_default_lora_preset_changes_perturbed_factors(cwb):
|
|
settings = cwb.resolve_cwb_settings(
|
|
"broad_sim_medn_rn_softcb", cwb.LORA_CWB_PRESETS
|
|
)
|
|
reference_down = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
|
|
reference_up = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
|
|
source_down = torch.tensor([[1.0, 0.1], [0.1, 1.0]])
|
|
source_up = torch.tensor([[1.1, 0.1], [0.1, 0.9]])
|
|
diagnostics = cwb.CWBDiagnostics()
|
|
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
|
|
[reference_down, source_down],
|
|
[reference_up, source_up],
|
|
settings,
|
|
reference_index=0,
|
|
diagnostics=diagnostics,
|
|
)
|
|
assert diagnostics.alignment_matches == 2
|
|
assert diagnostics.anchor_only_groups == 0
|
|
assert not torch.equal(merged_down, reference_down)
|
|
assert not torch.equal(merged_up, reference_up)
|
|
|
|
|
|
def test_eight_input_lora_merge_normalizes_before_variable_rank_alignment(
|
|
monkeypatch, tmp_path, cwb
|
|
):
|
|
paths = {}
|
|
events = []
|
|
original_normalize = cwb.normalize_lora_pair
|
|
original_merge = cwb.merge_cwb_lora_pairs
|
|
original_mark_processed = cwb.MemoryEfficientSafeOpen.mark_processed
|
|
received_ranks = []
|
|
processed = []
|
|
for index in range(8):
|
|
rank = 1 if index < 4 else 2
|
|
name = f"lora_{index}"
|
|
path = tmp_path / f"{name}.safetensors"
|
|
save_file({
|
|
"diffusion_model.foo.lora_A.weight": torch.full(
|
|
(rank, 2), 1.0 + index * 0.01
|
|
),
|
|
"diffusion_model.foo.lora_B.weight": torch.full(
|
|
(2, rank), 1.0 + index * 0.02
|
|
),
|
|
"diffusion_model.foo.alpha": torch.tensor(float(rank) * 0.5),
|
|
}, str(path))
|
|
paths[name] = str(path)
|
|
|
|
def record_normalize(*args, **kwargs):
|
|
events.append("normalize")
|
|
return original_normalize(*args, **kwargs)
|
|
|
|
def record_merge(downs, ups, *args, **kwargs):
|
|
events.append("merge")
|
|
received_ranks.extend(down.shape[0] for down in downs)
|
|
for down, up in zip(downs, ups):
|
|
assert down.shape[0] == up.shape[1]
|
|
return original_merge(downs, ups, *args, **kwargs)
|
|
|
|
def record_processed(loader, key):
|
|
processed.append(key)
|
|
return original_mark_processed(loader, key)
|
|
|
|
monkeypatch.setattr(cwb, "normalize_lora_pair", record_normalize)
|
|
monkeypatch.setattr(cwb, "merge_cwb_lora_pairs", record_merge)
|
|
monkeypatch.setattr(
|
|
cwb.MemoryEfficientSafeOpen, "mark_processed", record_processed
|
|
)
|
|
_patch_io(monkeypatch, cwb, tmp_path, paths)
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
list(paths), "loras", _params("eight_input"), lora_mode=True
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
|
|
assert tensors["diffusion_model.foo.lora_A.weight"].shape[0] == 2
|
|
assert "diffusion_model.foo.alpha" not in tensors
|
|
assert events.count("normalize") == 8
|
|
assert events.index("merge") > max(
|
|
index for index, event in enumerate(events) if event == "normalize"
|
|
)
|
|
assert received_ranks == [1, 1, 1, 1, 2, 2, 2, 2]
|
|
assert len(processed) == 24
|
|
assert processed.count("diffusion_model.foo.lora_A.weight") == 8
|
|
assert processed.count("diffusion_model.foo.lora_B.weight") == 8
|
|
assert processed.count("diffusion_model.foo.alpha") == 8
|
|
|
|
|
|
def test_lora_mismatch_zeros_remains_explicit_across_full_output_rank(
|
|
monkeypatch, tmp_path, cwb
|
|
):
|
|
a = tmp_path / "mismatch_zero_a.safetensors"
|
|
b = tmp_path / "mismatch_zero_b.safetensors"
|
|
save_file({
|
|
"diffusion_model.foo.lora_A.weight": torch.eye(2),
|
|
"diffusion_model.foo.lora_B.weight": torch.eye(2),
|
|
}, str(a))
|
|
save_file({"unrelated": 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(
|
|
"mismatch_zero",
|
|
mismatch_mode="zeros",
|
|
consensus_type="mean",
|
|
alignment_method="index",
|
|
rescale_norm=True,
|
|
save_dtype="fp32",
|
|
),
|
|
lora_mode=True,
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
|
|
|
|
torch.testing.assert_close(
|
|
tensors["diffusion_model.foo.lora_A.weight"], torch.eye(2) * 0.5
|
|
)
|
|
torch.testing.assert_close(
|
|
tensors["diffusion_model.foo.lora_B.weight"], torch.eye(2) * 0.5
|
|
)
|
|
|
|
|
|
def test_lora_multi_wrapper_validates_count_and_forwards_equal_prior_names(
|
|
monkeypatch, cwb
|
|
):
|
|
kwargs = {
|
|
"execution_mode": "MERGE",
|
|
"lora_count": "3",
|
|
"lora_1": "a.safetensors",
|
|
"lora_2": "b.safetensors",
|
|
"lora_3": "None",
|
|
"lora_4": "None",
|
|
"lora_5": "None",
|
|
"lora_6": "None",
|
|
"lora_7": "None",
|
|
"lora_8": "None",
|
|
"cwb_preset": "broad_sim_medn_rn_softcb",
|
|
"cwb_config": None,
|
|
}
|
|
monkeypatch.setattr(cwb, "load_documentation_from_file", lambda _: "DOCS")
|
|
with pytest.raises(ValueError, match="unselected input.*3"):
|
|
cwb.CWBLoRAMultiMerger.execute(**kwargs)
|
|
|
|
captured = {}
|
|
|
|
def execute(names, model_type, params, **options):
|
|
captured.update(
|
|
names=names,
|
|
model_type=model_type,
|
|
params=params,
|
|
options=options,
|
|
)
|
|
return "merged.safetensors", "REPORT"
|
|
|
|
monkeypatch.setattr(cwb.ConsensusMergerLogic, "execute", execute)
|
|
kwargs["lora_3"] = "c.safetensors"
|
|
result = cwb.CWBLoRAMultiMerger.execute(**kwargs)
|
|
assert captured["names"] == [
|
|
"a.safetensors", "b.safetensors", "c.safetensors"
|
|
]
|
|
assert captured["model_type"] == "loras"
|
|
assert captured["options"] == {"lora_mode": True}
|
|
assert result.result == ("merged.safetensors", "DOCS", "REPORT")
|
|
|
|
|
|
def test_embedding_coalescing_uses_mutual_local_pairs(cwb):
|
|
rows = torch.tensor([
|
|
[1.0, 0.0],
|
|
[0.99, 0.01],
|
|
[0.0, 1.0],
|
|
])
|
|
merged, pairs = cwb.coalesce_embedding_rows(
|
|
rows,
|
|
_settings(cwb),
|
|
target_vector_count=2,
|
|
similarity_threshold=0.95,
|
|
position_window=1.0,
|
|
)
|
|
assert merged.shape == (2, 2)
|
|
assert pairs == ((0, 1),)
|
|
torch.testing.assert_close(merged[1], rows[2])
|
|
|
|
|
|
def test_visual_boundary_preparation_validates_and_searches_legacy_rows(cwb):
|
|
start = torch.tensor([1.0, 0.0])
|
|
end = torch.tensor([0.0, 1.0])
|
|
body = torch.tensor([[0.6, 0.4]])
|
|
legacy = torch.stack((torch.tensor([-1.0, 0.0]), start, body[0], end, torch.tensor([0.0, -1.0])))
|
|
prepared, scores, trimmed = cwb._prepare_vision_embedding(
|
|
legacy,
|
|
start_reference=start,
|
|
end_reference=end,
|
|
legacy_search=True,
|
|
threshold=0.99,
|
|
)
|
|
torch.testing.assert_close(prepared, body)
|
|
assert scores == pytest.approx((1.0, 1.0))
|
|
assert trimmed == 2
|
|
|
|
with pytest.raises(ValueError, match="Vision boundary validation failed"):
|
|
cwb._prepare_vision_embedding(
|
|
torch.stack((start, body[0], -end)),
|
|
start_reference=start,
|
|
end_reference=end,
|
|
legacy_search=False,
|
|
threshold=0.99,
|
|
)
|
|
|
|
|
|
def test_embedding_self_coalesce_streams_and_restores_visual_boundaries(
|
|
monkeypatch, tmp_path, cwb
|
|
):
|
|
source = tmp_path / "source.safetensors"
|
|
start = torch.tensor([1.0, 0.0])
|
|
end = torch.tensor([0.0, 1.0])
|
|
save_file({"visual": torch.stack((start, torch.tensor([0.99, 0.01]), torch.tensor([0.98, 0.02]), end))}, str(source))
|
|
_patch_io(monkeypatch, cwb, tmp_path, {"source": str(source)})
|
|
monkeypatch.setattr(cwb, "load_documentation_from_file", lambda _: "DOCS")
|
|
result = cwb.CWBEmbeddingSelfCoalesce.execute(
|
|
embedding="source",
|
|
execution_mode="MERGE",
|
|
cwb_preset="balanced_sim_mean",
|
|
cwb_config=_settings(cwb),
|
|
target_vector_count=1,
|
|
coalesce_similarity_threshold=0.95,
|
|
coalesce_position_window=1.0,
|
|
vision_boundary_embeddings=True,
|
|
legacy_boundary_search=False,
|
|
boundary_reference_embedding="None",
|
|
boundary_similarity_threshold=0.99,
|
|
output_filename="visual_coalesced",
|
|
save_dtype="fp32",
|
|
process_device="cpu",
|
|
lazy_load=True,
|
|
force_clear_cache=False,
|
|
override_dtype=False,
|
|
)
|
|
output = load_file(str(_result_path(tmp_path, "embeddings", result.result[0])))
|
|
assert output["visual"].shape == (3, 2)
|
|
torch.testing.assert_close(output["visual"][0], start)
|
|
torch.testing.assert_close(output["visual"][-1], end)
|
|
assert "CWB EMBEDDING COALESCING" in result.result[2]
|
|
assert "Similarity threshold: 0.9500" in result.result[2]
|
|
|
|
|
|
def test_embedding_multi_wrapper_forwards_selected_inputs(monkeypatch, cwb):
|
|
kwargs = {
|
|
"execution_mode": "MERGE",
|
|
"embedding_count": "3",
|
|
"embedding_1": "a.safetensors",
|
|
"embedding_2": "b.safetensors",
|
|
"embedding_3": "c.safetensors",
|
|
"embedding_4": "None",
|
|
"embedding_5": "None",
|
|
"embedding_6": "None",
|
|
"embedding_7": "None",
|
|
"embedding_8": "None",
|
|
"cwb_preset": "balanced_sim_mean",
|
|
"cwb_config": None,
|
|
"target_vector_count": 0,
|
|
"coalesce_similarity_threshold": 0.95,
|
|
"coalesce_position_window": 0.1,
|
|
"vision_boundary_embeddings": False,
|
|
"legacy_boundary_search": False,
|
|
"boundary_reference_embedding": "None",
|
|
"boundary_similarity_threshold": 0.95,
|
|
"output_filename": "merged",
|
|
"save_dtype": "fp32",
|
|
"process_device": "cpu",
|
|
"lazy_load": True,
|
|
"force_clear_cache": False,
|
|
"override_dtype": False,
|
|
}
|
|
monkeypatch.setattr(cwb, "load_documentation_from_file", lambda _: "DOCS")
|
|
captured = {}
|
|
|
|
def execute(names, model_type, params, **options):
|
|
captured.update(names=names, model_type=model_type, params=params, options=options)
|
|
return "merged.safetensors", "REPORT"
|
|
|
|
monkeypatch.setattr(cwb.ConsensusMergerLogic, "execute", execute)
|
|
result = cwb.CWBEmbeddingMultiMerger.execute(**kwargs)
|
|
assert captured["names"] == ["a.safetensors", "b.safetensors", "c.safetensors"]
|
|
assert captured["model_type"] == "embeddings"
|
|
assert captured["options"] == {"embedding_union": True}
|
|
assert captured["params"]["embedding_coalesce"] is True
|
|
assert result.result == ("merged.safetensors", "DOCS", "REPORT")
|
|
|
|
|
|
def test_embedding_stack_and_stack_merger(monkeypatch, cwb):
|
|
stack_schema = cwb.CWBEmbeddingStack.define_schema()
|
|
assert stack_schema.node_id == "CWBEmbeddingStack"
|
|
assert [inp.id for inp in stack_schema.inputs] == ["add_embedding", "embedding_list"]
|
|
assert stack_schema.outputs[0].get_io_type() == "COMFY_DYNAMICGROUP_V3"
|
|
|
|
# 1. Active recursive DynamicCombo stack
|
|
stack_data = {
|
|
"add_embedding": "True",
|
|
"embedding_name": "emb1.safetensors",
|
|
"enabled": True,
|
|
"add_embedding_2": {
|
|
"add_embedding_2": "True",
|
|
"embedding_name": "emb2.safetensors",
|
|
"enabled": False,
|
|
"add_embedding_3": {
|
|
"add_embedding_3": "True",
|
|
"embedding_name": "emb3.safetensors",
|
|
"enabled": True,
|
|
"add_embedding_4": {"add_embedding_4": "False"},
|
|
},
|
|
},
|
|
}
|
|
stack_out = cwb.CWBEmbeddingStack.execute(
|
|
add_embedding=stack_data, embedding_list="ignored1.safetensors, ignored2.safetensors"
|
|
)
|
|
assert stack_out.result[0] == [
|
|
{"embedding_name": "emb1.safetensors", "enabled": True},
|
|
{"embedding_name": "emb2.safetensors", "enabled": False},
|
|
{"embedding_name": "emb3.safetensors", "enabled": True},
|
|
]
|
|
|
|
# 2. Stack set to False -> uses newline/comma-separated string list
|
|
list_out = cwb.CWBEmbeddingStack.execute(
|
|
add_embedding="False",
|
|
embedding_list="first.safetensors, second.safetensors\nthird.safetensors, \nfourth.safetensors",
|
|
)
|
|
assert list_out.result[0] == [
|
|
{"embedding_name": "first.safetensors", "enabled": True},
|
|
{"embedding_name": "second.safetensors", "enabled": True},
|
|
{"embedding_name": "third.safetensors", "enabled": True},
|
|
{"embedding_name": "fourth.safetensors", "enabled": True},
|
|
]
|
|
|
|
# 3. CWBEmbeddingStackMerger is a child of CWBEmbeddingMultiMerger and executes enabled entries
|
|
assert issubclass(cwb.CWBEmbeddingStackMerger, cwb.CWBEmbeddingMultiMerger)
|
|
monkeypatch.setattr(cwb, "load_documentation_from_file", lambda _: "DOCS")
|
|
captured = {}
|
|
|
|
def execute(names, model_type, params, **options):
|
|
captured.update(names=names, model_type=model_type, params=params, options=options)
|
|
return "merged.safetensors", "REPORT"
|
|
|
|
monkeypatch.setattr(cwb.ConsensusMergerLogic, "execute", execute)
|
|
base_kwargs = {
|
|
"execution_mode": "MERGE",
|
|
"cwb_preset": "balanced_sim_mean",
|
|
"cwb_config": None,
|
|
"target_vector_count": 0,
|
|
"coalesce_similarity_threshold": 0.95,
|
|
"coalesce_position_window": 0.1,
|
|
"vision_boundary_embeddings": False,
|
|
"legacy_boundary_search": False,
|
|
"boundary_reference_embedding": "None",
|
|
"boundary_similarity_threshold": 0.95,
|
|
"output_filename": "merged",
|
|
"save_dtype": "fp32",
|
|
"process_device": "cpu",
|
|
"lazy_load": True,
|
|
"force_clear_cache": False,
|
|
"override_dtype": False,
|
|
}
|
|
res = cwb.CWBEmbeddingStackMerger.execute(
|
|
embeddings=stack_out.result[0], **base_kwargs
|
|
)
|
|
assert captured["names"] == ["emb1.safetensors", "emb3.safetensors"]
|
|
assert res.result == ("merged.safetensors", "DOCS", "REPORT")
|
|
|
|
with pytest.raises(ValueError, match="At least 2 enabled embeddings"):
|
|
cwb.CWBEmbeddingStackMerger.execute(
|
|
embeddings=[{"embedding_name": "only_one.safetensors", "enabled": True}],
|
|
**base_kwargs,
|
|
)
|
|
|
|
|
|
def test_streaming_dtype_nonfloat_and_secondary_preservation(
|
|
monkeypatch, tmp_path, cwb
|
|
):
|
|
a = tmp_path / "a.safetensors"
|
|
b = tmp_path / "b.safetensors"
|
|
save_file({
|
|
"shared": torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32),
|
|
"only_a": torch.tensor([5.0, 6.0], dtype=torch.float16),
|
|
"metadata_tensor": torch.tensor([1, 2], dtype=torch.int64),
|
|
}, str(a), metadata={"source": "A"})
|
|
save_file({
|
|
"shared": torch.tensor([[3.0, 0.0], [0.0, 3.0]], dtype=torch.float16),
|
|
"extra": torch.ones((2, 2)),
|
|
"metadata_tensor": torch.tensor([9, 9], dtype=torch.int64),
|
|
}, str(b))
|
|
paths = {"a": str(a), "b": str(b)}
|
|
_patch_io(monkeypatch, cwb, tmp_path, paths)
|
|
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"], "diffusion_models", _params("anchored")
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
|
|
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"],
|
|
"diffusion_models",
|
|
_params("anchored_override", override_dtype=True),
|
|
)
|
|
overridden_tensors = load_file(
|
|
str(_result_path(tmp_path, "diffusion_models", overridden))
|
|
)
|
|
assert overridden_tensors["shared"].dtype == torch.float16
|
|
assert overridden_tensors["metadata_tensor"].dtype == torch.int64
|
|
|
|
|
|
def test_failed_merge_preserves_existing_output_and_removes_temporary_file(
|
|
monkeypatch, tmp_path, cwb
|
|
):
|
|
a = tmp_path / "atomic_a.safetensors"
|
|
b = tmp_path / "atomic_b.safetensors"
|
|
output = tmp_path / "diffusion_models" / "atomic_output.safetensors"
|
|
output.parent.mkdir(parents=True)
|
|
save_file({"layer": torch.ones((2, 2))}, str(a))
|
|
save_file({"layer": torch.full((2, 2), 2.0)}, str(b))
|
|
save_file({"existing": torch.tensor([7.0])}, str(output))
|
|
paths = {"a": str(a), "b": str(b)}
|
|
_patch_io(monkeypatch, cwb, tmp_path, paths)
|
|
|
|
def fail_merge(*args, **kwargs):
|
|
raise RuntimeError("injected CWB failure")
|
|
|
|
monkeypatch.setattr(cwb, "merge_cwb_tensors", fail_merge)
|
|
with pytest.raises(RuntimeError, match="injected CWB failure"):
|
|
cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"], "diffusion_models", _params("atomic_output")
|
|
)
|
|
|
|
existing = load_file(str(output))
|
|
assert set(existing) == {"existing"}
|
|
torch.testing.assert_close(existing["existing"], torch.tensor([7.0]))
|
|
assert not list(output.parent.glob(".atomic_output.safetensors.*.tmp"))
|
|
|
|
|
|
def test_embedding_union_uses_longest_first_dimension(monkeypatch, tmp_path, cwb):
|
|
a = tmp_path / "embed_a.safetensors"
|
|
b = tmp_path / "embed_b.safetensors"
|
|
save_file({"emb": torch.tensor([[1.0, 0.0], [0.0, 1.0]])}, str(a))
|
|
save_file({
|
|
"emb": torch.tensor([[3.0, 0.0], [0.0, 3.0], [4.0, 4.0]]),
|
|
"secondary_only": torch.tensor([[7.0, 0.0]]),
|
|
}, str(b))
|
|
paths = {"a": str(a), "b": str(b)}
|
|
_patch_io(monkeypatch, cwb, tmp_path, paths)
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"],
|
|
"embeddings",
|
|
_params("embedding_union"),
|
|
embedding_union=True,
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, "embeddings", result)))
|
|
assert tensors["emb"].shape == (3, 2)
|
|
torch.testing.assert_close(tensors["emb"][:2], torch.tensor([[2.0, 0.0], [0.0, 2.0]]))
|
|
torch.testing.assert_close(tensors["emb"][2], torch.tensor([4.0, 4.0]))
|
|
torch.testing.assert_close(tensors["secondary_only"], torch.tensor([[7.0, 0.0]]))
|
|
|
|
|
|
def test_output_name_preserves_category_relative_subdirectory(
|
|
monkeypatch, tmp_path, cwb
|
|
):
|
|
a = tmp_path / "relative_a.safetensors"
|
|
b = tmp_path / "relative_b.safetensors"
|
|
save_file({"layer": torch.tensor([[1.0]])}, str(a))
|
|
save_file({"layer": 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("groupfolder/modelname"),
|
|
)
|
|
|
|
assert result == "groupfolder/modelname.safetensors"
|
|
assert (tmp_path / "loras" / "groupfolder" / "modelname.safetensors").is_file()
|
|
|
|
|
|
def test_three_input_streaming_merge(monkeypatch, tmp_path, cwb):
|
|
paths = {}
|
|
for name, value in (("a", 1.0), ("b", 3.0), ("c", 5.0)):
|
|
path = tmp_path / f"three_{name}.safetensors"
|
|
save_file({"layer": torch.tensor([[value, 0.0]])}, str(path))
|
|
paths[name] = str(path)
|
|
_patch_io(monkeypatch, cwb, tmp_path, paths)
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b", "c"], "diffusion_models", _params("three_inputs")
|
|
)
|
|
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["layer"]
|
|
torch.testing.assert_close(tensor, torch.tensor([[3.0, 0.0]]))
|
|
|
|
|
|
@pytest.mark.parametrize("model_count", [2, 3])
|
|
@pytest.mark.parametrize("prefixed_input", ["a", "b"])
|
|
def test_diffusion_cwb_matches_prefixed_quantized_and_bare_layers(
|
|
monkeypatch, tmp_path, cwb, model_count, prefixed_input,
|
|
):
|
|
names = ["a", "b", "c"][:model_count]
|
|
paths = {}
|
|
for index, name in enumerate(names):
|
|
path = tmp_path / f"mixed_{name}.safetensors"
|
|
key = f"{'model.diffusion_model.' if name == prefixed_input else ''}layer.weight"
|
|
values = torch.tensor([2.0, 4.0]) * (index + 1)
|
|
if name == prefixed_input:
|
|
stem = key[:-len(".weight")]
|
|
save_file({
|
|
key: values.to(torch.int8),
|
|
f"{stem}.weight_scale": torch.tensor(1.0),
|
|
f"{stem}.comfy_quant": torch.tensor(
|
|
list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
|
|
),
|
|
}, str(path))
|
|
else:
|
|
save_file({key: values}, str(path))
|
|
paths[name] = str(path)
|
|
_patch_io(monkeypatch, cwb, tmp_path, paths)
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
names, "diffusion_models", _params("mixed_cwb")
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
|
|
assert list(tensors) == ["layer.weight"]
|
|
torch.testing.assert_close(tensors["layer.weight"].float(), torch.tensor([2.0, 4.0]) * (model_count + 1) / 2)
|
|
|
|
|
|
def test_diffusion_cwb_excluded_quantized_layer_is_saved_dense(monkeypatch, tmp_path, cwb):
|
|
a = tmp_path / "excluded_quant_a.safetensors"
|
|
b = tmp_path / "excluded_quant_b.safetensors"
|
|
save_file({
|
|
"model.diffusion_model.layer.weight": torch.tensor([1, 2], dtype=torch.int8),
|
|
"model.diffusion_model.layer.weight_scale": torch.tensor(2.0),
|
|
"model.diffusion_model.layer.comfy_quant": torch.tensor(
|
|
list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
|
|
),
|
|
}, str(a))
|
|
save_file({"layer.weight": torch.tensor([9.0, 9.0])}, str(b))
|
|
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"], "diffusion_models", _params("excluded_quant", include_mode=True),
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
|
|
assert list(tensors) == ["layer.weight"]
|
|
torch.testing.assert_close(tensors["layer.weight"].float(), torch.tensor([2.0, 4.0]))
|
|
|
|
|
|
def test_diffusion_cwb_guarded_quantized_primary_is_saved_dense(monkeypatch, tmp_path, cwb):
|
|
a = tmp_path / "guarded_quant_a.safetensors"
|
|
b = tmp_path / "guarded_quant_b.safetensors"
|
|
save_file({
|
|
"model.diffusion_model.layer.weight": torch.tensor([1, 2], dtype=torch.int8),
|
|
"model.diffusion_model.layer.weight_scale": torch.tensor(2.0),
|
|
"model.diffusion_model.layer.comfy_quant": torch.tensor(
|
|
list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
|
|
),
|
|
}, str(a))
|
|
save_file({"layer.weight": torch.tensor([9, 9], dtype=torch.uint8)}, str(b))
|
|
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"], "diffusion_models", _params("guarded_quant"),
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
|
|
assert list(tensors) == ["layer.weight"]
|
|
torch.testing.assert_close(tensors["layer.weight"].float(), torch.tensor([2.0, 4.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]))
|
|
|
|
|
|
@pytest.mark.parametrize("lora_mode", [False, True])
|
|
@pytest.mark.parametrize("patterns, glob_patterns, selected", [
|
|
(r"\.selected", False, True),
|
|
("*.selected*", True, True),
|
|
("", False, False),
|
|
("no_match", False, False),
|
|
])
|
|
def test_include_filter_preserves_nonmatches_and_discards_first(
|
|
monkeypatch, tmp_path, cwb, lora_mode, patterns, glob_patterns, selected,
|
|
):
|
|
paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
|
|
for name, value in (("a", 2.0), ("b", 6.0)):
|
|
save_file({
|
|
f"diffusion_model.{layer}.diff": torch.full((2, 2), value)
|
|
for layer in ("selected", "other", "selected_discard")
|
|
}, paths[name])
|
|
_patch_io(monkeypatch, cwb, tmp_path, paths)
|
|
category = "loras" if lora_mode else "diffusion_models"
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"], category,
|
|
_params("include_filter", exclude_patterns=patterns, include_mode=True,
|
|
glob_patterns=glob_patterns, discard_patterns="*discard*" if glob_patterns else "discard",
|
|
power_alpha=0.0),
|
|
lora_mode=lora_mode,
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, category, result)))
|
|
prefix = "diffusion_model." if lora_mode else ""
|
|
assert set(tensors) == {f"{prefix}selected.diff", f"{prefix}other.diff"}
|
|
torch.testing.assert_close(tensors[f"{prefix}selected.diff"], torch.full((2, 2), 4.0 if selected else 2.0))
|
|
torch.testing.assert_close(tensors[f"{prefix}other.diff"], torch.full((2, 2), 2.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
|
|
):
|
|
a = tmp_path / "lora_a.safetensors"
|
|
b = tmp_path / "lora_b.safetensors"
|
|
save_file({
|
|
"diffusion_model.foo.lora_A.weight": torch.tensor([[1.0, 2.0]], dtype=torch.float32),
|
|
"diffusion_model.foo.lora_B.weight": torch.tensor([[1.0], [2.0]], dtype=torch.float32),
|
|
"diffusion_model.foo.dora_scale": torch.tensor([1.0, 1.0]),
|
|
}, str(a))
|
|
save_file({
|
|
"lora_unet_foo.lora_down.weight": torch.tensor([[3.0, 4.0], [5.0, 6.0]]),
|
|
"lora_unet_foo.lora_up.weight": torch.tensor([[3.0, 5.0], [4.0, 6.0]]),
|
|
"lora_unet_secondary.lora_down.weight": torch.ones((1, 2)),
|
|
"lora_unet_secondary.lora_up.weight": torch.ones((2, 1)),
|
|
}, str(b))
|
|
paths = {"a": str(a), "b": str(b)}
|
|
_patch_io(monkeypatch, cwb, tmp_path, paths)
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"], "loras", _params("lora_cross_format"), lora_mode=True
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
|
|
assert set(tensors) == {
|
|
"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)
|
|
assert tensors["diffusion_model.foo.lora_A.weight"].dtype == torch.float32
|
|
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_alpha_normalized(
|
|
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(0.5),
|
|
}, 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([[2.0], [2.5]]),
|
|
)
|
|
assert tensors["diffusion_model.extra.lora_A.weight"].dtype == torch.float32
|
|
assert "diffusion_model.extra.alpha" not in tensors
|
|
|
|
|
|
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):
|
|
a = tmp_path / "peft_a.safetensors"
|
|
b = tmp_path / "peft_b.safetensors"
|
|
save_file({
|
|
"base_model.model.diffusion_model.foo.lora_A.weight": torch.ones((1, 2)),
|
|
"base_model.model.diffusion_model.foo.lora_B.weight": torch.ones((2, 1)),
|
|
"base_model.model.diffusion_model.foo.alpha": torch.tensor(1.0),
|
|
}, str(a))
|
|
save_file({
|
|
"diffusion_model.foo.lora_down.weight": torch.ones((2, 2)),
|
|
"diffusion_model.foo.lora_up.weight": torch.ones((2, 2)),
|
|
"diffusion_model.foo.alpha": torch.tensor(2.0),
|
|
}, str(b))
|
|
paths = {"a": str(a), "b": str(b)}
|
|
_patch_io(monkeypatch, cwb, tmp_path, paths)
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"], "loras", _params("peft_prefix"), lora_mode=True
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
|
|
assert set(tensors) == {
|
|
"diffusion_model.foo.lora_A.weight",
|
|
"diffusion_model.foo.lora_B.weight",
|
|
}
|
|
assert "diffusion_model.foo.alpha" not in tensors
|
|
|
|
|
|
def test_lora_mochi_inputs_emit_preferred_canonical_output(
|
|
monkeypatch, tmp_path, cwb
|
|
):
|
|
a = tmp_path / "mochi_a.safetensors"
|
|
b = tmp_path / "mochi_b.safetensors"
|
|
for path, value in ((a, 1.0), (b, 3.0)):
|
|
save_file({
|
|
"diffusion_model.foo.lora_A": torch.full((1, 2), value),
|
|
"diffusion_model.foo.lora_B": torch.full((2, 1), value),
|
|
"diffusion_model.foo.alpha": torch.tensor(1.0),
|
|
}, str(path))
|
|
paths = {"a": str(a), "b": str(b)}
|
|
_patch_io(monkeypatch, cwb, tmp_path, paths)
|
|
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"], "loras", _params("mochi_canonical"), lora_mode=True
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
|
|
|
|
assert set(tensors) == {
|
|
"diffusion_model.foo.lora_A.weight",
|
|
"diffusion_model.foo.lora_B.weight",
|
|
}
|
|
|
|
|
|
def test_lora_1d_direct_switch_is_default_fallback_and_fp32(
|
|
monkeypatch, tmp_path, cwb
|
|
):
|
|
a = tmp_path / "direct_a.safetensors"
|
|
b = tmp_path / "direct_b.safetensors"
|
|
save_file({"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16)}, str(a))
|
|
save_file({"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.bfloat16)}, str(b))
|
|
paths = {"a": str(a), "b": str(b)}
|
|
_patch_io(monkeypatch, cwb, tmp_path, paths)
|
|
|
|
default = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"], "loras", _params("direct_default"), lora_mode=True
|
|
)
|
|
default_tensor = load_file(str(_result_path(tmp_path, "loras", default)))[
|
|
"diffusion_model.norm.diff"
|
|
]
|
|
torch.testing.assert_close(default_tensor, torch.tensor([1.0, 2.0], dtype=torch.bfloat16))
|
|
assert default_tensor.dtype == torch.bfloat16
|
|
|
|
enabled = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"],
|
|
"loras",
|
|
_params("direct_enabled", include_1d_diffs=True, override_dtype=True),
|
|
lora_mode=True,
|
|
)
|
|
enabled_tensor = load_file(str(_result_path(tmp_path, "loras", enabled)))[
|
|
"diffusion_model.norm.diff"
|
|
]
|
|
torch.testing.assert_close(enabled_tensor, torch.tensor([2.0, 3.0]))
|
|
assert enabled_tensor.dtype == torch.float32
|
|
|
|
|
|
def test_invalid_quantization_sidecar_errors_before_output(monkeypatch, tmp_path, cwb):
|
|
a = tmp_path / "quant_a.safetensors"
|
|
b = tmp_path / "quant_b.safetensors"
|
|
save_file({
|
|
"layer.weight": torch.ones((2, 2)),
|
|
"layer.comfy_quant": torch.tensor([1], dtype=torch.uint8),
|
|
}, str(a))
|
|
save_file({"layer.weight": torch.ones((2, 2))}, str(b))
|
|
paths = {"a": str(a), "b": str(b)}
|
|
_patch_io(monkeypatch, cwb, tmp_path, paths)
|
|
with pytest.raises(ValueError, match="Invalid quantization sidecar"):
|
|
cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"], "diffusion_models", _params("must_not_exist")
|
|
)
|
|
assert not (tmp_path / "diffusion_models" / "must_not_exist.safetensors").exists()
|
|
|
|
|
|
def test_isolated_low_bit_preserves_anchored_tensor(monkeypatch, tmp_path, cwb):
|
|
a = tmp_path / "guard_a.safetensors"
|
|
b = tmp_path / "guard_b.safetensors"
|
|
save_file({"layer": torch.tensor([1.0, 2.0])}, str(a))
|
|
save_file({"layer": torch.tensor([8, 9], dtype=torch.uint8)}, str(b))
|
|
paths = {"a": str(a), "b": str(b)}
|
|
_patch_io(monkeypatch, cwb, tmp_path, paths)
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"], "diffusion_models", _params("guarded_generic")
|
|
)
|
|
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["layer"]
|
|
torch.testing.assert_close(tensor, torch.tensor([1.0, 2.0]))
|
|
assert tensor.dtype == torch.float32
|
|
|
|
|
|
def test_isolated_low_bit_preserves_complete_lora_layer(monkeypatch, tmp_path, cwb):
|
|
a = tmp_path / "guard_lora_a.safetensors"
|
|
b = tmp_path / "guard_lora_b.safetensors"
|
|
save_file({
|
|
"diffusion_model.foo.lora_A.weight": torch.tensor([[1.0, 2.0]]),
|
|
"diffusion_model.foo.lora_B.weight": torch.tensor([[3.0], [4.0]]),
|
|
"diffusion_model.foo.alpha": torch.tensor(1.0),
|
|
"diffusion_model.foo.dora_scale": torch.tensor([5.0, 6.0]),
|
|
}, str(a))
|
|
save_file({
|
|
"diffusion_model.foo.lora_A.weight": torch.tensor([[9.0, 9.0]]),
|
|
"diffusion_model.foo.lora_B.weight": torch.tensor([[9.0], [9.0]]),
|
|
"diffusion_model.foo.alpha": torch.tensor(1, dtype=torch.uint8),
|
|
}, str(b))
|
|
paths = {"a": str(a), "b": str(b)}
|
|
_patch_io(monkeypatch, cwb, tmp_path, paths)
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"], "loras", _params("guarded_lora"), lora_mode=True
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
|
|
torch.testing.assert_close(
|
|
tensors["diffusion_model.foo.lora_A.weight"], torch.tensor([[1.0, 2.0]])
|
|
)
|
|
torch.testing.assert_close(
|
|
tensors["diffusion_model.foo.lora_B.weight"], torch.tensor([[3.0], [4.0]])
|
|
)
|
|
assert "diffusion_model.foo.alpha" not in tensors
|
|
torch.testing.assert_close(
|
|
tensors["diffusion_model.foo.dora_scale"], torch.tensor([5.0, 6.0])
|
|
)
|
|
|
|
|
|
def test_alpha_bearing_low_bit_lora_factors_are_rejected(monkeypatch, tmp_path, cwb):
|
|
a = tmp_path / "alpha_low_bit_a.safetensors"
|
|
b = tmp_path / "alpha_low_bit_b.safetensors"
|
|
save_file({
|
|
"diffusion_model.foo.lora_A.weight": torch.tensor([[1.0, 2.0]]),
|
|
"diffusion_model.foo.lora_B.weight": torch.tensor([[3.0], [4.0]]),
|
|
"diffusion_model.foo.alpha": torch.tensor(1.0),
|
|
}, str(a))
|
|
save_file({
|
|
"diffusion_model.foo.lora_A.weight": torch.tensor([[1, 2]], dtype=torch.uint8),
|
|
"diffusion_model.foo.lora_B.weight": torch.tensor([[3.0], [4.0]]),
|
|
"diffusion_model.foo.alpha": torch.tensor(1.0),
|
|
}, str(b))
|
|
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
|
|
|
|
with pytest.raises(ValueError, match="Cannot alpha-normalize low-bit LoRA factors"):
|
|
cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"], "loras", _params("alpha_low_bit"), lora_mode=True
|
|
)
|