silveroxides
3bd4977d2a
Runtime conditioning scale reduction must not be carried into persisted model or LoRA weights. Keep named merger presets at 1.0 and reserve scaling for explicit custom configuration.
846 lines
31 KiB
Python
846 lines
31 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 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_preserved_without_scaling(
|
|
monkeypatch, tmp_path, cwb
|
|
):
|
|
a = tmp_path / "lora_secondary_a.safetensors"
|
|
b = tmp_path / "lora_secondary_b.safetensors"
|
|
save_file({
|
|
"diffusion_model.base.diff": torch.tensor([[1.0]]),
|
|
}, str(a))
|
|
save_file({
|
|
"transformer.extra.lora_down.weight": torch.tensor([[2.0, 3.0]]),
|
|
"transformer.extra.lora_up.weight": torch.tensor([[4.0], [5.0]]),
|
|
"transformer.extra.alpha": torch.tensor(1.0),
|
|
}, str(b))
|
|
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
|
|
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"],
|
|
"loras",
|
|
_params(
|
|
"lora_secondary_exact",
|
|
global_scale=2.0,
|
|
override_dtype=True,
|
|
save_dtype="fp16",
|
|
),
|
|
lora_mode=True,
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
|
|
torch.testing.assert_close(
|
|
tensors["diffusion_model.extra.lora_A.weight"],
|
|
torch.tensor([[2.0, 3.0]]),
|
|
)
|
|
torch.testing.assert_close(
|
|
tensors["diffusion_model.extra.lora_B.weight"],
|
|
torch.tensor([[4.0], [5.0]]),
|
|
)
|
|
assert tensors["diffusion_model.extra.lora_A.weight"].dtype == torch.float32
|
|
assert tensors["diffusion_model.extra.alpha"].item() == pytest.approx(1.0)
|
|
|
|
|
|
def test_lora_three_input_merges_b_and_c_group_without_a(monkeypatch, tmp_path, cwb):
|
|
a = tmp_path / "lora_union_a.safetensors"
|
|
b = tmp_path / "lora_union_b.safetensors"
|
|
c = tmp_path / "lora_union_c.safetensors"
|
|
save_file({"diffusion_model.base.diff": torch.tensor([[1.0]])}, str(a))
|
|
save_file({
|
|
"diffusion_model.extra.lora_A.weight": torch.full((1, 2), 1.0),
|
|
"diffusion_model.extra.lora_B.weight": torch.full((2, 1), 1.0),
|
|
}, str(b))
|
|
save_file({
|
|
"transformer.extra.lora_down.weight": torch.full((1, 2), 3.0),
|
|
"transformer.extra.lora_up.weight": torch.full((2, 1), 3.0),
|
|
}, str(c))
|
|
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b), "c": str(c)})
|
|
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b", "c"], "loras", _params("lora_secondary_union"), lora_mode=True
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
|
|
torch.testing.assert_close(
|
|
tensors["diffusion_model.extra.lora_A.weight"], torch.full((1, 2), 2.0)
|
|
)
|
|
torch.testing.assert_close(
|
|
tensors["diffusion_model.extra.lora_B.weight"], torch.full((2, 1), 2.0)
|
|
)
|
|
|
|
|
|
def test_lora_secondary_direct_and_companion_groups_are_atomic(
|
|
monkeypatch, tmp_path, cwb
|
|
):
|
|
a = tmp_path / "lora_atomic_a.safetensors"
|
|
b = tmp_path / "lora_atomic_b.safetensors"
|
|
save_file({"diffusion_model.base.diff": torch.tensor([[1.0]])}, str(a))
|
|
save_file({
|
|
"diffusion_model.norm.diff": torch.tensor([2.0, 3.0], dtype=torch.bfloat16),
|
|
"diffusion_model.extra.lora_A.weight": torch.tensor([[4.0, 5.0]]),
|
|
"diffusion_model.extra.lora_B.weight": torch.tensor([[6.0], [7.0]]),
|
|
"diffusion_model.extra.dora_scale": torch.tensor([8.0, 9.0]),
|
|
}, str(b))
|
|
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
|
|
|
|
result = cwb.ConsensusMergerLogic.execute(
|
|
["a", "b"],
|
|
"loras",
|
|
_params("lora_secondary_atomic", global_scale=3.0),
|
|
lora_mode=True,
|
|
)
|
|
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
|
|
torch.testing.assert_close(
|
|
tensors["diffusion_model.norm.diff"],
|
|
torch.tensor([2.0, 3.0], dtype=torch.bfloat16),
|
|
)
|
|
assert tensors["diffusion_model.norm.diff"].dtype == torch.bfloat16
|
|
torch.testing.assert_close(
|
|
tensors["diffusion_model.extra.lora_A.weight"], torch.tensor([[4.0, 5.0]])
|
|
)
|
|
torch.testing.assert_close(
|
|
tensors["diffusion_model.extra.lora_B.weight"], torch.tensor([[6.0], [7.0]])
|
|
)
|
|
torch.testing.assert_close(
|
|
tensors["diffusion_model.extra.dora_scale"], torch.tensor([8.0, 9.0])
|
|
)
|
|
|
|
|
|
def test_lora_rejects_duplicate_normalized_cores_in_secondary_input(cwb):
|
|
with pytest.raises(ValueError, match="same logical layer 'foo'.*input 2"):
|
|
cwb._secondary_lora_map(
|
|
{
|
|
"diffusion_model.foo": {},
|
|
"transformer.foo": {},
|
|
},
|
|
input_index=1,
|
|
)
|
|
|
|
|
|
def test_lora_peft_prefix_matching_and_existing_alpha(monkeypatch, tmp_path, cwb):
|
|
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",
|
|
"diffusion_model.foo.alpha",
|
|
}
|
|
assert tensors["diffusion_model.foo.alpha"].item() == pytest.approx(2.0)
|
|
|
|
|
|
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",
|
|
"diffusion_model.foo.alpha",
|
|
}
|
|
|
|
|
|
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 tensors["diffusion_model.foo.alpha"].item() == pytest.approx(1.0)
|
|
torch.testing.assert_close(
|
|
tensors["diffusion_model.foo.dora_scale"], torch.tensor([5.0, 6.0])
|
|
)
|