From bed9c6e8ffb2df12eeb19ef688ffccf36f44f4e3 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Sat, 1 Aug 2026 13:14:14 +0200 Subject: [PATCH] feat(merge): add dedicated CWB mergers --- __init__.py | 3 + docs/consensus_mergers.md | 76 +++ nodes/consensus_merger.py | 1010 ++++++++++++++++++++++++++++++++ tests/test_consensus_merger.py | 410 +++++++++++++ 4 files changed, 1499 insertions(+) create mode 100644 docs/consensus_mergers.md create mode 100644 nodes/consensus_merger.py create mode 100644 tests/test_consensus_merger.py diff --git a/__init__.py b/__init__.py index 89bbcb0..2e1e381 100644 --- a/__init__.py +++ b/__init__.py @@ -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, diff --git a/docs/consensus_mergers.md b/docs/consensus_mergers.md new file mode 100644 index 0000000..8df729c --- /dev/null +++ b/docs/consensus_mergers.md @@ -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. diff --git a/nodes/consensus_merger.py b/nodes/consensus_merger.py new file mode 100644 index 0000000..2c18c80 --- /dev/null +++ b/nodes/consensus_merger.py @@ -0,0 +1,1010 @@ +"""Dedicated streaming Consensus-Weighted Blending merger nodes.""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass +from typing import Iterable + +import comfy.utils +import folder_paths +import torch +import torch.nn.functional as F +from comfy_api.latest import io +from tqdm import tqdm +from unifiedefficientloader import ( + IncrementalSafetensorsWriter, + MemoryEfficientSafeOpen, + transfer_to_gpu_pinned, +) + +from .device_utils import ( + cleanup_after_operation, + estimate_model_size, + prepare_for_large_operation, +) +from .lora_resize import parse_lora_layers, select_output_dtype +from .merger import ( + _compile_patterns, + _matches_any_pattern, + load_documentation_from_file, +) +from .quantization_guard import ( + inspect_low_bit_input, + layer_has_low_bit, + write_preserved_tensor, +) + + +CWB_PRESETS = { + "baseline": { + "consensus_type": "median", "power_alpha": 2.0, + "similarity_threshold": 0.0, "diversity_beta": 0.0, + "rescale_norm": False, "global_scale": 1.0, + }, + "power_blend": { + "consensus_type": "median", "alignment_method": "similarity", + "alignment_threshold": 0.9, "power_alpha": 8.0, + "similarity_threshold": 0.75, "diversity_beta": 0.0, + "rescale_norm": True, "global_scale": 1.0, + "dynamic_similarity_contrast": True, + }, + "high_clarity": { + "consensus_type": "median", "power_alpha": 3.0, + "similarity_threshold": 0.3, "diversity_beta": 0.0, + "rescale_norm": False, "global_scale": 1.0, + }, + "smooth": { + "consensus_type": "mean", "power_alpha": 1.5, + "similarity_threshold": 0.0, "diversity_beta": 0.0, + "rescale_norm": False, "global_scale": 1.0, + }, + "varied_merge": { + "consensus_type": "median", "power_alpha": 2.0, + "similarity_threshold": 0.0, "diversity_beta": 0.0, + "rescale_norm": True, "global_scale": 0.7, + }, + "diverse_concept": { + "consensus_type": "median", "power_alpha": 2.0, + "similarity_threshold": 0.0, "diversity_beta": 1.0, + "rescale_norm": True, "global_scale": 0.7, + }, + "high_diversity_concept": { + "consensus_type": "median", "power_alpha": 2.0, + "similarity_threshold": 0.0, "diversity_beta": 2.0, + "rescale_norm": True, "global_scale": 0.7, + }, + "dsc_baseline": { + "consensus_type": "median", "power_alpha": 2.0, + "similarity_threshold": 0.0, "diversity_beta": 0.0, + "rescale_norm": False, "global_scale": 1.0, + "dynamic_similarity_contrast": True, "soft_comfort_bandpass": True, + }, + "dsc_high_clarity": { + "consensus_type": "median", "power_alpha": 4.0, + "similarity_threshold": 0.3, "diversity_beta": 0.0, + "rescale_norm": False, "global_scale": 1.0, + "dynamic_similarity_contrast": True, "soft_comfort_bandpass": True, + }, + "dsc_smooth": { + "consensus_type": "mean", "power_alpha": 1.0, + "similarity_threshold": 0.0, "diversity_beta": 0.0, + "rescale_norm": False, "global_scale": 1.0, + "dynamic_similarity_contrast": True, "soft_comfort_bandpass": True, + }, + "dsc_varied_merge": { + "consensus_type": "median", "power_alpha": 2.5, + "similarity_threshold": 0.0, "diversity_beta": 0.0, + "rescale_norm": True, "global_scale": 0.7, + "dynamic_similarity_contrast": True, "soft_comfort_bandpass": True, + }, + "dsc_diverse_concept": { + "consensus_type": "median", "power_alpha": 2.0, + "similarity_threshold": 0.0, "diversity_beta": 1.5, + "rescale_norm": True, "global_scale": 0.7, + "dynamic_similarity_contrast": True, "soft_comfort_bandpass": True, + }, + "dsc_high_diversity_concept": { + "consensus_type": "median", "power_alpha": 2.0, + "similarity_threshold": 0.0, "diversity_beta": 3.0, + "rescale_norm": True, "global_scale": 0.7, + "dynamic_similarity_contrast": True, "soft_comfort_bandpass": True, + }, +} + +CWB_PRESET_OPTIONS = ["custom", *CWB_PRESETS] +LORA_PREFIXES = ( + "base_model.model.", + "lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", + "lora_te_", "lycoris_", "diffusion_model.", "transformer.", "unet.", +) + + +@dataclass(frozen=True) +class CWBSettings: + consensus_type: str + alignment_method: str + alignment_threshold: float + similarity_threshold: float + power_alpha: float + diversity_beta: float + rescale_norm: bool + global_scale: float + dynamic_similarity_contrast: bool + soft_comfort_bandpass: bool + position_weight: float + preserve_common_prefix: bool + + +def resolve_cwb_settings(**values) -> CWBSettings: + """Resolve a named preset without hiding position/prefix controls.""" + preset_name = values.get("cwb_preset", "baseline") + resolved = { + "consensus_type": values.get("consensus_type", "median"), + "alignment_method": values.get("alignment_method", "similarity"), + "alignment_threshold": float(values.get("alignment_threshold", 0.4)), + "similarity_threshold": float(values.get("similarity_threshold", 0.0)), + "power_alpha": float(values.get("power_alpha", 2.0)), + "diversity_beta": float(values.get("diversity_beta", 0.0)), + "rescale_norm": bool(values.get("rescale_norm", False)), + "global_scale": float(values.get("global_scale", 1.0)), + "dynamic_similarity_contrast": bool( + values.get("dynamic_similarity_contrast", False) + ), + "soft_comfort_bandpass": bool(values.get("soft_comfort_bandpass", False)), + "position_weight": float(values.get("position_weight", 0.0)), + "preserve_common_prefix": bool(values.get("preserve_common_prefix", False)), + } + if not 0.0 <= resolved["position_weight"] <= 1.0: + raise ValueError("Position weight must be between 0.0 and 1.0.") + if preset_name != "custom": + try: + preset = CWB_PRESETS[preset_name] + except KeyError as exc: + raise ValueError(f"Unsupported CWB preset: {preset_name}") from exc + user_scale = resolved["global_scale"] + resolved.update( + alignment_method="similarity", + dynamic_similarity_contrast=False, + soft_comfort_bandpass=False, + ) + resolved.update(preset) + if user_scale != 1.0: + resolved["global_scale"] = user_scale + return CWBSettings(**resolved) + + +def _common_prefix_length(tensors: list[torch.Tensor]) -> int: + if not tensors: + return 0 + limit = min(tensor.shape[0] for tensor in tensors) + if limit == 0: + return 0 + reference = tensors[0][:limit] + common = torch.ones(limit, dtype=torch.bool, device=reference.device) + for tensor in tensors[1:]: + comparison = torch.isclose(reference, tensor[:limit], rtol=1e-5, atol=1e-6) + common &= comparison.reshape(limit, -1).all(dim=1) + mismatch = torch.nonzero(~common, as_tuple=False) + return limit if mismatch.numel() == 0 else int(mismatch[0].item()) + + +def _position_biased_scores(scores: torch.Tensor, weight: float) -> torch.Tensor: + if weight <= 0.0: + return scores + n_ref, n_source = scores.shape + if n_ref == 0 or n_source == 0: + return scores + ref_positions = torch.linspace(0.0, 1.0, n_ref, device=scores.device) + source_positions = torch.linspace(0.0, 1.0, n_source, device=scores.device) + distance = ref_positions[:, None] - source_positions[None, :] + sigma = max(1.0 - min(weight, 1.0), 1.0 / max(n_ref, n_source, 1)) + affinity = torch.exp(-0.5 * (distance / sigma) ** 2).to(scores) + return scores * (1.0 - weight) + affinity * weight + + +def merge_consensus_group( + stacked: torch.Tensor, settings: CWBSettings +) -> torch.Tensor: + """Merge one aligned vector group using CWB.""" + if stacked.shape[0] == 1: + merged = stacked[0].clone() + else: + consensus = ( + torch.median(stacked, dim=0).values + if settings.consensus_type == "median" + else torch.mean(stacked, dim=0) + ) + stacked_norm = F.normalize(stacked, p=2, dim=1, eps=1e-8) + consensus_norm = F.normalize(consensus, p=2, dim=0, eps=1e-8) + similarities = torch.mv(stacked_norm, consensus_norm) + weighted = similarities + if settings.dynamic_similarity_contrast: + min_sim = similarities.min() + max_sim = similarities.max() + if max_sim > min_sim: + weighted = 0.7 + 0.3 * ( + similarities - min_sim + ) / (max_sim - min_sim + 1e-8) + + row_weights = torch.zeros_like(similarities) + mask = similarities >= settings.similarity_threshold + if mask.any(): + safe = weighted[mask].clamp(min=0.0, max=1.0) + row_weights[mask] = torch.pow(safe, settings.power_alpha) + if settings.diversity_beta > 0.0: + distance_base = 1.5 if settings.soft_comfort_bandpass else 1.001 + row_weights[mask] *= torch.pow( + (distance_base - safe).clamp(min=0.0), + settings.diversity_beta, + ) + weight_sum = row_weights.sum() + if weight_sum > 0: + row_weights /= weight_sum + else: + row_weights.fill_(1.0 / len(similarities)) + else: + row_weights.fill_(1.0 / len(similarities)) + merged = (stacked * row_weights.unsqueeze(1)).sum(dim=0) + + if settings.rescale_norm: + average_norm = torch.norm(stacked, p=2, dim=1).mean() + merged_norm = torch.norm(merged, p=2) + if merged_norm > 0: + merged = (merged / merged_norm) * average_norm + if settings.global_scale != 1.0: + merged *= settings.global_scale + return merged + + +def merge_cwb_tensors( + tensors: list[torch.Tensor], + settings: CWBSettings, + *, + reference_index: int = 0, +) -> torch.Tensor: + """CWB tensors with a shared rank and trailing vector dimensions.""" + if not tensors: + raise ValueError("CWB requires at least one tensor.") + if any(tensor.ndim != tensors[0].ndim for tensor in tensors): + raise ValueError("CWB source tensors must have matching ranks.") + + if tensors[0].ndim == 0: + stacked = torch.stack(tensors) + merged = ( + torch.median(stacked, dim=0).values + if settings.consensus_type == "median" + else torch.mean(stacked, dim=0) + ) + return merged * settings.global_scale + + if tensors[0].ndim == 1: + target = tensors[reference_index].shape[0] + aligned = [] + for tensor in tensors: + if tensor.shape[0] < target: + tensor = F.pad(tensor, (0, target - tensor.shape[0])) + else: + tensor = tensor[:target] + aligned.append(tensor) + stacked = torch.stack(aligned) + merged = ( + torch.median(stacked, dim=0).values + if settings.consensus_type == "median" + else torch.mean(stacked, dim=0) + ) + return merged * settings.global_scale + + trailing_shapes = {tuple(tensor.shape[1:]) for tensor in tensors} + if len(trailing_shapes) != 1: + raise ValueError("CWB vector dimensions must match after alignment.") + + prefix_length = _common_prefix_length(tensors) if settings.preserve_common_prefix else 0 + preserved_prefix = tensors[0][:prefix_length] + bodies = [tensor[prefix_length:] for tensor in tensors] + adjusted_reference = reference_index + reference = bodies[adjusted_reference] + if reference.shape[0] == 0: + return preserved_prefix.clone() + + groups: list[list[torch.Tensor]] = [[reference[row]] for row in range(reference.shape[0])] + reference_flat = reference.reshape(reference.shape[0], -1) + + for index, tensor in enumerate(bodies): + if index == adjusted_reference or tensor.shape[0] == 0: + continue + flat = tensor.reshape(tensor.shape[0], -1) + if settings.alignment_method == "index": + for row in range(min(reference.shape[0], tensor.shape[0])): + groups[row].append(tensor[row]) + continue + + similarities = torch.mm( + F.normalize(reference_flat, p=2, dim=1, eps=1e-8), + F.normalize(flat, p=2, dim=1, eps=1e-8).t(), + ) + scores = _position_biased_scores(similarities, settings.position_weight) + scores = scores.masked_fill( + similarities < settings.alignment_threshold, + -100.0, + ) + matched = [-1] * reference.shape[0] + working = scores.clone() + for _ in range(min(reference.shape[0], tensor.shape[0])): + flat_index = torch.argmax(working) + best = working.flatten()[flat_index].item() + if best <= -100.0: + break + ref_row = int((flat_index // tensor.shape[0]).item()) + source_row = int((flat_index % tensor.shape[0]).item()) + matched[ref_row] = source_row + working[ref_row, :] = -100.0 + working[:, source_row] = -100.0 + for ref_row, source_row in enumerate(matched): + if source_row >= 0: + groups[ref_row].append(tensor[source_row]) + + merged_rows = [] + for group in groups: + stacked = torch.stack([row.reshape(-1) for row in group]) + merged_rows.append(merge_consensus_group(stacked, settings).reshape(group[0].shape)) + merged = torch.stack(merged_rows) + if prefix_length: + merged = torch.cat([preserved_prefix, merged], dim=0) + return merged + + +def _copy_to_target_shape(tensor: torch.Tensor, target: torch.Size) -> torch.Tensor | None: + if tensor.ndim != len(target): + return None + if tensor.ndim == 0: + return tensor if tensor.shape == target else None + output_shape = list(target) + output = tensor.new_zeros(output_shape) + slices = tuple(slice(0, min(source, wanted)) for source, wanted in zip(tensor.shape, output_shape)) + output[slices] = tensor[slices] + return output + + +def _to_compute(tensor: torch.Tensor, device: str) -> torch.Tensor: + if device == "cuda": + return transfer_to_gpu_pinned(tensor, device, torch.float32) + return tensor.to(device=device, dtype=torch.float32) + + +def _clear_previous_layer(params: dict) -> None: + if not params["force_clear_cache"]: + return + import gc + + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def _requested_dtype(name: str) -> torch.dtype: + return { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, + }[name] + + +def _is_float_dtype(dtype: torch.dtype) -> bool: + return dtype in { + torch.float16, torch.bfloat16, torch.float32, torch.float64, + } + + +def _normalized_lora_core(block_name: str) -> str: + core = block_name + changed = True + while changed: + changed = False + for prefix in LORA_PREFIXES: + if core.startswith(prefix): + core = core[len(prefix):] + changed = True + break + if core.endswith(".lora"): + core = core[:-5] + return core.replace(".", "_") + + +def _secondary_lora_map(pairs: dict[str, dict[str, str]]) -> dict[str, str]: + result = {} + for block_name in pairs: + result.setdefault(_normalized_lora_core(block_name), block_name) + return result + + +def _pad_rank_axis(tensor: torch.Tensor, axis: int, rank: int) -> torch.Tensor: + if tensor.shape[axis] == rank: + return tensor + padding = [0] * (tensor.ndim * 2) + reverse_axis = tensor.ndim - 1 - axis + padding[reverse_axis * 2 + 1] = rank - tensor.shape[axis] + return F.pad(tensor, tuple(padding)) + + +class ConsensusMergerLogic: + """Shared streaming execution for the dedicated CWB node family.""" + + @classmethod + def execute( + cls, + model_names: list[str], + model_type: str, + params: dict, + *, + embedding_union: bool = False, + lora_mode: bool = False, + ) -> str: + paths = [] + for name in model_names: + path = folder_paths.get_full_path(model_type, name) + if not path: + raise FileNotFoundError(f"{model_type} input '{name}' was not found.") + paths.append(path) + + total_size = sum(estimate_model_size(path) for path in paths) + device = params["process_device"] + if total_size: + prepare_for_large_operation(total_size * 1.2, torch.device(device)) + + handlers = [MemoryEfficientSafeOpen(path, low_memory=params["lazy_load"]) for path in paths] + 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)) + ] + settings = resolve_cwb_settings(**params) + if lora_mode: + return cls._merge_loras(handlers, low_bit_sets, model_type, params, settings) + return cls._merge_generic( + handlers, + low_bit_sets, + model_type, + params, + settings, + embedding_union=embedding_union, + ) + finally: + for handler in handlers: + handler.__exit__(None, None, None) + cleanup_after_operation() + + @staticmethod + def _output_path(model_type: str, output_filename: str) -> str: + output_dir = folder_paths.get_folder_paths(model_type)[-1] + os.makedirs(output_dir, exist_ok=True) + return os.path.join(output_dir, f"{output_filename.strip()}.safetensors") + + @classmethod + def _merge_generic( + cls, + handlers, + low_bit_sets, + model_type, + params, + settings, + *, + embedding_union, + ): + primary = handlers[0] + keys = set(primary.keys()) + if embedding_union: + for handler in handlers[1:]: + keys.update(handler.keys()) + keys = sorted(keys) + output_path = cls._output_path(model_type, params["output_filename"]) + requested_dtype = _requested_dtype(params["save_dtype"]) + mismatch_mode = params["mismatch_mode"] + glob_mode = params["glob_patterns"] + exclude = _compile_patterns(params["exclude_patterns"], glob_mode=glob_mode) + discard = _compile_patterns(params["discard_patterns"], glob_mode=glob_mode) + pbar = comfy.utils.ProgressBar(len(keys)) + + writer = IncrementalSafetensorsWriter(output_path, metadata=primary.metadata()) + writer.__enter__() + try: + with torch.no_grad(): + for key in tqdm(keys, desc="CWB merging tensors", unit="tensors"): + _clear_previous_layer(params) + if _matches_any_pattern(key, discard, glob_mode=glob_mode): + pbar.update(1) + continue + source_indices = [i for i, handler in enumerate(handlers) if key in handler.keys()] + preserve_index = 0 if key in primary.keys() else source_indices[0] + guarded = any(key in low_bit_sets[i] for i in source_indices) + excluded = _matches_any_pattern(key, exclude, glob_mode=glob_mode) + if guarded or excluded: + write_preserved_tensor(writer, key, handlers[preserve_index]) + pbar.update(1) + continue + + source_dtypes = [handlers[i].get_dtype(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]) + pbar.update(1) + continue + + if not embedding_union and len(source_indices) != len(handlers): + 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) + pbar.update(1) + continue + + raw = {i: handlers[i].get_tensor(key) for i in source_indices} + reference_index = 0 + if embedding_union: + reference_source = max( + source_indices, + key=lambda i: raw[i].shape[0] if raw[i].ndim else 1, + ) + reference_shape = raw[reference_source].shape + else: + reference_source = 0 + reference_shape = raw[0].shape + + tensors = [] + actual_dtypes = [] + reference_index = 0 + for source_index, handler in enumerate(handlers): + if source_index not in raw: + if mismatch_mode == "error": + raise ValueError(f"Tensor '{key}' is missing from input {source_index + 1}.") + if mismatch_mode == "zeros": + tensors.append(torch.zeros( + reference_shape, + dtype=torch.float32, + device=params["process_device"], + )) + continue + aligned = raw[source_index] + if not embedding_union: + aligned = _copy_to_target_shape(aligned, reference_shape) + elif aligned.ndim != len(reference_shape) or ( + aligned.ndim > 1 and tuple(aligned.shape[1:]) != tuple(reference_shape[1:]) + ): + aligned = None + if aligned is None: + if mismatch_mode == "error": + raise ValueError(f"Tensor shape mismatch for '{key}'.") + if mismatch_mode == "zeros": + tensors.append(torch.zeros( + reference_shape, + dtype=torch.float32, + device=params["process_device"], + )) + continue + if source_index == reference_source: + reference_index = len(tensors) + tensors.append(_to_compute(aligned, params["process_device"])) + actual_dtypes.append(handler.get_dtype(key)) + + if not tensors: + pbar.update(1) + continue + merged = merge_cwb_tensors( + tensors, + settings, + reference_index=reference_index, + ) + target_dtype = select_output_dtype( + actual_dtypes, + requested_dtype, + force=params["override_dtype"], + ) + writer.write(key, merged.to(target_dtype).cpu().contiguous()) + pbar.update(1) + finally: + writer.__exit__(None, None, None) + return os.path.basename(output_path) + + @classmethod + def _merge_loras(cls, handlers, low_bit_sets, model_type, params, settings): + parsed = [parse_lora_layers(handler.keys()) for handler in handlers] + primary_pairs, primary_passthrough = parsed[0] + secondary_maps = [_secondary_lora_map(pairs) for pairs, _ in parsed] + output_path = cls._output_path(model_type, params["output_filename"]) + requested_dtype = _requested_dtype(params["save_dtype"]) + mismatch_mode = params["mismatch_mode"] + include_1d = params.get("include_1d_diffs", False) + glob_mode = params["glob_patterns"] + exclude = _compile_patterns(params["exclude_patterns"], glob_mode=glob_mode) + discard = _compile_patterns(params["discard_patterns"], glob_mode=glob_mode) + written = set() + pbar = comfy.utils.ProgressBar(len(primary_pairs) + len(primary_passthrough)) + + def preserve_keys(writer, keys: Iterable[str]): + for key in keys: + if key not in written: + write_preserved_tensor(writer, key, handlers[0]) + written.add(key) + + writer = IncrementalSafetensorsWriter(output_path, metadata=handlers[0].metadata()) + writer.__enter__() + try: + with torch.no_grad(): + for primary_block, primary_keys in tqdm( + primary_pairs.items(), desc="CWB merging LoRA layers", unit="layers" + ): + _clear_previous_layer(params) + core = _normalized_lora_core(primary_block) + matches = [] + for index, (pairs, _) in enumerate(parsed): + block = primary_block if index == 0 else secondary_maps[index].get(core) + matches.append((index, block, pairs.get(block) if block else None)) + + layer_keys = [ + key for key in handlers[0].keys() + if key in primary_keys.values() or key.startswith(f"{primary_block}.") + ] + if any(_matches_any_pattern(key, discard, glob_mode=glob_mode) for key in layer_keys): + pbar.update(1) + continue + guarded = any( + keys is not None and layer_has_low_bit(keys, low_bit_sets[index]) + for index, _, keys in matches + ) + excluded = any( + _matches_any_pattern(key, exclude, glob_mode=glob_mode) + for key in layer_keys + ) + if guarded or excluded: + preserve_keys(writer, layer_keys) + pbar.update(1) + continue + + logical_dtypes = [] + for index, _, keys in matches: + if keys: + logical_dtypes.extend( + handlers[index].get_dtype(key) + for name, key in keys.items() + if name in {"down", "up", "alpha", "diff", "diff_b"} + ) + + if "down" in primary_keys and "up" in primary_keys: + pair_sources = [] + pair_failed = False + primary_down_template = handlers[0].get_tensor(primary_keys["down"]) + primary_up_template = handlers[0].get_tensor(primary_keys["up"]) + for index, _, keys in matches: + if not keys or "down" not in keys or "up" not in keys: + if mismatch_mode == "error": + raise ValueError(f"LoRA pair '{primary_block}' is missing from input {index + 1}.") + if mismatch_mode == "skip": + pair_failed = True + break + pair_sources.append(( + index, + torch.zeros_like(primary_down_template), + torch.zeros_like(primary_up_template), + )) + continue + down = handlers[index].get_tensor(keys["down"]) + up = handlers[index].get_tensor(keys["up"]) + if down.ndim < 2 or up.ndim < 2 or down.shape[0] != up.shape[1]: + if mismatch_mode == "error" or index == 0: + raise ValueError(f"Invalid LoRA rank dimensions for '{primary_block}'.") + if mismatch_mode == "skip": + pair_failed = True + break + pair_sources.append(( + index, + torch.zeros_like(primary_down_template), + torch.zeros_like(primary_up_template), + )) + continue + pair_sources.append((index, down, up)) + if pair_failed: + preserve_keys(writer, [primary_keys["down"], primary_keys["up"]]) + if "alpha" in primary_keys: + preserve_keys(writer, [primary_keys["alpha"]]) + elif pair_sources: + primary_down = pair_sources[0][1] + primary_up = pair_sources[0][2] + compatible = [] + for index, down, up in pair_sources: + valid = ( + tuple(down.shape[1:]) == tuple(primary_down.shape[1:]) + and up.shape[0] == primary_up.shape[0] + and tuple(up.shape[2:]) == tuple(primary_up.shape[2:]) + ) + if not valid: + if mismatch_mode == "error": + raise ValueError(f"LoRA pair shape mismatch for '{primary_block}'.") + if mismatch_mode == "skip": + compatible = [] + break + compatible.append(( + index, + torch.zeros_like(primary_down), + torch.zeros_like(primary_up), + )) + continue + compatible.append((index, down, up)) + if not compatible: + preserve_keys(writer, [primary_keys["down"], primary_keys["up"]]) + if "alpha" in primary_keys: + preserve_keys(writer, [primary_keys["alpha"]]) + else: + max_rank = max(down.shape[0] for _, down, _ in compatible) + downs = [ + _to_compute(_pad_rank_axis(down, 0, max_rank), params["process_device"]) + for _, down, _ in compatible + ] + ups = [ + _to_compute(_pad_rank_axis(up, 1, max_rank), params["process_device"]) + for _, _, up in compatible + ] + merged_down = merge_cwb_tensors(downs, settings) + merged_up = merge_cwb_tensors(ups, settings) + target_dtype = select_output_dtype( + logical_dtypes, + requested_dtype, + force=params["override_dtype"], + ) + writer.write( + primary_keys["down"], + merged_down.to(target_dtype).cpu().contiguous(), + ) + writer.write( + primary_keys["up"], + merged_up.to(target_dtype).cpu().contiguous(), + ) + written.update({primary_keys["down"], primary_keys["up"]}) + if "alpha" in primary_keys: + writer.write( + primary_keys["alpha"], + torch.tensor(float(max_rank), dtype=target_dtype), + ) + written.add(primary_keys["alpha"]) + + for kind in ("diff", "diff_b"): + if kind not in primary_keys: + continue + primary_key = primary_keys[kind] + primary_tensor = handlers[0].get_tensor(primary_key) + if primary_tensor.ndim == 1 and not include_1d: + preserve_keys(writer, [primary_key]) + continue + direct = [] + direct_dtypes = [] + failed = False + for index, _, keys in matches: + if not keys or kind not in keys: + if mismatch_mode == "error": + raise ValueError(f"Direct LoRA layer '{primary_block}.{kind}' is missing from input {index + 1}.") + if mismatch_mode == "skip": + failed = True + break + direct.append(torch.zeros_like(primary_tensor, dtype=torch.float32, device=params["process_device"])) + continue + tensor = handlers[index].get_tensor(keys[kind]) + if tensor.shape != primary_tensor.shape: + if mismatch_mode == "error": + raise ValueError(f"Direct LoRA shape mismatch for '{primary_block}.{kind}'.") + if mismatch_mode == "skip": + failed = True + break + direct.append(torch.zeros_like(primary_tensor, dtype=torch.float32, device=params["process_device"])) + continue + direct.append(_to_compute(tensor, params["process_device"])) + direct_dtypes.append(handlers[index].get_dtype(keys[kind])) + if failed: + preserve_keys(writer, [primary_key]) + continue + merged = merge_cwb_tensors(direct, settings) + target_dtype = select_output_dtype( + direct_dtypes, + requested_dtype, + force=params["override_dtype"], + is_1d_diff=merged.ndim == 1, + ) + writer.write(primary_key, merged.to(target_dtype).cpu().contiguous()) + written.add(primary_key) + + preserve_keys(writer, [key for key in layer_keys if key not in written]) + pbar.update(1) + + for key in primary_passthrough: + _clear_previous_layer(params) + if key in written: + continue + if _matches_any_pattern(key, discard, glob_mode=glob_mode): + pbar.update(1) + continue + preserve_keys(writer, [key]) + pbar.update(1) + finally: + writer.__exit__(None, None, None) + return os.path.basename(output_path) + + +def _common_inputs(model_type: str, input_count: int, default_filename: str): + label = model_type + inputs = [ + io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]), + io.Combo.Input("model_a", options=folder_paths.get_filename_list(label)), + io.Combo.Input("model_b", options=folder_paths.get_filename_list(label)), + ] + if input_count == 3: + inputs.append(io.Combo.Input("model_c", options=folder_paths.get_filename_list(label))) + inputs.extend([ + io.Combo.Input("cwb_preset", options=CWB_PRESET_OPTIONS, default="baseline"), + io.Combo.Input("consensus_type", options=["mean", "median"], default="median"), + io.Combo.Input("alignment_method", options=["index", "similarity"], default="similarity"), + io.Float.Input("alignment_threshold", default=0.4, min=0.0, max=1.0, step=0.01), + io.Float.Input("similarity_threshold", default=0.0, min=-1.0, max=1.0, step=0.01), + io.Float.Input("power_alpha", default=2.0, min=0.0, max=10.0, step=0.1), + io.Float.Input("diversity_beta", default=0.0, min=0.0, max=10.0, step=0.1), + io.Boolean.Input("rescale_norm", default=False), + io.Float.Input("global_scale", default=1.0, min=0.0, max=10.0, step=0.01), + io.Boolean.Input("dynamic_similarity_contrast", default=False), + io.Boolean.Input("soft_comfort_bandpass", default=False), + io.Float.Input("position_weight", default=0.0, min=0.0, max=1.0, step=0.01), + io.Boolean.Input("preserve_common_prefix", default=False), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), + io.String.Input("output_filename", default=default_filename), + io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), + io.Combo.Input("process_device", options=["cuda", "cpu"]), + io.String.Input("exclude_patterns", default="", multiline=True), + io.String.Input("discard_patterns", default="", multiline=True), + io.Boolean.Input("glob_patterns", default=False), + io.Boolean.Input("lazy_load", default=True), + io.Boolean.Input("force_clear_cache", default=True), + io.Boolean.Input( + "override_dtype", + default=False, + tooltip="Force generated tensors to save_dtype; guarded tensors and enabled 1D direct diffs are exempt.", + ), + ]) + return inputs + + +class _CWBMergerNode(io.ComfyNode): + MODEL_TYPE = "diffusion_models" + INPUT_COUNT = 2 + NODE_ID = "" + DISPLAY_NAME = "" + DEFAULT_FILENAME = "cwb_merged" + EMBEDDING_UNION = False + LORA_MODE = False + + @classmethod + def define_schema(cls): + inputs = _common_inputs(cls.MODEL_TYPE, cls.INPUT_COUNT, cls.DEFAULT_FILENAME) + if cls.LORA_MODE: + inputs.append(io.Boolean.Input( + "include_1d_diffs", + default=False, + tooltip="CWB-merge 1D direct diffs as FP32. Disabled preserves Model A.", + )) + return io.Schema( + node_id=cls.NODE_ID, + display_name=cls.DISPLAY_NAME, + category="ModelUtils/Merging", + inputs=inputs, + outputs=[ + io.String.Output(display_name="output_filename"), + io.String.Output(display_name="documentation"), + ], + is_experimental=True, + ) + + @classmethod + def execute(cls, **kwargs): + documentation = load_documentation_from_file("consensus_mergers.md") + if kwargs["execution_mode"] == "DOCUMENTATION ONLY": + return io.NodeOutput("Documentation mode active. No merge performed.", documentation) + names = [kwargs["model_a"], kwargs["model_b"]] + if cls.INPUT_COUNT == 3: + names.append(kwargs["model_c"]) + filename = ConsensusMergerLogic.execute( + names, + cls.MODEL_TYPE, + kwargs, + embedding_union=cls.EMBEDDING_UNION, + lora_mode=cls.LORA_MODE, + ) + return io.NodeOutput(filename, documentation) + + +class CWBCheckpointTwoMerger(_CWBMergerNode): + MODEL_TYPE = "checkpoints" + NODE_ID = "CWBCheckpointTwoMerger" + DISPLAY_NAME = "CWB Merge Checkpoints (2 Models)" + DEFAULT_FILENAME = "cwb_merged_2_checkpoint" + + +class CWBCheckpointThreeMerger(CWBCheckpointTwoMerger): + INPUT_COUNT = 3 + NODE_ID = "CWBCheckpointThreeMerger" + DISPLAY_NAME = "CWB Merge Checkpoints (3 Models)" + DEFAULT_FILENAME = "cwb_merged_3_checkpoint" + + +class CWBModelTwoMerger(_CWBMergerNode): + NODE_ID = "CWBModelTwoMerger" + DISPLAY_NAME = "CWB Merge Models (2 Models)" + DEFAULT_FILENAME = "cwb_merged_2_model" + + +class CWBModelThreeMerger(CWBModelTwoMerger): + INPUT_COUNT = 3 + NODE_ID = "CWBModelThreeMerger" + DISPLAY_NAME = "CWB Merge Models (3 Models)" + DEFAULT_FILENAME = "cwb_merged_3_model" + + +class CWBTextEncoderTwoMerger(_CWBMergerNode): + MODEL_TYPE = "text_encoders" + NODE_ID = "CWBTextEncoderTwoMerger" + DISPLAY_NAME = "CWB Merge Text Encoders (2 Models)" + DEFAULT_FILENAME = "cwb_merged_2_text_encoder" + + +class CWBTextEncoderThreeMerger(CWBTextEncoderTwoMerger): + INPUT_COUNT = 3 + NODE_ID = "CWBTextEncoderThreeMerger" + DISPLAY_NAME = "CWB Merge Text Encoders (3 Models)" + DEFAULT_FILENAME = "cwb_merged_3_text_encoder" + + +class CWBLoRATwoMerger(_CWBMergerNode): + MODEL_TYPE = "loras" + NODE_ID = "CWBLoRATwoMerger" + DISPLAY_NAME = "CWB Merge LoRAs (2 Models)" + DEFAULT_FILENAME = "cwb_merged_2_lora" + LORA_MODE = True + + +class CWBLoRAThreeMerger(CWBLoRATwoMerger): + INPUT_COUNT = 3 + NODE_ID = "CWBLoRAThreeMerger" + DISPLAY_NAME = "CWB Merge LoRAs (3 Models)" + DEFAULT_FILENAME = "cwb_merged_3_lora" + + +class CWBEmbeddingTwoMerger(_CWBMergerNode): + MODEL_TYPE = "embeddings" + NODE_ID = "CWBEmbeddingTwoMerger" + DISPLAY_NAME = "CWB Merge Embeddings (2 Models)" + DEFAULT_FILENAME = "cwb_merged_2_embedding" + EMBEDDING_UNION = True + + +class CWBEmbeddingThreeMerger(CWBEmbeddingTwoMerger): + INPUT_COUNT = 3 + NODE_ID = "CWBEmbeddingThreeMerger" + DISPLAY_NAME = "CWB Merge Embeddings (3 Models)" + DEFAULT_FILENAME = "cwb_merged_3_embedding" + + +CWB_MERGER_NODES = [ + CWBCheckpointTwoMerger, + CWBCheckpointThreeMerger, + CWBModelTwoMerger, + CWBModelThreeMerger, + CWBTextEncoderTwoMerger, + CWBTextEncoderThreeMerger, + CWBLoRATwoMerger, + CWBLoRAThreeMerger, + CWBEmbeddingTwoMerger, + CWBEmbeddingThreeMerger, +] diff --git a/tests/test_consensus_merger.py b/tests/test_consensus_merger.py new file mode 100644 index 0000000..8c49f44 --- /dev/null +++ b/tests/test_consensus_merger.py @@ -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]) + )