867 lines
32 KiB
Python
867 lines
32 KiB
Python
import importlib
|
|
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):
|
|
params = {
|
|
"cwb_preset": "custom",
|
|
"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,
|
|
"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 test_preset_resolution_matches_extended_contract(cwb):
|
|
assert all(
|
|
preset["global_scale"] == pytest.approx(1.0)
|
|
for preset in cwb.CWB_PRESETS.values()
|
|
)
|
|
|
|
power = cwb.resolve_cwb_settings(cwb_preset="power_blend")
|
|
assert power.consensus_type == "median"
|
|
assert power.alignment_threshold == pytest.approx(0.9)
|
|
assert power.power_alpha == pytest.approx(8.0)
|
|
assert power.dynamic_similarity_contrast is True
|
|
assert power.soft_comfort_bandpass is False
|
|
|
|
baseline = cwb.resolve_cwb_settings(
|
|
cwb_preset="baseline",
|
|
alignment_method="index",
|
|
dynamic_similarity_contrast=True,
|
|
soft_comfort_bandpass=True,
|
|
)
|
|
assert baseline.alignment_method == "similarity"
|
|
assert baseline.dynamic_similarity_contrast is False
|
|
assert baseline.soft_comfort_bandpass is False
|
|
|
|
varied = cwb.resolve_cwb_settings(
|
|
cwb_preset="varied_merge",
|
|
global_scale=1.25,
|
|
position_weight=0.2,
|
|
preserve_common_prefix=True,
|
|
)
|
|
assert varied.global_scale == pytest.approx(1.25)
|
|
assert varied.position_weight == pytest.approx(0.2)
|
|
assert varied.preserve_common_prefix is True
|
|
|
|
|
|
def test_cwb_math_index_prefix_and_scale(cwb):
|
|
settings = cwb.resolve_cwb_settings(
|
|
cwb_preset="custom",
|
|
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 = cwb.resolve_cwb_settings(
|
|
cwb_preset="custom",
|
|
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 = cwb.resolve_cwb_settings(
|
|
cwb_preset="custom",
|
|
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 = cwb.resolve_cwb_settings(
|
|
cwb_preset="custom",
|
|
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 = cwb.resolve_cwb_settings(
|
|
cwb_preset="custom",
|
|
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 = cwb.resolve_cwb_settings(
|
|
cwb_preset="custom",
|
|
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_lora_alpha_and_global_scale_apply_once_to_pair(cwb):
|
|
settings = cwb.resolve_cwb_settings(
|
|
cwb_preset="custom",
|
|
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(
|
|
[(0, down, up, 2.0), (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 = cwb.resolve_cwb_settings(
|
|
cwb_preset="custom",
|
|
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 = cwb.resolve_cwb_settings(
|
|
cwb_preset="custom",
|
|
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)
|
|
down = torch.ones((1, 2))
|
|
up = torch.ones((2, 1))
|
|
with caplog.at_level("WARNING"):
|
|
merged_down, merged_up, rank = cwb._merge_lora_pair_to_cpu(
|
|
[(0, down, up, 1.0), (1, down, up, 1.0)],
|
|
settings,
|
|
"cuda",
|
|
torch.float32,
|
|
operation_label="diffusion_model.foo",
|
|
)
|
|
|
|
assert rank == 1
|
|
torch.testing.assert_close(merged_down, down)
|
|
torch.testing.assert_close(merged_up, up)
|
|
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
|
|
|
|
|
|
def test_norm_rescale_and_dsc_bandpass_are_finite(cwb):
|
|
rescale = cwb.resolve_cwb_settings(
|
|
cwb_preset="custom",
|
|
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 = cwb.resolve_cwb_settings(
|
|
cwb_preset="custom",
|
|
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_ten_schemas_and_lora_switch_contract(cwb):
|
|
assert len(cwb.CWB_MERGER_NODES) == 10
|
|
assert len({node.NODE_ID for node in cwb.CWB_MERGER_NODES}) == 10
|
|
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)
|
|
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[-1].id == "include_1d_diffs"
|
|
assert schema.inputs[-1].default is False
|
|
else:
|
|
assert "include_1d_diffs" not in ids
|
|
|
|
|
|
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_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]]))
|
|
|
|
|
|
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
|
|
):
|
|
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_quantized_marker_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="ComfyUI quantization metadata"):
|
|
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
|
|
)
|