Support mixed-prefix quantized Consensus merges

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)

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")
         )