From 66f795de974d3c6c4bac0617d4b4ef4bf2333def Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 23 Sep 2026 02:42:27 +0200 Subject: [PATCH] Support mixed-prefix quantized Consensus merges --- nodes/consensus_merger.py | 163 ++++++++++++++++++++++++--------- tests/test_consensus_merger.py | 84 ++++++++++++++++- 2 files changed, 198 insertions(+), 49 deletions(-) diff --git a/nodes/consensus_merger.py b/nodes/consensus_merger.py index a3485a1..b4fce44 100644 --- a/nodes/consensus_merger.py +++ b/nodes/consensus_merger.py @@ -39,6 +39,8 @@ from .merger import ( load_documentation_from_file, ) from .quantization_guard import ( + DiffusionQuantization, + diffusion_key_map, inspect_low_bit_input, write_preserved_tensor, ) @@ -1716,14 +1718,31 @@ class ConsensusMergerLogic: boundary_path, low_memory=params["lazy_load"] ) try: - low_bit_sets = [ - inspect_low_bit_input( - handler, - f"Input {index + 1} ({path})", - "CWB Merge", - ) - for index, (handler, path) in enumerate(zip(handlers, paths)) - ] + diffusion_quantizers = [] + diffusion_maps = [] + if model_type == "diffusion_models": + diffusion_quantizers = [ + DiffusionQuantization(handler, path, f"CWB input {index + 1}") + for index, (handler, path) in enumerate(zip(handlers, paths)) + ] + diffusion_maps = [ + diffusion_key_map(quantizer.data_keys, f"CWB input {index + 1}") + for index, quantizer in enumerate(diffusion_quantizers) + ] + low_bit_sets = [ + {logical for logical, physical in mapping.items() + if physical in quantizer.isolated_low_bit_keys} + for mapping, quantizer in zip(diffusion_maps, diffusion_quantizers) + ] + else: + low_bit_sets = [ + inspect_low_bit_input( + handler, + f"Input {index + 1} ({path})", + "CWB Merge", + ) + for index, (handler, path) in enumerate(zip(handlers, paths)) + ] settings = resolve_cwb_settings( params["cwb_preset"], _presets_for( @@ -1746,6 +1765,8 @@ class ConsensusMergerLogic: embedding_union=embedding_union, boundary_handler=boundary_handler, source_models=model_names, + diffusion_quantizers=diffusion_quantizers, + diffusion_maps=diffusion_maps, ) if params.get("_return_cwb_report"): report = diagnostics.render( @@ -1786,11 +1807,17 @@ class ConsensusMergerLogic: embedding_union, boundary_handler=None, source_models, + diffusion_quantizers, + diffusion_maps, ): primary = handlers[0] keys = set() - for handler in handlers: - keys.update(handler.keys()) + if diffusion_maps: + for mapping in diffusion_maps: + keys.update(mapping) + else: + for handler in handlers: + keys.update(handler.keys()) keys = sorted(keys) output_path = cls._output_path(model_type, params["output_filename"]) requested_dtype = _requested_dtype(params["save_dtype"]) @@ -1804,23 +1831,73 @@ class ConsensusMergerLogic: secondary_only_merged = 0 handler_map = dict(enumerate(handlers)) + def source_indices_for(key): + return [ + i for i, handler in enumerate(handlers) + if key in (diffusion_maps[i] if diffusion_maps else handler.keys()) + ] + + def physical_key(index, key): + return diffusion_maps[index][key] if diffusion_maps else key + + def source_dtype(index, key): + source_key = physical_key(index, key) + return ( + diffusion_quantizers[index].logical_dtype(source_key, requested_dtype) + if diffusion_maps else handlers[index].get_dtype(source_key) + ) + + def matches(key, patterns): + aliases = [key] + if diffusion_maps: + aliases.extend(mapping[key] for mapping in diffusion_maps if key in mapping) + return any( + _matches_any_pattern(alias, patterns, glob_mode=glob_mode) + for alias in aliases + ) + work_units = [] for key in keys: - source_indices = [i for i, handler in enumerate(handlers) if key in handler.keys()] - discarded = _matches_any_pattern(key, discard, glob_mode=glob_mode) + source_indices = source_indices_for(key) + discarded = matches(key, discard) guarded = any(key in low_bit_sets[i] for i in source_indices) - entries = {} if discarded or guarded else { - i: [key] for i in source_indices - } + if discarded: + entries = {} + elif guarded: + preserve_index = 0 if 0 in source_indices else source_indices[0] + source_key = physical_key(preserve_index, key) + entries = ( + {preserve_index: diffusion_quantizers[preserve_index].required_keys(source_key)} + if diffusion_maps and source_key in diffusion_quantizers[preserve_index].quantized_keys + else {} + ) + else: + entries = { + i: (diffusion_quantizers[i].required_keys(physical_key(i, key)) + if diffusion_maps else [key]) for i in source_indices + } work_units.append((key, entries)) - output_metadata = (primary.metadata() or {}).copy() + output_metadata = ( + diffusion_quantizers[0].output_metadata() + if diffusion_maps else (primary.metadata() or {}).copy() + ) output_metadata["cwb.merge"] = _cwb_merge_metadata( source_models, params, settings, operation="cwb_embedding_merge" if embedding_union else "cwb_tensor_merge", ) + def preserve_tensor(writer, index, key, loaded, *, force_raw=False): + source_key = physical_key(index, key) + if diffusion_maps and source_key in diffusion_quantizers[index].quantized_keys: + writer.write(key, loaded[(index, key)].to(requested_dtype).cpu().contiguous()) + else: + write_preserved_tensor( + writer, source_key, handlers[index], output_key=key, + force_raw=force_raw, tensor=loaded.get((index, key)), + ) + with atomic_uel_writer(output_path, output_metadata) as writer, torch.no_grad(), closing( stream_work_units( handler_map, work_units, @@ -1829,49 +1906,52 @@ class ConsensusMergerLogic: ) as streamed: for key, loaded in tqdm(streamed, total=len(work_units), desc="CWB merging tensors", unit="tensors"): _clear_previous_layer(params) - if _matches_any_pattern(key, discard, glob_mode=glob_mode): + if matches(key, discard): pbar.update(1) continue - source_indices = [i for i, handler in enumerate(handlers) if key in handler.keys()] - primary_owned = key in primary.keys() + source_indices = source_indices_for(key) + primary_owned = key in (diffusion_maps[0] if diffusion_maps else primary.keys()) secondary_only = not primary_owned preserve_index = 0 if primary_owned else source_indices[0] guarded = any(key in low_bit_sets[i] for i in source_indices) - matched = _matches_any_pattern(key, exclude, glob_mode=glob_mode) + if diffusion_maps: + for index in source_indices: + source_key = physical_key(index, key) + if (index, source_key) not in loaded: + continue + if source_key in diffusion_quantizers[index].quantized_keys: + parts = { + part: loaded[(index, part)] + for part in diffusion_quantizers[index].required_keys(source_key) + } + loaded[(index, key)] = diffusion_quantizers[index].decode( + source_key, parts, torch.float32, device="cpu", + ) + else: + loaded[(index, key)] = loaded[(index, source_key)] + matched = matches(key, exclude) excluded = not matched if include_mode else matched if guarded or excluded: - write_preserved_tensor( - writer, - key, - handlers[preserve_index], - force_raw=guarded, - tensor=loaded.get((preserve_index, key)), - ) + preserve_tensor(writer, preserve_index, key, loaded, force_raw=guarded) if secondary_only: secondary_only_copied += 1 pbar.update(1) continue if secondary_only and len(source_indices) == 1: - write_preserved_tensor( - writer, key, handlers[preserve_index], - tensor=loaded[(preserve_index, key)], - ) + preserve_tensor(writer, preserve_index, key, loaded) secondary_only_copied += 1 pbar.update(1) continue - source_dtypes = [handlers[i].get_dtype(key) for i in source_indices] + source_dtypes = [source_dtype(i, key) for i in source_indices] if not all(_is_float_dtype(dtype) for dtype in source_dtypes): logging.warning( "[CWB Merge] Preserving non-floating tensor '%s' from input %d.", key, preserve_index + 1, ) - write_preserved_tensor( - writer, key, handlers[preserve_index], - tensor=loaded[(preserve_index, key)], - ) + preserve_tensor(writer, preserve_index, key, loaded) if secondary_only: secondary_only_copied += 1 pbar.update(1) @@ -1881,9 +1961,7 @@ class ConsensusMergerLogic: if mismatch_mode == "error": raise ValueError(f"Tensor '{key}' is missing from a CWB input.") if mismatch_mode == "skip": - write_preserved_tensor( - writer, key, primary, tensor=loaded[(0, key)] - ) + preserve_tensor(writer, 0, key, loaded) pbar.update(1) continue @@ -1972,13 +2050,10 @@ class ConsensusMergerLogic: if source_index == reference_source: reference_index = len(tensors) tensors.append(aligned) - actual_dtypes.append(handler.get_dtype(key)) + actual_dtypes.append(source_dtype(source_index, key)) if preserve_for_mismatch: - write_preserved_tensor( - writer, key, handlers[preserve_index], - tensor=loaded[(preserve_index, key)], - ) + preserve_tensor(writer, preserve_index, key, loaded) if secondary_only: secondary_only_copied += 1 pbar.update(1) diff --git a/tests/test_consensus_merger.py b/tests/test_consensus_merger.py index 3a82d7f..e52aa4e 100644 --- a/tests/test_consensus_merger.py +++ b/tests/test_consensus_merger.py @@ -1,4 +1,5 @@ import importlib +import json import sys import types from pathlib import Path @@ -1194,6 +1195,78 @@ def test_three_input_streaming_merge(monkeypatch, tmp_path, cwb): 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" @@ -1289,9 +1362,10 @@ def test_include_filter_preserves_nonmatches_and_discards_first( lora_mode=lora_mode, ) tensors = load_file(str(_result_path(tmp_path, category, result))) - assert set(tensors) == {"diffusion_model.selected.diff", "diffusion_model.other.diff"} - torch.testing.assert_close(tensors["diffusion_model.selected.diff"], torch.full((2, 2), 4.0 if selected else 2.0)) - torch.testing.assert_close(tensors["diffusion_model.other.diff"], torch.full((2, 2), 2.0)) + 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): @@ -1556,7 +1630,7 @@ def test_lora_1d_direct_switch_is_default_fallback_and_fp32( assert enabled_tensor.dtype == torch.float32 -def test_quantized_marker_errors_before_output(monkeypatch, tmp_path, cwb): +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({ @@ -1566,7 +1640,7 @@ def test_quantized_marker_errors_before_output(monkeypatch, tmp_path, cwb): 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"): + with pytest.raises(ValueError, match="Invalid quantization sidecar"): cwb.ConsensusMergerLogic.execute( ["a", "b"], "diffusion_models", _params("must_not_exist") )