feat(merge): add dedicated CWB mergers
This commit is contained in:
silveroxides
@@ -20,6 +20,7 @@ from .nodes.merger import (
|
||||
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
|
||||
CheckpointThreeMerger, EmbeddingThreeMerger
|
||||
)
|
||||
from .nodes.consensus_merger import CWB_MERGER_NODES
|
||||
from .nodes.lora_extract_svd import (
|
||||
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
|
||||
LoRAExtractKnee, LoRAExtractFrobenius
|
||||
@@ -68,6 +69,8 @@ class ModelUtilsExtension(ComfyExtension):
|
||||
# Three-Model Mergers
|
||||
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
|
||||
CheckpointThreeMerger, EmbeddingThreeMerger,
|
||||
# Dedicated CWB Mergers
|
||||
*CWB_MERGER_NODES,
|
||||
# LoRA Extraction
|
||||
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
|
||||
LoRAExtractKnee, LoRAExtractFrobenius,
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
# Consensus-Weighted Blending Mergers
|
||||
|
||||
The dedicated CWB nodes merge two or three safetensors files using
|
||||
Consensus-Weighted Blending. They calculate an element-wise mean or median
|
||||
consensus, compare each aligned source vector to that consensus, and normalize
|
||||
similarity-derived weights before producing the output vector.
|
||||
|
||||
## Node families
|
||||
|
||||
Two-input and three-input variants are provided for checkpoints, standalone
|
||||
diffusion models, text encoders, LoRAs, and embeddings.
|
||||
|
||||
Checkpoint, diffusion-model, text-encoder, and LoRA outputs are anchored to
|
||||
Model A. Model A supplies the output key set, metadata, and naming. Keys found
|
||||
only in later inputs are ignored. Embedding mergers instead use the union of
|
||||
all keys and use the longest compatible first dimension as the alignment
|
||||
reference.
|
||||
|
||||
## Presets and custom controls
|
||||
|
||||
`baseline` is the default. The standard clarity, smooth, varied, and diversity
|
||||
presets are joined by `power_blend` and the DSC presets demonstrated by the CWB
|
||||
conditioning implementation. Select `custom` to make every manual CWB control
|
||||
authoritative.
|
||||
|
||||
Similarity alignment greedily matches source rows to the reference by cosine
|
||||
similarity. Index alignment keeps absolute row positions. Position weight can
|
||||
bias similarity matches toward nearby normalized positions. Preserve Common
|
||||
Prefix copies the longest numerically identical leading row span directly from
|
||||
the first input.
|
||||
|
||||
Inputs have equal prior weight. The similarity calculation supplies the
|
||||
effective contribution for each aligned vector. If every candidate is rejected
|
||||
or receives zero weight, CWB falls back to equal weighting.
|
||||
|
||||
## Missing tensors and filtering
|
||||
|
||||
- `skip` preserves Model A when an anchored input is missing or incompatible.
|
||||
In embedding-union mode it omits that source for the affected key.
|
||||
- `zeros` supplies a compatible zero tensor for the missing source.
|
||||
- `error` stops the merge.
|
||||
- Exclude patterns preserve the anchored source tensor.
|
||||
- Discard patterns remove the matching output tensor or logical LoRA layer.
|
||||
|
||||
Patterns use regular expressions by default. Enable Glob Patterns to use glob
|
||||
syntax instead.
|
||||
|
||||
## LoRA behavior
|
||||
|
||||
LoRA inputs use the same parser as the resize and multi-merge nodes. This
|
||||
includes the repository's plain `.lora_A.weight` / `.lora_B.weight` extraction
|
||||
output, already-supported alternative A/B or down/up suffixes, and direct
|
||||
`.diff` / `.diff_b` layers.
|
||||
|
||||
Model A's exact keys and suffix convention are retained. Secondary formats are
|
||||
matched by normalized logical layer name. Differing ranks are zero-padded like
|
||||
the DARE/TIES mergers: A/down tensors on dimension 0 and B/up tensors on
|
||||
dimension 1. An existing Model A alpha key is updated to the resulting maximum
|
||||
rank; no alpha key is invented when Model A has none. Incomplete pairs,
|
||||
auxiliary tensors, and unrecognized Model A tensors are copied unchanged.
|
||||
|
||||
`include_1d_diffs` is disabled by default. Disabled 1D direct layers preserve
|
||||
Model A. Enabled 1D direct layers participate as complete tensors and are
|
||||
always saved in FP32.
|
||||
|
||||
## Precision and quantized inputs
|
||||
|
||||
CWB computes in FP32. Unless Override Dtype is enabled, any participating FP32
|
||||
source keeps the generated logical layer in FP32. Guarded low-bit tensors and
|
||||
enabled 1D direct diffs are never downcast by the override.
|
||||
|
||||
ComfyUI/comfy-kitchen quantized inputs are rejected before an output writer is
|
||||
opened. Files with three or more bare INT8, UINT8, FP8, or FP4 tensors are also
|
||||
rejected. One or two isolated low-bit tensors produce one warning per input and
|
||||
are preserved without CWB arithmetic; an affected LoRA causes its complete
|
||||
logical layer to be preserved.
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,410 @@
|
||||
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, "prepare_for_large_operation", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
|
||||
|
||||
|
||||
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):
|
||||
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_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_model_a_anchored_streaming_dtype_and_nonfloat_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(tmp_path / result))
|
||||
assert set(tensors) == {"shared", "only_a", "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]))
|
||||
|
||||
overridden = cwb.ConsensusMergerLogic.execute(
|
||||
["a", "b"],
|
||||
"diffusion_models",
|
||||
_params("anchored_override", override_dtype=True),
|
||||
)
|
||||
overridden_tensors = load_file(str(tmp_path / overridden))
|
||||
assert overridden_tensors["shared"].dtype == torch.float16
|
||||
assert overridden_tensors["metadata_tensor"].dtype == torch.int64
|
||||
|
||||
|
||||
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(tmp_path / 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(tmp_path / result))["layer"]
|
||||
torch.testing.assert_close(tensor, torch.tensor([[3.0, 0.0]]))
|
||||
|
||||
|
||||
def test_lora_cross_format_rank_padding_preserves_model_a_names(
|
||||
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(tmp_path / result))
|
||||
assert set(tensors) == {
|
||||
"diffusion_model.foo.lora_A.weight",
|
||||
"diffusion_model.foo.lora_B.weight",
|
||||
"diffusion_model.foo.dora_scale",
|
||||
}
|
||||
assert tensors["diffusion_model.foo.lora_A.weight"].shape == (2, 2)
|
||||
assert tensors["diffusion_model.foo.lora_B.weight"].shape == (2, 2)
|
||||
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])
|
||||
)
|
||||
|
||||
|
||||
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(tmp_path / result))
|
||||
assert set(tensors) == {
|
||||
"base_model.model.diffusion_model.foo.lora_A.weight",
|
||||
"base_model.model.diffusion_model.foo.lora_B.weight",
|
||||
"base_model.model.diffusion_model.foo.alpha",
|
||||
}
|
||||
assert tensors["base_model.model.diffusion_model.foo.alpha"].item() == pytest.approx(2.0)
|
||||
|
||||
|
||||
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(tmp_path / 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(tmp_path / 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 / "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(tmp_path / 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(tmp_path / 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])
|
||||
)
|
||||
Reference in New Issue
Block a user